diff --git a/base/abstractarray.jl b/base/abstractarray.jl index d460af2ed1632..d3e5c3e6befb0 100644 --- a/base/abstractarray.jl +++ b/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) diff --git a/base/docs/helpdb.jl b/base/docs/helpdb.jl index e562757f4a91f..aed011621ae4b 100644 --- a/base/docs/helpdb.jl +++ b/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]) diff --git a/base/exports.jl b/base/exports.jl index cfc1574f57046..7e60a58356ea3 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -548,6 +548,7 @@ export linspace, logspace, mapslices, + mapranked, max, maxabs, maxabs!, diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 56f882338ec36..d67a6732c5287 100644 --- a/base/multidimensional.jl +++ b/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}) diff --git a/test/arrayops.jl b/test/arrayops.jl index 392938a19d6ed..350f0ee114454 100644 --- a/test/arrayops.jl +++ b/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