Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP: adds AxisArray setindex! methods #11

Closed
wants to merge 1 commit into from
Closed

WIP: adds AxisArray setindex! methods #11

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 42 additions & 62 deletions src/indexing.jl
View file
Open in desktop
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