Enable reduce_indices over AbstractArrays

.gitignore

@@ -2,3 +2,4 @@
 *.mem
 docs/build
 docs/site
+Manifest.toml

Project.toml

@@ -1,6 +1,6 @@
 name = "AxisArrays"
 uuid = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"

src/core.jl

@@ -359,6 +359,11 @@ reduced_indices(axs::Tuple{Vararg{Axis,N}}, region::Dims) where {N} =
 reduced_indices0(axs::Tuple{Vararg{Axis,N}}, region::Dims) where {N} =
     map((ax,d)->d∈region ? reduced_axis0(ax) : ax, axs, ntuple(identity, N))
 
+reduced_indices(axs::Tuple{Vararg{Axis}}, regions::AbstractArray) =
+    reduced_indices(axs, Tuple(regions))
+reduced_indices0(axs::Tuple{Vararg{Axis}}, regions::AbstractArray) =
+    reduced_indices0(axs, Tuple(regions))
+
 @inline reduced_indices(axs::Tuple{Vararg{Axis}}, region::Type{<:Axis}) =
     _reduced_indices(reduced_axis, (), region, axs...)
 @inline reduced_indices0(axs::Tuple{Vararg{Axis}}, region::Type{<:Axis}) =
@@ -427,6 +432,31 @@ function Base.map(f, As::AxisArray{T,N,D,Ax}...) where {T,N,D,Ax<:Tuple{Vararg{A
     return AxisArray(map(f, data...), As[1].axes...)
 end
 
+function Base.mapslices(f, A::AxisArray; dims)
+    new_axes = Axis[axes(A)...]
+
+    if (dims isa Integer || dims isa Vector{<:Integer})
+        for i in dims
+            ax = axes(A)[i]
+            new_axes[i] = ax(Base.OneTo(1))
+        end
+
+        return AxisArray(mapslices(f, A.data; dims=dims), new_axes...)
+    else
+        if (dims isa Axis || dims isa Type{<:Axis})
+            dims = [dims]
+        end
+
+        for ax in dims
+            i = axisdim(A, ax)
+            new_axes[i] = ax(Base.OneTo(1))
+        end
+
+        int_dims = axisdim.(Ref(A), dims)
+        return AxisArray(mapslices(f, A.data; dims=int_dims), new_axes...)
+    end
+end
+
 permutation(to::Union{AbstractVector{Int},Tuple{Int,Vararg{Int}}}, from::Symbols) = to
 
 """

test/core.jl

@@ -303,9 +303,30 @@ for C in arrays
         @test (op(C, dims=Axis{:y}())) == C1
         @test (op(C, dims=Axis{:x}())) == C2
         @test (op(C, dims=(Axis{:y}(),Axis{:x}()))) == C12
+        @test (op(C, dims=[Axis{:y}(),Axis{:x}()])) == C12
     end
 end
 
+A = AxisArray(collect(reshape(1:15,3,5)), y=["a", "b", "c"], x=[1, 2, 3, 4, 5])
+expected1 = AxisArray([6 15 24 33 42], y=Base.OneTo(1), x=[1, 2, 3, 4, 5])
+expected2 = AxisArray([35, 40, 45][:, :], y=["a", "b", "c"], x=Base.OneTo(1))
+expected12 = AxisArray([120][:, :], :y, :x)
+
+@test mapslices(sum, A; dims=1) == expected1
+@test mapslices(sum, A; dims=2) == expected2
+@test mapslices(sum, A; dims=[1, 2]) == expected12
+
+@test mapslices(sum, A; dims=Axis{:y}) == expected1
+@test mapslices(sum, A; dims=Axis{:x}) == expected2
+@test mapslices(sum, A; dims=[Axis{:y}, Axis{:x}]) == expected12
+
+@test mapslices(sum, A; dims=Axis{:y}()) == expected1
+@test mapslices(sum, A; dims=Axis{:x}()) == expected2
+@test mapslices(sum, A; dims=[Axis{:y}(), Axis{:x}()]) == expected12
+
+@test mapslices(sum, A; dims=[]) == A
+
+
 function typeof_noaxis(::AxisArray{T,N,D}) where {T,N,D}
     AxisArray{T,N,D}
 end
@@ -343,6 +364,7 @@ for op in functions  # together, cover both reduced_indices and reduced_indices0
     @test (op(C, dims=Axis{:y}())) == C1
     @test (op(C, dims=Axis{:x}())) == C2
     @test (op(C, dims=(Axis{:y}(),Axis{:x}()))) == C12
+    @test (op(C, dims=[Axis{:y}(),Axis{:x}()])) == C12
 end
 
 C = AxisArray(collect(reshape(1:15,3,5)), Axis{:y}([:a,:b,:c]), Axis{:x}(["a","b","c","d","e"]))