diff --git a/Project.toml b/Project.toml index 70b32d00e..1f330fe61 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Compat" uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.1.0" +version = "4.2.0" [deps] Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" diff --git a/README.md b/README.md index b38de05e0..d4c7ce558 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,8 @@ changes in `julia`. ## Supported features +* `stack` combines a collection of slices into one array ([#43334]). (since Compat 4.2.0) + * `keepat!` removes the items at all the indices which are not given and returns the modified source ([#36229], [#42351]). (since Compat 4.1.0) @@ -150,3 +152,4 @@ Note that you should specify the correct minimum version for `Compat` in the [#42125]: https://github.com/JuliaLang/julia/issues/42125 [#42351]: https://github.com/JuliaLang/julia/issues/42351 [#43354]: https://github.com/JuliaLang/julia/issues/43354 +[#43334]: https://github.com/JuliaLang/julia/issues/43334 diff --git a/src/Compat.jl b/src/Compat.jl index 0dec8e775..27f05d2a0 100644 --- a/src/Compat.jl +++ b/src/Compat.jl @@ -362,6 +362,241 @@ end end end +# https://github.com/JuliaLang/julia/pull/43334 +if VERSION < v"1.9.0-DEV.1163" + import Base: IteratorSize, HasLength, HasShape, OneTo + export stack + + """ + stack(iter; [dims]) + + Combine a collection of arrays (or other iterable objects) of equal size + into one larger array, by arranging them along one or more new dimensions. + + By default the axes of the elements are placed first, + giving `size(result) = (size(first(iter))..., size(iter)...)`. + This has the same order of elements as [`Iterators.flatten`](@ref)`(iter)`. + + With keyword `dims::Integer`, instead the `i`th element of `iter` becomes the slice + [`selectdim`](@ref)`(result, dims, i)`, so that `size(result, dims) == length(iter)`. + In this case `stack` reverses the action of [`eachslice`](@ref) with the same `dims`. + + The various [`cat`](@ref) functions also combine arrays. However, these all + extend the arrays' existing (possibly trivial) dimensions, rather than placing + the arrays along new dimensions. + They also accept arrays as separate arguments, rather than a single collection. + + !!! compat "Julia 1.9" + This function is available in Julia 1.9, or in Compat 4.2. + + # Examples + ```jldoctest + julia> vecs = (1:2, [30, 40], Float32[500, 600]); + + julia> mat = stack(vecs) + 2×3 Matrix{Float32}: + 1.0 30.0 500.0 + 2.0 40.0 600.0 + + julia> mat == hcat(vecs...) == reduce(hcat, collect(vecs)) + true + + julia> vec(mat) == vcat(vecs...) == reduce(vcat, collect(vecs)) + true + + julia> stack(zip(1:4, 10:99)) # accepts any iterators of iterators + 2×4 Matrix{Int64}: + 1 2 3 4 + 10 11 12 13 + + julia> vec(ans) == collect(Iterators.flatten(zip(1:4, 10:99))) + true + + julia> stack(vecs; dims=1) # unlike any cat function, 1st axis of vecs[1] is 2nd axis of result + 3×2 Matrix{Float32}: + 1.0 2.0 + 30.0 40.0 + 500.0 600.0 + + julia> x = rand(3,4); + + julia> x == stack(eachcol(x)) == stack(eachrow(x), dims=1) # inverse of eachslice + true + ``` + + Higher-dimensional examples: + + ```jldoctest + julia> A = rand(5, 7, 11); + + julia> E = eachslice(A, dims=2); # a vector of matrices + + julia> (element = size(first(E)), container = size(E)) + (element = (5, 11), container = (7,)) + + julia> stack(E) |> size + (5, 11, 7) + + julia> stack(E) == stack(E; dims=3) == cat(E...; dims=3) + true + + julia> A == stack(E; dims=2) + true + + julia> M = (fill(10i+j, 2, 3) for i in 1:5, j in 1:7); + + julia> (element = size(first(M)), container = size(M)) + (element = (2, 3), container = (5, 7)) + + julia> stack(M) |> size # keeps all dimensions + (2, 3, 5, 7) + + julia> stack(M; dims=1) |> size # vec(container) along dims=1 + (35, 2, 3) + + julia> hvcat(5, M...) |> size # hvcat puts matrices next to each other + (14, 15) + ``` + """ + stack(iter; dims=:) = _stack(dims, iter) + + """ + stack(f, args...; [dims]) + + Apply a function to each element of a collection, and `stack` the result. + Or to several collections, [`zip`](@ref)ped together. + + The function should return arrays (or tuples, or other iterators) all of the same size. + These become slices of the result, each separated along `dims` (if given) or by default + along the last dimensions. + + See also [`mapslices`](@ref), [`eachcol`](@ref). + + # Examples + ```jldoctest + julia> stack(c -> (c, c-32), "julia") + 2×5 Matrix{Char}: + 'j' 'u' 'l' 'i' 'a' + 'J' 'U' 'L' 'I' 'A' + + julia> stack(eachrow([1 2 3; 4 5 6]), (10, 100); dims=1) do row, n + vcat(row, row .* n, row ./ n) + end + 2×9 Matrix{Float64}: + 1.0 2.0 3.0 10.0 20.0 30.0 0.1 0.2 0.3 + 4.0 5.0 6.0 400.0 500.0 600.0 0.04 0.05 0.06 + ``` + """ + stack(f, iter; dims=:) = _stack(dims, f(x) for x in iter) + stack(f, xs, yzs...; dims=:) = _stack(dims, f(xy...) for xy in zip(xs, yzs...)) + + _stack(dims::Union{Integer, Colon}, iter) = _stack(dims, IteratorSize(iter), iter) + + _stack(dims, ::IteratorSize, iter) = _stack(dims, collect(iter)) + + function _stack(dims, ::Union{HasShape, HasLength}, iter) + S = Base.@default_eltype iter + T = S != Union{} ? eltype(S) : Any # Union{} occurs for e.g. stack(1,2), postpone the error + if isconcretetype(T) + _typed_stack(dims, T, S, iter) + else # Need to look inside, but shouldn't run an expensive iterator twice: + array = iter isa Union{Tuple, AbstractArray} ? iter : collect(iter) + isempty(array) && return _empty_stack(dims, T, S, iter) + T2 = mapreduce(eltype, promote_type, array) + _typed_stack(dims, T2, eltype(array), array) + end + end + + function _typed_stack(::Colon, ::Type{T}, ::Type{S}, A, Aax=_iterator_axes(A)) where {T, S} + xit = iterate(A) + nothing === xit && return _empty_stack(:, T, S, A) + x1, _ = xit + ax1 = _iterator_axes(x1) + B = similar(_ensure_array(x1), T, ax1..., Aax...) + off = firstindex(B) + len = length(x1) + while xit !== nothing + x, state = xit + _stack_size_check(x, ax1) + copyto!(B, off, x) + off += len + xit = iterate(A, state) + end + B + end + + _iterator_axes(x) = _iterator_axes(x, IteratorSize(x)) + _iterator_axes(x, ::HasLength) = (OneTo(length(x)),) + _iterator_axes(x, ::IteratorSize) = axes(x) + + # For some dims values, stack(A; dims) == stack(vec(A)), and the : path will be faster + _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, A) where {T,S} = + _typed_stack(dims, T, S, IteratorSize(S), A) + _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::HasLength, A) where {T,S} = + _typed_stack(dims, T, S, HasShape{1}(), A) + function _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::HasShape{N}, A) where {T,S,N} + if dims == N+1 + _typed_stack(:, T, S, A, (_vec_axis(A),)) + else + _dim_stack(dims, T, S, A) + end + end + _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::IteratorSize, A) where {T,S} = + _dim_stack(dims, T, S, A) + + _vec_axis(A, ax=_iterator_axes(A)) = length(ax) == 1 ? only(ax) : OneTo(prod(length, ax; init=1)) + + @constprop :aggressive function _dim_stack(dims::Integer, ::Type{T}, ::Type{S}, A) where {T,S} + xit = Iterators.peel(A) + nothing === xit && return _empty_stack(dims, T, S, A) + x1, xrest = xit + ax1 = _iterator_axes(x1) + N1 = length(ax1)+1 + dims in 1:N1 || throw(ArgumentError(string("cannot stack slices ndims(x) = ", N1-1, " along dims = ", dims))) + + newaxis = _vec_axis(A) + outax = ntuple(d -> d==dims ? newaxis : ax1[d - (d>dims)], N1) + B = similar(_ensure_array(x1), T, outax...) + + if dims == 1 + _dim_stack!(Val(1), B, x1, xrest) + elseif dims == 2 + _dim_stack!(Val(2), B, x1, xrest) + else + _dim_stack!(Val(dims), B, x1, xrest) + end + B + end + + function _dim_stack!(::Val{dims}, B::AbstractArray, x1, xrest) where {dims} + before = ntuple(d -> Colon(), dims - 1) + after = ntuple(d -> Colon(), ndims(B) - dims) + + i = firstindex(B, dims) + copyto!(view(B, before..., i, after...), x1) + + for x in xrest + _stack_size_check(x, _iterator_axes(x1)) + i += 1 + @inbounds copyto!(view(B, before..., i, after...), x) + end + end + + @inline function _stack_size_check(x, ax1::Tuple) + if _iterator_axes(x) != ax1 + uax1 = map(UnitRange, ax1) + uaxN = map(UnitRange, axes(x)) + throw(DimensionMismatch( + string("stack expects uniform slices, got axes(x) == ", uaxN, " while first had ", uax1))) + end + end + + _ensure_array(x::AbstractArray) = x + _ensure_array(x) = 1:0 # passed to similar, makes stack's output an Array + + _empty_stack(_...) = throw(ArgumentError("`stack` on an empty collection is not allowed")) +end + include("deprecated.jl") end # module Compat diff --git a/test/runtests.jl b/test/runtests.jl index e45f32dd9..eb0f85b61 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -447,3 +447,112 @@ end keepat!(ea, Bool[]) @test isempty(ea) end + +# https://github.com/JuliaLang/julia/pull/43334 +@testset "stack" begin + # Basics + for args in ([[1, 2]], [1:2, 3:4], [[1 2; 3 4], [5 6; 7 8]], + AbstractVector[1:2, [3.5, 4.5]], Vector[[1,2], [3im, 4im]], + [[1:2, 3:4], [5:6, 7:8]], [fill(1), fill(2)]) + X = stack(args) + Y = cat(args...; dims=ndims(args[1])+1) + @test X == Y + @test typeof(X) === typeof(Y) + + X2 = stack(x for x in args) + @test X2 == Y + @test typeof(X2) === typeof(Y) + + X3 = stack(x for x in args if true) + @test X3 == Y + @test typeof(X3) === typeof(Y) + + if isconcretetype(eltype(args)) + @inferred stack(args) + @inferred stack(x for x in args) + end + end + + # Higher dims + @test size(stack([rand(2,3) for _ in 1:4, _ in 1:5])) == (2,3,4,5) + @test size(stack(rand(2,3) for _ in 1:4, _ in 1:5)) == (2,3,4,5) + @test size(stack(rand(2,3) for _ in 1:4, _ in 1:5 if true)) == (2, 3, 20) + @test size(stack([rand(2,3) for _ in 1:4, _ in 1:5]; dims=1)) == (20, 2, 3) + @test size(stack(rand(2,3) for _ in 1:4, _ in 1:5; dims=2)) == (2, 20, 3) + + # Tuples + @test stack([(1,2), (3,4)]) == [1 3; 2 4] + @test stack(((1,2), (3,4))) == [1 3; 2 4] + @test stack(Any[(1,2), (3,4)]) == [1 3; 2 4] + @test stack([(1,2), (3,4)]; dims=1) == [1 2; 3 4] + @test stack(((1,2), (3,4)); dims=1) == [1 2; 3 4] + @test stack(Any[(1,2), (3,4)]; dims=1) == [1 2; 3 4] + @test size(@inferred stack(Iterators.product(1:3, 1:4))) == (2,3,4) + @test @inferred(stack([('a', 'b'), ('c', 'd')])) == ['a' 'c'; 'b' 'd'] + @test @inferred(stack([(1,2+3im), (4, 5+6im)])) isa Matrix{Number} + + # stack(f, iter) + @test @inferred(stack(x -> [x, 2x], 3:5)) == [3 4 5; 6 8 10] + @test @inferred(stack(x -> x*x'/2, [1:2, 3:4])) == reshape([0.5, 1.0, 1.0, 2.0, 4.5, 6.0, 6.0, 8.0], 2, 2, 2) + @test @inferred(stack(*, [1:2, 3:4], 5:6)) == [5 18; 10 24] + + # Iterators + @test stack([(a=1,b=2), (a=3,b=4)]) == [1 3; 2 4] + @test stack([(a=1,b=2), (c=3,d=4)]) == [1 3; 2 4] + @test stack([(a=1,b=2), (c=3,d=4)]; dims=1) == [1 2; 3 4] + @test stack([(a=1,b=2), (c=3,d=4)]; dims=2) == [1 3; 2 4] + @test stack((x/y for x in 1:3) for y in 4:5) == (1:3) ./ (4:5)' + @test stack((x/y for x in 1:3) for y in 4:5; dims=1) == (1:3)' ./ (4:5) + + # Exotic + ips = ((Iterators.product([i,i^2], [2i,3i,4i], 1:4)) for i in 1:5) + @test size(stack(ips)) == (2, 3, 4, 5) + @test stack(ips) == cat(collect.(ips)...; dims=4) + ips_cat2 = cat(reshape.(collect.(ips), Ref((2,1,3,4)))...; dims=2) + @test stack(ips; dims=2) == ips_cat2 + @test stack(collect.(ips); dims=2) == ips_cat2 + ips_cat3 = cat(reshape.(collect.(ips), Ref((2,3,1,4)))...; dims=3) + @test stack(ips; dims=3) == ips_cat3 # path for non-array accumulation on non-final dims + @test stack(collect, ips; dims=3) == ips_cat3 # ... and for array accumulation + @test stack(collect.(ips); dims=3) == ips_cat3 + + # Trivial, because numbers are iterable: + @test stack(abs2, 1:3) == [1, 4, 9] == collect(Iterators.flatten(abs2(x) for x in 1:3)) + + # Allocation tests + xv = [rand(10) for _ in 1:100] + xt = Tuple.(xv) + for dims in (1, 2, :) + @test stack(xv; dims) == stack(xt; dims) + @test_skip 9000 > @allocated stack(xv; dims) + @test_skip 9000 > @allocated stack(xt; dims) + end + xr = (reshape(1:1000,10,10,10) for _ = 1:1000) + for dims in (1, 2, 3, :) + stack(xr; dims) + @test_skip 8.1e6 > @allocated stack(xr; dims) + end + + # Mismatched sizes + @test_throws DimensionMismatch stack([1:2, 1:3]) + @test_throws DimensionMismatch stack([1:2, 1:3]; dims=1) + @test_throws DimensionMismatch stack([1:2, 1:3]; dims=2) + @test_throws DimensionMismatch stack([(1,2), (3,4,5)]) + @test_throws DimensionMismatch stack([(1,2), (3,4,5)]; dims=1) + @test_throws DimensionMismatch stack(x for x in [1:2, 1:3]) + @test_throws DimensionMismatch stack([[5 6; 7 8], [1, 2, 3, 4]]) + @test_throws DimensionMismatch stack([[5 6; 7 8], [1, 2, 3, 4]]; dims=1) + @test_throws DimensionMismatch stack(x for x in [[5 6; 7 8], [1, 2, 3, 4]]) + # Inner iterator of unknown length + @test_throws MethodError stack((x for x in 1:3 if true) for _ in 1:4) + @test_throws MethodError stack((x for x in 1:3 if true) for _ in 1:4; dims=1) + + @test_throws ArgumentError stack([1:3, 4:6]; dims=0) + @test_throws ArgumentError stack([1:3, 4:6]; dims=3) + @test_throws ArgumentError stack(abs2, 1:3; dims=2) + + # Empty + @test_throws ArgumentError stack(()) + @test_throws ArgumentError stack([]) + @test_throws ArgumentError stack(x for x in 1:3 if false) +end