Add stack (Julia PR 43334)

Backport of PR 777 to the `release-3` branch.

(cherry picked from commit dce2f96)

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "Compat"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.45.0"
+version = "3.46.0"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

README.md

@@ -72,6 +72,8 @@ changes in `julia`.
 * `popat!` removes the item at the given `i` and returns it ([#36070]). (since
   Compat 3.45.0)
 
+* `stack` combines a collection of slices into one array ([#43334]). (since Compat 3.46.0, 4.2.0)
+
 * `keepat!` removes the items at all the indices which are not given and returns
   the modified source ([#36229], [#42351]). (since Compat 3.44.0, 4.1.0)
 
@@ -322,3 +324,4 @@ Note that you should specify the correct minimum version for `Compat` in the
 [#36229]: https://github.com/JuliaLang/julia/issues/36229
 [#39245]: https://github.com/JuliaLang/julia/issues/39245
 [#42351]: https://github.com/JuliaLang/julia/issues/42351
+[#43334]: https://github.com/JuliaLang/julia/issues/43334

src/Compat.jl

@@ -1387,6 +1387,248 @@ end
 end
 
 include("iterators.jl")
+
+# https://github.com/JuliaLang/julia/pull/43334
+if VERSION < v"1.9.0-DEV.1163"
+    export stack
+
+    if VERSION < v"1.1.0-DEV.590"
+        _tuple_axes(x) = axes(x)
+        _tuple_axes(@nospecialize t::Tuple) = (axes(t),)
+    else
+        const _tuple_axes = axes
+    end
+
+    """
+        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, Base.IteratorSize(iter), iter)
+
+    _stack(dims, ::Base.IteratorSize, iter) = _stack(dims, collect(iter))
+
+    function _stack(dims, ::Union{Base.HasShape, Base.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, Base.IteratorSize(x))
+    _iterator_axes(x, ::Base.HasLength) = (Base.OneTo(length(x)),)
+    _iterator_axes(x, ::Base.IteratorSize) = _tuple_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, Base.IteratorSize(S), A)
+    _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::Base.HasLength, A) where {T,S} =
+        _typed_stack(dims, T, S, Base.HasShape{1}(), A)
+    function _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::Base.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}, ::Base.IteratorSize, A) where {T,S} =
+        _dim_stack(dims, T, S, A)
+
+    _vec_axis(A, ax=_iterator_axes(A)) = length(ax) == 1 ? only(ax) : Base.OneTo(mapreduce(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, _tuple_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

src/iterators.jl

@@ -32,7 +32,7 @@ for n in names(Base.Iterators)
 end
 
 # Import unexported public APIs
-using Base.Iterators: filter
+using Base.Iterators: filter, peel
 
 # https://github.com/JuliaLang/julia/pull/33437
 if VERSION < v"1.4.0-DEV.291"  # 5f013d82f92026f7dfbe4234f283658beb1f8a2a

test/runtests.jl

@@ -1466,3 +1466,112 @@ end
     badpop() = @inbounds popat!([1], 2)
     @test_throws BoundsError badpop()
 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=dims) == stack(xt; dims=dims)
+        @test_skip 9000 > @allocated stack(xv; dims=dims)
+        @test_skip 9000 > @allocated stack(xt; dims=dims)
+    end
+    xr = (reshape(1:1000,10,10,10) for _ = 1:1000)
+    for dims in (1, 2, 3, :)
+        stack(xr; dims=dims)
+        @test_skip 8.1e6 > @allocated stack(xr; dims=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