Multi-symbol indexing. Closes #71. Closes #129.

Author: jonniedie

NEWS.md

@@ -1,7 +1,23 @@
 # ComponentArrays.jl NEWS
 Notes on new features (minor releases). For more details on bugfixes and non-feature-adding changes (patch releases), check out the [releases page](https://github.com/jonniedie/ComponentArrays.jl/releases).
 
+### v0.12.0
+- Multiple symbol indexing!
+  - Use either an `Array` or `Tuple` of `Symbol`s to extract multiple named components into a new `ComponentArray
+  - It's fast!
+```julia
+julia> ca = ComponentArray(a=5, b=[4, 1], c=(a=2, b=[6, 30.0]))
+ComponentVector{Float64}(a = 5.0, b = [4.0, 1.0], c = (a = 2.0, b = [6.0, 30.0]))
+
+julia> ca[(:c, :a)]
+ComponentVector{Float64}(c = (a = 2.0, b = [6.0, 30.0]), a = 5.0)
 
+julia> ca[[:c, :a]] == ca[(:c, :a)]
+true
+
+julia> @view ca[(:c, :a)]
+ComponentVector{Float64,SubArray...}(c = (a = 2.0, b = [6.0, 30.0]), a = 5.0)
+```
 ### v0.11.0
 - Calling `axes` on a `ComponentArray` returns a new `CombinedAxis` type!
   - Doing things The Right Way™!

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <[email protected]>"]
-version = "0.11.17"
+version = "0.12.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

docs/src/indexing_behavior.md

@@ -37,6 +37,26 @@ julia> ca
 ComponentVector{Int64}(a = 0, b = [0, 1])
 ```
 
+## Indexing with multiple symbols
+It is often useful to create a new `ComponentArray` with only select fields of an old one. For this reason, `ComponentArray`s can be indexed with multiple symbolic names:
+```julia
+julia> ca = ComponentArray(a=5, b=[4, 1], c=(a=2, b=[6, 30.0]))
+ComponentVector{Float64}(a = 5.0, b = [4.0, 1.0], c = (a = 2.0, b = [6.0, 30.0]))
+
+julia> ca[(:c, :a)]
+ComponentVector{Float64}(c = (a = 2.0, b = [6.0, 30.0]), a = 5.0)
+
+julia> @view ca[(:c, :a)]
+ComponentVector{Float64,SubArray...}(c = (a = 2.0, b = [6.0, 30.0]), a = 5.0)
+```
+We see here that the new `ComponentArray` has the order of the `a` and `c` fields switched according to the order they were indexed by.
+
+Multi-symbol indexing can be performed by passing either a `Tuple` or an `Array` of `Symbol`s.
+```julia
+julia> ca[[:c, :a]] == ca[(:c, :a)]
+true
+```
+
 ## Retaining component labels through index operations
 Sometimes you might want to index into a `ComponentArray` without dropping the component name. Let's look at a new example with a more deeply nested structure:
 ```julia

src/ComponentArrays.jl

@@ -6,7 +6,7 @@ using ArrayInterface
 using LinearAlgebra
 using Requires
 
-const FlatIdx = Union{Int, CartesianIndex, AbstractArray}
+const FlatIdx = Union{Integer, CartesianIndex, AbstractArray{<:Integer}}
 const FlatOrColonIdx = Union{FlatIdx, Colon}
 
 

src/axis.jl

@@ -158,6 +158,23 @@ reindex(ax::ViewAxis, offset) = ViewAxis(viewindex(ax) .+ offset, indexmap(ax))
     ComponentIndex(getproperty(IdxMap, s))
 @inline Base.getindex(::AbstractAxis{IdxMap}, ::Val{s}) where {IdxMap, s} =
     ComponentIndex(getproperty(IdxMap, s))
+function Base.getindex(ax::AbstractAxis, syms::Union{NTuple{N,Symbol}, <:AbstractArray{Symbol}}) where {N}
+    @assert allunique(syms) "Indexing symbols must all be unique. Got $syms"
+    c_inds = getindex.((ax,), syms)
+    inds = map(x->x.idx, c_inds)
+    axs = map(x->x.ax, c_inds)
+    last_index = 0
+    new_axs = map(inds, axs) do i, ax
+        first_index = last_index + 1
+        last_index = last_index + length(i)
+        _maybe_view_axis(first_index:last_index, ax)
+    end
+    new_ax = Axis(NamedTuple(syms .=> new_axs))
+    return ComponentIndex(reduce(vcat, inds), new_ax)
+end
+
+_maybe_view_axis(inds, ax::Axis) = ViewAxis(inds, ax)
+_maybe_view_axis(inds, ::NullAxis) = inds[1]
 
 struct CombinedAxis{C,A} <: AbstractUnitRange{Int}
     component_axis::C

src/componentindex.jl

@@ -9,6 +9,8 @@ ComponentIndex(vax::ViewAxis{Inds,IdxMap,Ax}) where {Inds,IdxMap,Ax} = Component
 const FlatComponentIndex{Idx} = ComponentIndex{Idx, FlatAxis}
 const NullComponentIndex{Idx} = ComponentIndex{Idx, NullAxis}
 
+Base.:(==)(ci1::ComponentIndex, ci2::ComponentIndex) = ci1.idx == ci2.idx && ci1.ax == ci2.ax
+
 
 """
     KeepIndex(idx)

src/utils.jl

@@ -5,6 +5,7 @@
 toval(x::Val) = x
 toval(x) = Val(x)
 toval(x::String) = Val(Symbol(x))
+toval(x::AbstractArray{Symbol}) = Val((x...,))
 
 # Get value from Val type
 getval(::Val{x}) where x = x

test/runtests.jl

@@ -170,6 +170,11 @@ end
     @test ca[:b] == ca["b"] == ca.b
     @test ca[:c] == ca["c"] == ca.c
 
+    @test ca[(:a, :c)].c == ca[(:c, :a)].c == ca.c
+    @test ca[(:a, :c)].a isa Number
+    @test ca[[:a, :c]] == ca[(:a, :c)]
+    @test_throws AssertionError ca[(:a, :a)]
+
     @test cmat[:a, :a] == cmat["a", "a"] == 10000.0
     @test cmat[:a, :b] == cmat["a", "b"] == [400, 130]
     @test all(cmat[:c, :c] .== ComponentArray(a[4:10] .* a[4:10]', Axis(ax_c), Axis(ax_c)))
@@ -298,6 +303,7 @@ end
             @test ax[:a] == ax[1] == ComponentArrays.ComponentIndex(1, ComponentArrays.NullAxis())
             @test ax[:c] == ax[3:4] == ComponentArrays.ComponentIndex(3:4, FlatAxis())
             @test ax[:d] == ComponentArrays.ComponentIndex(5:8, Axis(a = 1:3, b = 4))
+            @test ax[(:a, :c)] == ax[[:a, :c]] == ComponentArrays.ComponentIndex([1, 3, 4], Axis(a = 1, c = 2:3))
         end
 
         @testset "KeepIndex" begin