diff --git a/docs/src/api/reference.md b/docs/src/api/reference.md index 3665d92af..ae4c77302 100644 --- a/docs/src/api/reference.md +++ b/docs/src/api/reference.md @@ -105,6 +105,7 @@ Base methods ```@docs Base.cat +Base.stack Base.copy! Base.eachslice ``` diff --git a/src/array/methods.jl b/src/array/methods.jl index 0088751b9..c8a2fc7ab 100644 --- a/src/array/methods.jl +++ b/src/array/methods.jl @@ -529,6 +529,53 @@ $message on dimension $D. To fix for `AbstractDimArray`, pass new lookup values as `cat(As...; dims=$D(newlookupvals))` keyword or `dims=$D()` for empty `NoLookup`. """ +function check_stack_dims(iter) + comparedims(Bool, dims.(iter)...; order=true, val=true, msg=Dimensions.Warn(" Can't `stack` AbstractDimArray, applying to `parent` object.")) + iter +end + +_maybe_dimnum(x, dim) = hasdim(x, dim) ? dimnum(x, dim) : ndims(x) + 1 +function _maybe_dimnum(x, dim::Integer) + if dim < ndims(x) + 2 + return dim + else + throw(ArgumentError(LazyString("cannot stack slices ndims(x) = ", ndims(x) + 1, " along dims = ", dim))) + end +end +_maybe_dimnum(_, ::Colon) = Colon() +_maybe_dimnum(x, dims::Tuple) = map(d -> _maybe_dimnum(x, d), dims) + +""" + stack(x::AbstractDimArray; [dims]) + +Combine a nested `AbstractDimArray` into a single array by arranging the inner arrays +along one or more new dimensions taken from the outer array. +The inner `AbstractDimArray`s must all have the same dimensions. + +The keyword `dims` specifies where the new dimensions will be placed in the resulting array. +The index of all existing dimensions equal to or greater than `dims` will be increased to +accomodate the new dimensions. +The default `dims` places the new dimensions at the end of the resulting array. + +# Examples +```julia +using DimensionalData +a = DimArray([1 2 3; 4 5 6], (X(4.0:5.0), Y(6.0:8.0))) +b = DimArray([7 8 9; 10 11 12], (X(4.0:5.0), Y(6.0:8.0))) +x = DimArray([a, b], Z(4.0:5.0)) # Construct a nested DimArray +y = stack(x; dims=2) # Resulting array has dims (X, *Z*, Y) +``` +""" +function Base.stack(iter::AbstractArray{<:AbstractDimArray}; dims=:) + Base._stack(_maybe_dimnum(first(iter), dims), check_stack_dims(iter)) +end +function Base.stack(f, iter::AbstractArray{<:AbstractDimArray}; kwargs...) + Base.stack(f(x) for x in check_stack_dims(iter); kwargs...) +end +function Base.stack(f, xs::AbstractArray{<:AbstractDimArray}, yzs...; kwargs...) + Base.stack(f(xy...) for xy in zip(xs, yzs...); kwargs...) +end + function Base.inv(A::AbstractDimArray{T,2}) where T newdata = inv(parent(A)) newdims = reverse(dims(A)) diff --git a/test/methods.jl b/test/methods.jl index d874ec245..7bc943d4a 100644 --- a/test/methods.jl +++ b/test/methods.jl @@ -802,6 +802,47 @@ end end end +@testset "Base.stack" begin + a = [1 2 3; 4 5 6] + da = DimArray(a, (X(4.0:5.0), Y(6.0:8.0))) + b = [7 8 9; 10 11 12] + ca = DimArray(b, (X(4.0:5.0), Y(6.0:8.0))) + db = DimArray(b, (X(6.0:7.0), Y(6.0:8.0))) + + x = DimArray([da, db], (Z(4.0:5.0))) + @test_warn "Lookup values" stack(x) + + x = DimArray(fill(da, 2, 2), (Dim{:a}(4.0:5.0), Dim{:b}(6.0:7.0))) + sx = stack(x) + @test sx isa AbstractDimArray + @test dims(sx) == (dims(first(x))..., dims(x)...) + + @test_throws "ArgumentError" stack(x; dims=4) + + x = DimArray([da, ca], (Dim{:a}(1:2),)) + sx = stack(x; dims=1) + sy = @test_nowarn stack(x; dims=:) + sz = @test_nowarn stack(x; dims=X) + @test sx == sz + @test sx == stack([parent(da), parent(ca)], dims=1) + @test sx isa AbstractDimArray + @test dims(sx) == (dims(x)..., dims(first(x))...) + + for d = 1:3 + x = DimArray([da, ca], (AnonDim(),)) + dc = stack(x, dims=d) + @test dims(dc, d) isa AnonDim + @test parent(dc) == stack([da, db], dims=d) + + dc = stack(x -> x .^ 2, x; dims=d) + @test dims(dc, d) isa AnonDim + @test dc == stack(parent(x); dims=d) .^ 2 + dc = stack(+, x, x, x; dims=d) + @test dims(dc, d) isa AnonDim + @test dc == stack(x + x + x; dims=d) + end +end + @testset "unique" begin a = [1 1 6; 1 1 6] xs = (X, X(), :X)