From 402fb5b83d568dc12f2f85a97ce3dc90831c2198 Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Sat, 28 Feb 2015 14:39:34 -0500 Subject: [PATCH] WIP: adds AxisArray setindex! methods Requires JuliaLang/julia#10331. But I think it is also running into a MAX_TUPLE_LEN issue, which is messing up my workaround for JuliaLang/julia#10191. Needs tests, too. --- src/indexing.jl | 104 +++++++++++++++++++----------------------------- 1 file changed, 42 insertions(+), 62 deletions(-) diff --git a/src/indexing.jl b/src/indexing.jl index a78e0ac..d691785 100644 --- a/src/indexing.jl +++ b/src/indexing.jl @@ -22,10 +22,12 @@ Base.getindex{T,D,names,Ax}(A::AxisArray{T,1,D,names,Ax}, idx::Colon) = A # Linear indexing with an array Base.getindex{T,N,D,names,Ax,S<:Int}(A::AxisArray{T,N,D,names,Ax}, idx::AbstractArray{S}) = A.data[idx] +Base.setindex!{T,N,D,names,Ax,S<:Int}(A::AxisArray{T,N,D,names,Ax}, v, idx::AbstractArray{S}) = (A.data[idx] = v) # Cartesian iteration Base.eachindex(A::AxisArray) = eachindex(A.data) Base.getindex(A::AxisArray, idx::Base.IteratorsMD.CartesianIndex) = A.data[idx] +Base.setindex!(A::AxisArray, v, idx::Base.IteratorsMD.CartesianIndex) = (A.data[idx] = v) # More complicated cases where we must create a subindexed AxisArray # TODO: do we want to be dogmatic about using views? For the data? For the axes? @@ -60,6 +62,8 @@ stagedfunction Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, idxs: $(AxisArray{T,newdims,newdata,newnames,newaxes})(data, $axes) end end +# Setindex is so much simpler. Just assign it to the data: +Base.setindex!{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, v, idxs::Idx...) = (A.data[idxs...] = v) # Stolen and stripped down from the Base stagedfunction _sub: function _sub_type(A, I) @@ -82,15 +86,36 @@ function _sub_type(A, I) SubArray{T,N,A,It,LD} end - ### Fancier indexing capabilities provided only by AxisArrays ### -# First is the ability to index by named axis. +# Defining the fallbacks on get/setindex are tricky due to ambiguities with +# AbstractArray definitions... but they simply punt to to_index to convert the +# special indexing forms to integers and integer ranges. +# Even though all these splats look scary, they get inlined and don't allocate. +Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, idx::AbstractArray) = A[to_index(A,idx)...] +Base.setindex!{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, v, idx::AbstractArray) = (A[to_index(A,idx)...] = v) +let rargs = Expr[], aargs = Expr[], idxs = Symbol[] + for i = 1:4 + isym = symbol("i$i") + push!(rargs, :($isym::Real)) + push!(aargs, :($isym::Any)) + push!(idxs, isym) + @eval Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, $(rargs...)) = A[to_index(A,$(idxs...))...] + @eval Base.setindex!{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, v, $(rargs...)) = (A[to_index(A,$(idxs...))...] = v) + @eval Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, $(aargs...)) = A[to_index(A,$(idxs...))...] + @eval Base.setindex!{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, v, $(aargs...)) = (A[to_index(A,$(idxs...))...] = v) + end +end +Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, idxs...) = A[to_index(A,idxs...)...] +Base.setindex!{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, v, idxs...) = (A[to_index(A,idxs...)...] = v) + + +# First is indexing by named axis. We simply sort the axes and re-dispatch. # When indexing by named axis the shapes of omitted dimensions are preserved # TODO: should we handle multidimensional Axis indexes? It could be interpreted # as adding dimensions in the middle of an AxisArray. # TODO: should we allow repeated axes? As a union of indices of the duplicates? -stagedfunction Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I::Axis...) +stagedfunction to_index{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I::Axis...) dims = Int[axisdim(A, ax) for ax in I] idxs = Expr[:(Colon()) for d = 1:N] for i=1:length(dims) @@ -98,7 +123,8 @@ stagedfunction Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I::Ax idxs[dims[i]] = :(I[$i].I) end - return :(A[$(idxs...)]) + meta = Expr(:meta, :inline) + return :($meta; to_index(A, $(idxs...))) end ### Indexing along values of the axes ### @@ -121,70 +147,24 @@ function axisindexes{T}(::Type{Categorical}, ax::AbstractVector{T}, idx::Abstrac res end -# Defining the fallbacks on getindex are tricky due to ambiguities with -# AbstractArray definitions - -let args = Expr[], idxs = Symbol[] - for i = 1:4 - isym = symbol("i$i") - push!(args, :($isym::Real)) - push!(idxs, isym) - @eval Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, $(args...)) = fallback_getindex(A, $(idxs...)) - end -end -Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, idx::AbstractArray) = fallback_getindex(A, idx) -Base.getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, idxs...) = fallback_getindex(A, idxs...) - # These catch-all methods attempt to convert any axis-specific non-standard -# indexing types to their integer or integer range equivalents using the +# indexing types to their integer or integer range equivalents using axisindexes # They are separate from the `Base.getindex` function to help alleviate # ambiguity warnings from, e.g., `getindex(::AbstractArray, ::Real...)`. -# TODO: These could be generated with meta-meta-programming -stagedfunction fallback_getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I1) - ex = :(getindex(A)) - push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1))) - for _=2:N - push!(ex.args, :(Colon())) - end - ex -end -stagedfunction fallback_getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I1, I2) - ex = :(getindex(A)) - push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1))) - push!(ex.args, I2 <: Idx || length(Ax) < 2 ? :(I2) : :(axisindexes(A.axes[2], I2))) - for _=3:N - push!(ex.args, :(Colon())) - end - ex -end -stagedfunction fallback_getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I1, I2, I3) - ex = :(getindex(A)) - push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1))) - push!(ex.args, I2 <: Idx || length(Ax) < 2 ? :(I2) : :(axisindexes(A.axes[2], I2))) - push!(ex.args, I3 <: Idx || length(Ax) < 3 ? :(I3) : :(axisindexes(A.axes[3], I3))) - for _=4:N - push!(ex.args, :(Colon())) - end - ex -end -stagedfunction fallback_getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I1, I2, I3, I4) - ex = :(getindex(A)) - push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1))) - push!(ex.args, I2 <: Idx || length(Ax) < 2 ? :(I2) : :(axisindexes(A.axes[2], I2))) - push!(ex.args, I3 <: Idx || length(Ax) < 3 ? :(I3) : :(axisindexes(A.axes[3], I3))) - push!(ex.args, I4 <: Idx || length(Ax) < 4 ? :(I4) : :(axisindexes(A.axes[4], I4))) - for _=5:N - push!(ex.args, :(Colon())) - end - ex -end -stagedfunction fallback_getindex{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I...) - ex = :(getindex(A)) +stagedfunction to_index{T,N,D,names,Ax}(A::AxisArray{T,N,D,names,Ax}, I...) + ex = Expr(:tuple) for i=1:length(I) - push!(ex.args, I[i] <: Idx || length(Ax) < i ? :(I[$i]) : :(axisindexes(A.axes[$i], I[$i]))) + if I[i] <: Idx + push!(ex.args, :(I[$i])) + elseif i <= length(Ax) + push!(ex.args, :(axisindexes(A.axes[$i], I[$i]))) + else + push!(ex.args, :(error("dimension ", $i, " does not have an axis to index"))) + end end for _=length(I)+1:N push!(ex.args, :(Colon())) end - ex + meta = Expr(:meta, :inline) + return :($meta; $ex) end -