Skip to content

Commit

Permalink
add keyword indexing + constructor
Browse files Browse the repository at this point in the history
  • Loading branch information
Michael Abbott committed Nov 4, 2019
1 parent 032327a commit cbe4dbc
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 1 deletion.
5 changes: 5 additions & 0 deletions src/core.jl
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,11 @@ function AxisArray(A::AbstractArray{T,N}, names::NTuple{N,Symbol}, steps::NTuple
AxisArray(A, axs...)
end

# Alternative constructor, takes names as keywords:
AxisArray(A; kw...) = AxisArray(A, nt_to_axes(kw.data)...)
@generated nt_to_axes(nt::NamedTuple) =
Expr(:tuple, (:(Axis{$(QuoteNode(n))}(getfield(nt, $(QuoteNode(n))))) for n in nt.names)...)

AxisArray(A::AxisArray) = A
AxisArray(A::AxisArray, ax::Vararg{Axis, N}) where N =
AxisArray(A.data, ax..., last(Base.IteratorsMD.split(axes(A), Val(N)))...)
Expand Down
14 changes: 13 additions & 1 deletion src/indexing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,6 @@ end
@propagate_inbounds getindex_converted(A, idxs...) = A.data[idxs...]
@propagate_inbounds setindex!_converted(A, v, idxs...) = (A.data[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
Expand All @@ -155,6 +154,19 @@ function Base.reshape(A::AxisArray, ::Val{N}) where N
AxisArray(reshape(A.data, Val(N)), Base.front(axN))
end

# Keyword indexing, reconstructs the Axis{}() objects
@propagate_inbounds Base.view(A::AxisArray; kw...) =
view(A, kw_to_axes(parent(A), kw.data)...)
@propagate_inbounds Base.getindex(A::AxisArray; kw...) =
getindex(A, kw_to_axes(parent(A), kw.data)...)
@propagate_inbounds Base.setindex!(A::AxisArray, val; kw...) =
setindex!(A, val, kw_to_axes(parent(A), kw.data)...)

function kw_to_axes(A::AbstractArray, nt::NamedTuple)
length(nt) == 0 && throw(BoundsError(A, ())) # Trivial case A[] lands here
nt_to_axes(nt)
end

### Indexing along values of the axes ###

# Default axes indexing throws an error
Expand Down
5 changes: 5 additions & 0 deletions test/core.jl
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,11 @@ B = AxisArray([1 4; 2 5; 3 6], (:x, :y), (0.2, 100))
B = AxisArray([1 4; 2 5; 3 6], (:x, :y), (0.2, 100), (-3,14))
@test axisnames(B) == (:x, :y)
@test axisvalues(B) == (-3:0.2:-2.6, 14:100:114)
# Keyword constructor
C = AxisArray([1 4; 2 5; 3 6], x=10:10:30, y=[:a, :b])
@test axisnames(C) == (:x, :y)
@test axisvalues(C) == (10:10:30, [:a, :b])
@test @inferred(AxisArray(parent(C), x=1:3, y=1:2)) isa AxisArray

@test AxisArrays.HasAxes(A) == AxisArrays.HasAxes{true}()
@test AxisArrays.HasAxes([1]) == AxisArrays.HasAxes{false}()
Expand Down
9 changes: 9 additions & 0 deletions test/indexing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -302,3 +302,12 @@ a = [2, 3, 7]
@test a[idx] 6.2
aa = AxisArray(a, :x)
@test aa[idx] 6.2

# Keyword indexing
A = AxisArray([1 2; 3 4], Axis{:x}(10:10:20), Axis{:y}(["c", "d"]))
@test @inferred(A[x=1, y=1]) == 1
@test @inferred(A[x=1]) == [1, 2]
@test axisnames(A[x=1]) == (:y,)
@test @inferred(view(A, x=1)) == [1,2]
@test parent(view(A, x=1)) isa SubArray
@test @inferred(A[x=atvalue(20), y=atvalue("d")]) == 4

0 comments on commit cbe4dbc

Please sign in to comment.