Fix mapslices (#553)

* fix mapslices

* add some tests

* add import and move to proper location

* Require DIskArrays 0.4.18 and use dropdims

* add dropdims keyword to mapslices

Author: meggart

Project.toml

@@ -38,7 +38,7 @@ DataStructures = "0.17, 0.18, 0.19"
 DimensionalData = "0.27, 0.28, 0.29"
 DiskArrayEngine = "0.2"
 DiskArrayTools = "0.1.12"
-DiskArrays = "0.3, 0.4.10"
+DiskArrays = "0.4.18"
 DocStringExtensions = "0.8, 0.9"
 Glob = "1.3"
 Interpolations = "0.12, 0.13, 0.14, 0.15, 0.16"

src/DAT/xmap.jl

@@ -406,13 +406,20 @@ compute_to_zarr(ds, "output.zarr")
 function xmap end
 
 import Base.mapslices
-function mapslices(f, d::YAXArray, addargs...; dims, kwargs...)
+function mapslices(f, d::YAXArray, addargs...; dims, dropdims=false, kwargs...)
     !isa(dims, Tuple) && (dims = (dims,))
     dw = map(dims) do d
         Symbol(d)=>Whole()
     end
     w = windows(d,dw...)
-    xmap(f,w,inplace=false)    
+    outaxes = YAXArrays.getOutAxis((YAXArrays.ByInference(),), Symbol.(dims), (d,), addargs, f)
+
+    returncube = xmap(f, w, output=XOutput(outaxes...), inplace=false)
+    if dropdims
+        return Base.dropdims(returncube; dims=Symbol.(dims))
+    else
+        return returncube
+    end
 end
 
 struct XFunction{F,O,I} <: Function

src/helpers.jl

@@ -78,19 +78,21 @@ match_axis(a, ax) = match_axis(get_descriptor(a), ax)
 getOutAxis(desc, axlist, incubes, pargs, f) = getAxis(desc, unique(axlist))
 
 function getOutAxis(desc::Tuple{ByInference}, axlist, incubes, pargs, f)
-    inAxes = map(DD.dims, incubes)
-    inAxSmall = map(i -> filter(j -> in(j, axlist), i) |> collect, inAxes)
+    axlist = map(axlist) do ax
+        isa(ax, String) ? Symbol(ax) : ax
+    end
+    inAxSmall = map(i -> DD.dims(i, axlist), incubes)
     inSizes = map(i -> (map(length, i)...,), inAxSmall)
     intypes = map(eltype, incubes)
     testars = map((s, it) -> zeros(it, s...), inSizes, intypes)
-    map(testars) do ta
-        ta .= rand(Base.nonmissingtype(eltype(ta)), size(ta)...)
-        if eltype(ta) >: Missing
-            # Add some missings
-            randind = rand(1:length(ta), length(ta) ÷ 10)
-            ta[randind] .= missing
-        end
-    end
+    # map(testars) do ta
+    #     ta .= rand(Base.nonmissingtype(eltype(ta)), size(ta)...)
+    #     if eltype(ta) >: Missing
+    #         # Add some missings
+    #         randind = rand(1:length(ta), length(ta) ÷ 10)
+    #         ta[randind] .= missing
+    #     end
+    # end
     resu = f(testars..., pargs...)
     isa(resu, AbstractArray) ||
         isa(resu, Number) ||
@@ -102,16 +104,17 @@ function getOutAxis(desc::Tuple{ByInference}, axlist, incubes, pargs, f)
         end
     end
     outsizes = size(resu)
+    allAxes = reduce(union, inAxSmall)
     outaxes = map(outsizes, 1:length(outsizes)) do s, il
         if s > 2
-            i = findall(i -> i == s, length.(axlist))
+            i = findall(i -> i == s, length.(allAxes))
             if length(i) == 1
-                return axlist[i[1]]
+                return allAxes[i[1]]
             elseif length(i) > 1
                 @info "Found multiple matching axes for output dimension $il"
             end
         end
-        return Dim{Symbol("OutAxis$(il)")}( 1:s)
+        return Dim{Symbol("OutAxis$(il)")}(1:s)
     end
     if !allunique(outaxes)
         #TODO: fallback with axis renaming in this case

test/Datasets/datasets.jl

@@ -528,6 +528,30 @@ end
     mean_slice = mapslices(mean, cube; dims="Dim_1")
 
     @test mean_slice[1, :, :] == ones(20, 5)
+    import DimensionalData as DD
+    a = YAXArray(reshape(1:1000, 10, 20, 5))
+    b = mapslices(cumsum, a, dims="Dim_1")
+    @test size(b) == size(a)
+    @test DD.dims(b) == DD.dims(a)
+    @test b[3, 1, 1] == 6
+    @test b[2, 2, :].data == 23:400:1623
+
+    c = mapslices(sum, a, dims="Dim_2")
+    @test size(c) == (1, 10, 5)
+    @test c.Dim_2 == DD.rebuild(a.Dim_2, [a.Dim_2.val])
+    @test c.Dim_1 == a.Dim_1
+    @test c.Dim_3 == a.Dim_3
+    @test c[1, 3, 3] == 9960
+
+    #Test dropdims as well
+    d = dropdims(c, dims=:Dim_2)
+    DD.dims(d) == Base.tail(DD.dims(c))
+
+    d = mapslices(sum, a, dims="Dim_2", dropdims=true)
+    @test size(d) == (10, 5)
+    @test d.Dim_1 == a.Dim_1
+    @test d.Dim_3 == a.Dim_3
+    @test d[3, 3] == 9960
 end
 
 @testset "Making Cubes from heterogemous data types" begin