diff --git a/src/array/methods.jl b/src/array/methods.jl index 22726b59f..f6f442fd5 100644 --- a/src/array/methods.jl +++ b/src/array/methods.jl @@ -109,7 +109,7 @@ end ds = DD.dims(A, _astuple(dims)) # Run one slice with dimensions to get the transformed dim d_inds = map(d -> rebuild(d, 1), otherdims(A, ds)) - example_dims = length(d_inds) > 0 ? DD.dims(f(view(A, d_inds...))) : () + example_dims = DD.dims(f(view(A, d_inds...))) replacement_dims = if isnothing(example_dims) || length(example_dims) != length(ds) map(d -> rebuild(d, NoLookup()), ds) else @@ -547,6 +547,100 @@ $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 Base._typed_stack(::Colon, ::Type{T}, ::Type{S}, A, Aax=_iterator_axes(A)) where {T,S<:AbstractDimArray} + origdims = map(dims, A) + _A = parent.(A) + t = eltype(_A) + _A = Base._typed_stack(:, T, t, A) + + if !comparedims(Bool, origdims...; + order=true, val=true, warn=" Can't `stack` AbstractDimArray, applying to `parent` object." + ) + return _A + else + DimArray(_A, (first(origdims)..., AnonDim())) + end +end + +function Base._dim_stack(newdim::Integer, ::Type{T}, ::Type{S}, A) where {T,S<:AbstractDimArray} + origdims = dims.(A) + _A = parent.(A) + t = eltype(_A) + _A = Base._dim_stack(newdim, T, t, A) + + if !comparedims(Bool, origdims...; + order=true, val=true, warn=" Can't `stack` AbstractDimArray, applying to `parent` object." + ) + return _A + end + + newdims = first(origdims) + newdims = ntuple(length(newdims) + 1) do d + if d == newdim + AnonDim() + else # Return the old dimension, shifted across once if it comes after the new dim + newdims[d-(d>newdim)] + end + end + DimArray(_A, newdims) +end + +""" + Base.stack(A::AbstractVector{<:AbstractDimArray}; dims=Pair(ndims(A[1])+1, AnonDim())) + +Stack arrays along a new axis while preserving the dimensional information of other axes. + +The optional keyword argument `dims` has the following behavior: +- `dims isa Integer`: The dimension of the new axis is an `AnonDim` at position `dims` +- `dims isa Dimension`: The new axis is at `ndims(A[1])+1` and has a dimension of `dims`. +- `dims isa Pair{Integer, Dimension}`: The new axis is at `first(dims)` and has a dimension + of `last(dims)`. + +If `dims` contains a `Dimension`, that `Dimension` must have the same length as A. + +# Examples +```julia-repl +julia> da = DimArray([1 2 3; 4 5 6], (X(10:10:20), Y(300:-100:100))); +julia> db = DimArray([6 5 4; 3 2 1], (X(10:10:20), Y(300:-100:100))); + +# Stack along a new dimension `Z` +julia> dc = stack([da, db], dims=3=>Z(1:2)) +╭─────────────────────────╮ +│ 2×3×2 DimArray{Int64,3} │ +├─────────────────────────┴──────────────────────────────── dims ┐ + ↓ X Sampled{Int64} 10:10:20 ForwardOrdered Regular Points, + → Y Sampled{Int64} 300:-100:100 ReverseOrdered Regular Points, + ↗ Z 1:2 +└────────────────────────────────────────────────────────────────┘ + +julia> dims(dc, 3) == Z(1:2) +true +julia> parent(dc) == stack(map(parent, [da, db]), dims=3) +true +``` +""" +function Base.stack(A::AbstractVector{<:AbstractDimArray}; dims=Pair(ndims(A[1])+1, AnonDim())) + if dims isa Integer + dims = dims => AnonDim() + elseif dims isa Dimension + dims = ndims(A[1])+1 => dims + end + + B = Base._stack(first(dims), A) + + if B isa AbstractDimArray + newdims = ntuple(ndims(B)) do d + if d == first(dims) # Use the new provided dimension + last(dims) + else + DimensionalData.dims(B, d) + end + end + B = rebuild(B; dims=format(newdims, B)) + end + return B +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 5d1750abd..1ec034c33 100644 --- a/test/methods.jl +++ b/test/methods.jl @@ -366,12 +366,6 @@ end @test size(y) == size(dims(y)) @test dims(y) == dims(x[2:9, :, :]) end - - @testset "single dim" begin - x = X(1:10) - A = DimArray(ones(10), x) - @test mapslices(sum, A; dims=X) == fill(10.0, X(1)) - end end @testset "cumsum" begin @@ -661,6 +655,37 @@ 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))) + + @test stack([da, db]; dims=3) == stack([parent(da), parent(db)], dims=3) + @test stack([da, db]; dims=3) == stack([da, db]) # Test default dims + @test_warn "Lookup values for X" stack([da, db]; dims=3) + + @test stack([da, ca]; dims=1) == stack([parent(da), parent(ca)], dims=1) + @test_warn "Lookup values for X" stack([da, db]; dims=1) + + dc = stack([da, ca], dims=Z(1:2)) + @test dims(dc, ndims(da)+1) == Z(1:2) + @test parent(dc) == stack(map(parent, [da, db])) + + for d = 1:3 + dc = stack([da, ca], dims=d) + @test dims(dc, d) isa AnonDim + @test parent(dc) == stack(map(parent, [da, db]), dims=d) + end + + for d = 1:3 + dc = stack([da, ca], dims=d=>Z(1:2)) + @test dims(dc, d) == Z(1:2) + @test parent(dc) == stack(map(parent, [da, db]), dims=d) + end +end + @testset "unique" begin a = [1 1 6; 1 1 6] xs = (X, X(), :X)