From b6d7703df756aa199b47717930a284f3c91f8ec8 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Tue, 14 Apr 2015 07:28:06 -0500 Subject: [PATCH 1/2] Support indexing A[I1, i, I2] for CartesianIndexes I1, I2 This is useful for algorithms that perform a calculation along a particular dimension. --- base/multidimensional.jl | 20 ++++++++++++++++++++ test/arrayops.jl | 19 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 38b8ec3d231d1..c7bfb1f6a88ea 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -60,12 +60,20 @@ end stagedfunction getindex{N}(A::Array, i::Integer, index::CartesianIndex{N}) N==0 ? :(Base.arrayref(A, i)) : :(@ncall $(N+1) Base.arrayref A d->(d == 1 ? i : index[d-1])) end +stagedfunction getindex{M,N}(A::Array, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) + :(Base.arrayref(A, $(exprs...))) +end stagedfunction setindex!{T,N}(A::Array{T}, v, index::CartesianIndex{N}) N==0 ? :(Base.arrayset(A, convert($T,v), 1)) : :(@ncall $N Base.arrayset A convert($T,v) d->index[d]) end stagedfunction setindex!{T,N}(A::Array{T}, v, i::Integer, index::CartesianIndex{N}) N==0 ? :(Base.arrayset(A, convert($T,v), i)) : :(@ncall $(N+1) Base.arrayset A convert($T,v) d->(d == 1 ? i : index[d-1])) end +stagedfunction setindex!{T,M,N}(A::Array{T}, v, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) + :(Base.arrayset(A, convert($T,v), $(exprs...))) +end stagedfunction getindex{N}(A::AbstractArray, index::CartesianIndex{N}) :(@nref $N A d->index[d]) @@ -79,6 +87,18 @@ end stagedfunction setindex!{N}(A::AbstractArray, v, i::Integer, index::CartesianIndex{N}) :((@nref $(N+1) A d->(d == 1 ? i : index[d-1])) = v) end +for AT in (AbstractVector, AbstractMatrix, AbstractArray) # nix ambiguity warning + @eval begin + stagedfunction getindex{M,N}(A::$AT, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) + :(getindex(A, $(exprs...))) + end + stagedfunction setindex!{M,N}(A::$AT, v, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) + :(setindex!(A, v, $(exprs...))) + end + end +end # arithmetic, min/max for op in (:+, :-, :min, :max) diff --git a/test/arrayops.jl b/test/arrayops.jl index a2bb56a6c9aae..d52d6eec0e8f4 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -981,6 +981,25 @@ a = ones(5,0) b = sub(a, :, :) @test mdsum(b) == 0 +a = reshape(1:60, 3, 4, 5) +@test a[CartesianIndex{3}(2,3,4)] == 44 +a[CartesianIndex{3}(2,3,3)] = -1 +@test a[CartesianIndex{3}(2,3,3)] == -1 +@test a[2,CartesianIndex{2}(3,4)] == 44 +a[1,CartesianIndex{2}(3,4)] = -2 +@test a[1,CartesianIndex{2}(3,4)] == -2 +@test a[CartesianIndex{1}(2),3,CartesianIndex{1}(4)] == 44 +a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] = -3 +@test a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] == -3 + +a = sub(zeros(3, 4, 5), :, :, :) +a[CartesianIndex{3}(2,3,3)] = -1 +@test a[CartesianIndex{3}(2,3,3)] == -1 +a[1,CartesianIndex{2}(3,4)] = -2 +@test a[1,CartesianIndex{2}(3,4)] == -2 +a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] = -3 +@test a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] == -3 + I1 = CartesianIndex((2,3,0)) I2 = CartesianIndex((-1,5,2)) @test I1 + I2 == CartesianIndex((1,8,2)) From d88532cf8c7b5979a7c17f783ac89b0009537816 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Tue, 14 Apr 2015 10:27:33 -0500 Subject: [PATCH 2/2] Generalize multi-cartesian indexing --- base/multidimensional.jl | 53 +++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 20 deletions(-) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index c7bfb1f6a88ea..30e5659858e30 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -55,51 +55,64 @@ length{I<:CartesianIndex}(::Type{I})=length(super(I)) getindex(index::CartesianIndex, i::Integer) = getfield(index, i)::Int stagedfunction getindex{N}(A::Array, index::CartesianIndex{N}) - N==0 ? :(Base.arrayref(A, 1)) : :(@ncall $N Base.arrayref A d->index[d]) + :(Base.arrayref(A, $(cartindex_exprs((index,), (:index,))...))) end stagedfunction getindex{N}(A::Array, i::Integer, index::CartesianIndex{N}) - N==0 ? :(Base.arrayref(A, i)) : :(@ncall $(N+1) Base.arrayref A d->(d == 1 ? i : index[d-1])) + :(Base.arrayref(A, $(cartindex_exprs((i, index), (:i, :index))...))) end stagedfunction getindex{M,N}(A::Array, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) - exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) - :(Base.arrayref(A, $(exprs...))) + :(Base.arrayref(A, $(cartindex_exprs((index1, i, index2), (:index1, :i, :index2))...))) end stagedfunction setindex!{T,N}(A::Array{T}, v, index::CartesianIndex{N}) - N==0 ? :(Base.arrayset(A, convert($T,v), 1)) : :(@ncall $N Base.arrayset A convert($T,v) d->index[d]) + :(Base.arrayset(A, convert($T,v), $(cartindex_exprs((index,), (:index,))...))) end stagedfunction setindex!{T,N}(A::Array{T}, v, i::Integer, index::CartesianIndex{N}) - N==0 ? :(Base.arrayset(A, convert($T,v), i)) : :(@ncall $(N+1) Base.arrayset A convert($T,v) d->(d == 1 ? i : index[d-1])) + :(Base.arrayset(A, convert($T,v), $(cartindex_exprs((i, index), (:i, :index))...))) end stagedfunction setindex!{T,M,N}(A::Array{T}, v, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) - exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) - :(Base.arrayset(A, convert($T,v), $(exprs...))) + :(Base.arrayset(A, convert($T,v), $(cartindex_exprs((index1, i, index2), (:index1, :i, :index2))...))) end stagedfunction getindex{N}(A::AbstractArray, index::CartesianIndex{N}) - :(@nref $N A d->index[d]) + :(getindex(A, $(cartindex_exprs((index,), (:index,))...))) end stagedfunction getindex{N}(A::AbstractArray, i::Integer, index::CartesianIndex{N}) - :(@nref $(N+1) A d->(d == 1 ? i : index[d-1])) + :(getindex(A, $(cartindex_exprs((i, index), (:i, :index))...))) end -stagedfunction setindex!{N}(A::AbstractArray, v, index::CartesianIndex{N}) - :((@nref $N A d->index[d]) = v) +stagedfunction setindex!{T,N}(A::AbstractArray{T}, v, index::CartesianIndex{N}) + :(setindex!(A, v, $(cartindex_exprs((index,), (:index,))...))) end -stagedfunction setindex!{N}(A::AbstractArray, v, i::Integer, index::CartesianIndex{N}) - :((@nref $(N+1) A d->(d == 1 ? i : index[d-1])) = v) +stagedfunction setindex!{T,N}(A::AbstractArray{T}, v, i::Integer, index::CartesianIndex{N}) + :(setindex!(A, v, $(cartindex_exprs((i, index), (:i, :index))...))) end for AT in (AbstractVector, AbstractMatrix, AbstractArray) # nix ambiguity warning @eval begin - stagedfunction getindex{M,N}(A::$AT, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) - exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) - :(getindex(A, $(exprs...))) + stagedfunction getindex{M,N}(A::AbstractArray, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + :(getindex(A, $(cartindex_exprs((index1, i, index2), (:index1, :i, :index2))...))) end - stagedfunction setindex!{M,N}(A::$AT, v, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) - exprs = vcat([:(index1[$d]) for d = 1:M], :i, [:(index2[$d]) for d = 1:N]) - :(setindex!(A, v, $(exprs...))) + stagedfunction setindex!{T,M,N}(A::AbstractArray{T}, v, index1::CartesianIndex{M}, i::Integer, index2::CartesianIndex{N}) + :(setindex!(A, v, $(cartindex_exprs((index1, i, index2), (:index1, :i, :index2))...))) end end end +function cartindex_exprs(indexes, syms) + exprs = Any[] + for (i,ind) in enumerate(indexes) + if ind <: Number + push!(exprs, :($(syms[i]))) + else + for j = 1:length(ind) + push!(exprs, :($(syms[i])[$j])) + end + end + end + if isempty(exprs) + push!(exprs, 1) # Handle the zero-dimensional case + end + exprs +end + # arithmetic, min/max for op in (:+, :-, :min, :max) @eval begin