Extending Base.stack for DimArrays

src/array/methods.jl

@@ -547,6 +547,86 @@ $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([dim::Dimension], A::AbstractVector{<:AbstractDimArray}; dims=nothing)
+
+Stack arrays along a specified axis `dims`, while preserving the dimensional
+information of other axes.
+If the optional `Dimension` argument `dim` is supplied, it must have
+`length(dim) == length(A)`; the resulting axis `dims` is given the dimension `dim`.
+If `dim` is not supplied, the new dimension will be an `AnonDim`.
+
+# 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(Z(1:2), [da, db], dims=3)
+╭─────────────────────────╮
+│ 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(dim::Dimension, A::AbstractVector{<:AbstractDimArray}; dims=nothing, kwargs...)
+    B = Base.stack(A; dims, kwargs...)
+    newdims = ntuple(ndims(B)) do d
+        if d == dims # Use the new provided dimension
+            dim
+        else
+            DimensionalData.dims(B, d)
+        end
+    end
+    rebuild(B; dims=newdims)
+end
+
 function Base.inv(A::AbstractDimArray{T,2}) where T
     newdata = inv(parent(A))
     newdims = reverse(dims(A))

test/methods.jl

@@ -655,6 +655,26 @@ 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_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)
+
+    for d = 1:3
+        dc = stack(Z(1:2), [da, ca], dims=d)
+        @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)