RFC: Dyadic mapslices with broadcasting

base/abstractarray.jl

@@ -1227,6 +1227,73 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
     return R
 end
 
+function mapranked(f::Function, A::AbstractArray, rank::Int)
+    const dots = Ellipsis()
+    sizeA = size(A)[1+rank:end]
+    iter = CartesianRange(sizeA)
+    idx = first(iter)
+    r1 = f(A[dots,idx])
+    if !isa(r1, AbstractArray) || ndims(r1) == 0
+        R = similar([r1], sizeA...)
+    else
+        R = similar(r1, size(r1)..., sizeA...)
+    end
+
+    R[dots, idx] = r1
+    firsti = true
+    for idx in iter
+        if firsti
+            firsti = false
+        else
+            R[dots, idx] = f(A[dots,idx])
+        end
+    end
+    return R
+end
+
+function mapranked(f::Function, A::AbstractArray, rankA::Int, B::AbstractArray, rankB::Int)
+    const dots = Ellipsis()
+    const downcast = Base.IteratorsMD.downcast
+    sizeA = size(A)[1+rankA:end]
+    sizeB = size(B)[1+rankB:end]
+    if sizeA == sizeB
+        sizeR = sizeA
+    else # broadcast
+        na = length(sizeA)
+        nb = length(sizeB)
+        n = min(na,nb)
+        a = Int[sizeA[1:n]...]
+        b = Int[sizeB[1:n]...]
+        if !all((a .== b) | (a .== 1) | (b .== 1))
+            throw(DimensionMismatch("remaining dimensions do not agree and cannot be broadcasted"))
+        end
+        if na >= nb
+            sizeR = (max(a,b)..., sizeA[n+1:end]...)
+        else
+            sizeR = (max(a,b)..., sizeB[n+1:end]...)
+        end
+    end
+    iter = CartesianRange(sizeR)
+    idx = first(iter)
+    r1 = f(A[dots,downcast(idx, sizeA)], B[dots, downcast(idx, sizeB)])
+
+    if !isa(r1, AbstractArray) || ndims(r1) == 0
+        R = similar([r1], sizeR...)
+    else
+        R = similar(r1, size(r1)..., sizeR...)
+    end
+
+    R[dots, idx] = r1
+    firsti = true
+    for idx in iter
+        if firsti
+            firsti = false
+        else
+            R[dots, idx] = f(A[dots,downcast(idx, sizeA)], B[dots, downcast(idx, sizeB)])
+        end
+    end
+    return R
+end
 
 # using promote_type
 function promote_to!{T,F}(f::F, offs, dest::AbstractArray{T}, A::AbstractArray)

base/docs/helpdb.jl

@@ -2100,6 +2100,13 @@ Transform the given dimensions of array `A` using function `f`. `f` is called on
 """
 mapslices
 
+doc"""
+    mapranked(f, A, r, [B, s])
+
+Transform the leading subarrays or slices of rank (number of dimensions) `r` respective `s` of the arrays `A` and `B` using the two-argument function `f`. For example, if `r` is 2 and `A` is 4-dimensional and `s` is 1 and `B` is 3-dimensional, then `f` is called on `(A[:,:,i,j], B[:, i,j])` for all `i` and `j`. The results are concatenated along the remaining dimensions. If the dimensions of the arrays of `A`-slices or rank `r` and `B`-slices of rank `s` disagree, the function broadcasts these arrays to a common size by expanding singleton dimensions. If `B` and `s` are omitted, call `f` on the `r`-slices of `A`.
+"""
+mapranked
+
 doc"""
     spdiagm(B, d[, m, n])
 

base/exports.jl

@@ -548,6 +548,7 @@ export
     linspace,
     logspace,
     mapslices,
+    mapranked,
     max,
     maxabs,
     maxabs!,

base/multidimensional.jl

@@ -53,6 +53,8 @@ end
 end
 *(index::CartesianIndex,a::Integer)=*(a,index)
 
+@generated function downcast{N,M}(x::CartesianIndex{N}, r::NTuple{M, Int}) :(CartesianIndex{M}(tuple($([:(r[$i] == 1 ? 1 : x[$i]) for i in 1:M]...)))) end
+
 # Iteration
 immutable CartesianRange{I<:CartesianIndex}
     start::I
@@ -149,6 +151,7 @@ simd_index{I<:CartesianIndex{0}}(iter::CartesianRange{I}, ::CartesianIndex, I1::
     :($meta; CartesianIndex{$(N+1)}($(args...)))
 end
 
+
 end  # IteratorsMD
 
 using .IteratorsMD
@@ -379,6 +382,18 @@ end
 end
 
 
+## Ellipsis() represents a number of leading Colon()'s implied by the number of dimensions of the AbstractArray and the CartesianIndex
+
+type Ellipsis
+end
+
+
+@generated _setindex!{T,N,M}(l::LinearIndexing, A::AbstractArray{T,N}, x, ::Ellipsis, idx::Base.IteratorsMD.CartesianIndex{M}) = :($(Expr(:meta, :inline)); setindex!(A, x, $([Colon(:) for i in 1:N-M]...), idx))
+@generated _unsafe_setindex!{T,N,M}(l::LinearIndexing, A::AbstractArray{T,N}, x, ::Ellipsis, idx::Base.IteratorsMD.CartesianIndex{M}) = :($(Expr(:meta, :inline)); unsafe_setindex!(A, x, $([Colon(:) for i in 1:N-M]...), idx))
+@generated _getindex{T,N,M}(l::LinearIndexing, A::AbstractArray{T,N}, ::Ellipsis, idx::Base.IteratorsMD.CartesianIndex{M}) = :($(Expr(:meta, :inline)); getindex(A, $([Colon(:) for i in 1:N-M]...), idx))
+@generated _unsafe_getindex{T,N,M}(l::LinearIndexing, A::AbstractArray{T,N}, ::Ellipsis, idx::Base.IteratorsMD.CartesianIndex{M}) = :($(Expr(:meta, :inline)); unsafe_getindex(A, $([Colon(:) for i in 1:N-M]...), idx))
+
+
 ##
 
 @generated function findn{T,N}(A::AbstractArray{T,N})

test/arrayops.jl

@@ -696,6 +696,30 @@ let
     @test size(n1) == (6,1,4) && size(n2) == (6,3,1)  && size(n3) == (2,6,1)
 end
 
+# mapranked
+let
+
+    local a, b
+
+    b = mapranked((x,y) -> sum(x)+sum(y), ones(2,3,4), 1, ones(2,3,4),1)
+    @test size(b) === (3,4)
+    @test all(b.== 4)
+    b = mapranked((x,y) -> sum(x)+sum(y), ones(2,3,4), 2, ones(2,3,4),2)
+    @test size(b) === (4,)
+    @test all(b.==12)
+    a = rand(5,5)
+    @test mapranked(x-> maximum(-x), a, 0) == -a
+    @test mapranked(x-> maximum(-x), a, 1) == vec(maximum(-a,1))
+
+    b = rand(5)
+    @test mapranked(*,reverse(b)',0,b,0) == b*reverse(b)'
+    @test mapranked(*,reverse(b),0,b',0) == reverse(b)*b'
+
+    # other types than Number
+    @test mapranked((x,y)-> prod(x)*prod(y), ["1" "2"; "3" "4"], 1, ["a" "b"; "c" "d"], 1) == ["13ac", "24bd"]
+    @test mapranked(prod,["1"],1) == ["1"]
+
+end
 
 # single multidimensional index
 let