Merge pull request #111 from invenia/logical-indexing

Support logical indexing

Author: ajkeller34

src/indexing.jl

@@ -48,6 +48,10 @@ This internal function determines the new set of axes that are constructed upon
 indexing with I.
 """
 reaxis(A::AxisArray, I::Idx...) = _reaxis(make_axes_match(axes(A), I), I)
+function reaxis(A::AxisArray, I::AbstractArray{Bool})
+    vecI = vec(I)
+    _reaxis(make_axes_match(axes(A), (vecI,)), (vecI,))
+end
 # Ensure the number of axes matches the number of indexing dimensions
 @inline make_axes_match(axs, idxs) = _make_axes_match((), axs, Base.index_ndims(idxs...))
 # Move the axes into newaxes, until we run out of both simultaneously
@@ -106,6 +110,12 @@ using Base.AbstractCartesianIndex
 # Setindex is so much simpler. Just assign it to the data:
 @propagate_inbounds Base.setindex!(A::AxisArray, v, idxs::Idx...) = (A.data[idxs...] = v)
 
+# Logical indexing
+@propagate_inbounds function Base.getindex(A::AxisArray, idx::AbstractArray{Bool})
+    AxisArray(A.data[idx], reaxis(A, idx))
+end
+@propagate_inbounds Base.setindex!(A::AxisArray, v, idx::AbstractArray{Bool}) = (A.data[idx] = v)
+
 ### Fancier indexing capabilities provided only by AxisArrays ###
 @propagate_inbounds Base.getindex(A::AxisArray, idxs...) = A[to_index(A,idxs...)...]
 @propagate_inbounds Base.setindex!(A::AxisArray, v, idxs...) = (A[to_index(A,idxs...)...] = v)

test/indexing.jl

@@ -48,6 +48,19 @@ B = B2[1:2,:]
 @test B.axes[1].val == A.axes[1].val[1:2]
 @test B.axes[2].val == 1:Base.trailingsize(A,2)
 
+# Logical indexing
+all_inds = collect(1:length(A))
+odd_inds = collect(1:2:length(A))
+@test @inferred(A[trues(A)]) == A[:] == A[all_inds]
+@test axes(A[trues(A)]) == axes(A[all_inds])
+@test @inferred(A[isodd.(A)]) == A[1:2:length(A)] == A[odd_inds]
+@test axes(A[isodd.(A)]) == axes(A[odd_inds])
+@test @inferred(A[vec(trues(A))]) == A[:] == A[all_inds]
+@test axes(A[vec(trues(A))]) == axes(A[all_inds])
+@test @inferred(A[vec(isodd.(A))]) == A[1:2:length(A)] == A[odd_inds]
+@test axes(A[vec(isodd.(A))]) == axes(A[odd_inds])
+
+
 B = AxisArray(reshape(1:15, 5,3), .1:.1:0.5, [:a, :b, :c])
 
 # Test indexing by Intervals