Skip to content

Commit

Permalink
WIP: adds AxisArray setindex! methods
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@mbauman
mbauman committed Feb 28, 2015
1 parent de5d48b commit 402fb5b
Showing 1 changed file with 42 additions and 62 deletions.
104 changes: 42 additions & 62 deletions src/indexing.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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?
Expand Down Expand Up @@ -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)
Expand All @@ -82,23 +86,45 @@ 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)
idxs[dims[i]] == :(Colon()) || return :(error("multiple indices provided on axis ", $(names[dims[i]])))
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 ###
Expand All @@ -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

0 comments on commit 402fb5b

Please sign in to comment.