use maplayers in (un)mergedims (#959)

* use maplayers in (un)mergedims

* test unmergedims on a dimstack

Author: tiemvanderdeure

src/stack/stack.jl

@@ -251,20 +251,18 @@ function mergedims(st::AbstractDimStack, dim_pairs::Pair...)
     end
     isempty(dim_pairs) && return st
     # Extend missing dimensions in all layers
-    extended_layers = map(layers(st)) do layer
-        if all(map((ds...) -> all(hasdim(layer, ds)), map(first, dim_pairs)...))
+    maplayers(st) do layer
+        extended_layer = if all(map((ds...) -> all(hasdim(layer, ds)), map(first, dim_pairs)...))
             layer
         else
             DimExtensionArray(layer, dims(st))
         end
+        mergedims(extended_layer, dim_pairs...)
     end
-
-    vals = map(A -> mergedims(A, dim_pairs...), extended_layers)
-    return rebuild_from_arrays(st, vals)
 end
 
 function unmergedims(s::AbstractDimStack, original_dims)
-    return map(A -> unmergedims(A, original_dims), s)
+    return maplayers(A -> unmergedims(A, original_dims), s)
 end
 
 @noinline _stack_size_mismatch() = throw(ArgumentError("Arrays must have identical axes. For mixed dimensions, use DimArrays`"))

test/merged.jl

@@ -66,24 +66,27 @@ end
 
 @testset "unmerge" begin
     a = DimArray(rand(32, 32, 3), (X,Y,Dim{:band}))
-    merged = mergedims(a, (X, Y) => :geometry)
-    unmerged = unmergedims(merged, dims(a))
-    perm_unmerged = unmergedims(permutedims(merged, (2,1)), dims(a))
-    
-    # Test Merge
-    @test hasdim(merged, Dim{:band})
-    @test hasdim(merged, Dim{:geometry})
-    @test !hasdim(merged, X)
-    @test !hasdim(merged, Y)
-    @test size(merged) == (3, 32 * 32)
+    b = DimStack((;a))
+    for a in (a, b)
+        merged = mergedims(a, (X, Y) => :geometry)
+        unmerged = unmergedims(merged, dims(a))
+        perm_unmerged = unmergedims(permutedims(merged, (2,1)), dims(a))
+        
+        # Test Merge
+        @test hasdim(merged, Dim{:band})
+        @test hasdim(merged, Dim{:geometry})
+        @test !hasdim(merged, X)
+        @test !hasdim(merged, Y)
+        @test size(merged) == (3, 32 * 32)
 
-    # Test Unmerge
-    @test hasdim(unmerged, X)
-    @test hasdim(unmerged, Y)
-    @test hasdim(unmerged, Dim{:band})
-    @test !hasdim(unmerged, Dim{:geometry})
-    @test dims(unmerged) == dims(a)
-    @test size(unmerged) == size(a)
-    @test all(a .== unmerged)
-    @test all(a .== perm_unmerged)
+        # Test Unmerge
+        @test hasdim(unmerged, X)
+        @test hasdim(unmerged, Y)
+        @test hasdim(unmerged, Dim{:band})
+        @test !hasdim(unmerged, Dim{:geometry})
+        @test dims(unmerged) == dims(a)
+        @test size(unmerged) == size(a)
+        @test all(a .== unmerged)
+        @test all(a .== perm_unmerged)
+    end
 end