Skip to content

Commit 0afed35

Browse files
committed
Added KeepIndex. Closes #88
1 parent aeae94b commit 0afed35

File tree

8 files changed

+157
-15
lines changed

8 files changed

+157
-15
lines changed

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "ComponentArrays"
22
uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
33
authors = ["Jonnie Diegelman <47193959+jonniedie@users.noreply.github.com>"]
4-
version = "0.10.2"
4+
version = "0.10.3"
55

66
[deps]
77
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

docs/make.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ makedocs(;
88
pages=[
99
"Home" => "index.md",
1010
"Quick Start" => "quickstart.md",
11+
"Indexing Behavior" => "indexing_behavior.md"
1112
"Examples" => [
1213
"examples/DiffEqFlux.md",
1314
"examples/adaptive_control.md",

docs/src/examples/ODE_jac.md

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,8 @@
11
# ODE with Jacobian
22

3-
This example shows how to use ComponentArrays for composing Jacobian update functions as well as ODE functions. Note using plain symbols to index into ```ComponentArrays``` is still pretty slow. Until symbolic indexing is faster, the convenience function ```fastindices``` can be used to speed up simulation. The general syntax looks like
3+
This example shows how to use ComponentArrays for composing Jacobian update functions as well as ODE functions. For most practical purposes, it is generally easier to use automatic differentiation libraries like ForwardDiff.jl, ReverseDiff.jl, or Zygote.jl for calculating Jacobians. Although those libraries all work with ComponentArrays, this is a nice way to handle it if you already have derived analytical Jacobians.
44

5-
```julia
6-
_x, _y, _z = fastindices(:x, :y, :z)
7-
D[_x,_y] = σ
8-
```
9-
10-
For more information on fastindices, see its entry in the API section.
5+
Note using plain symbols to index into `ComponentArrays` is still pretty slow. For speed, all symbolic indices should be wrapped in a `Val` like `D[Val(:x), Val(:y)]`.
116

127
```julia
138
using ComponentArrays
@@ -91,11 +86,11 @@ function composed_jac!(D, u, p, t)
9186
c = p.c
9287
@unpack lorenz, lotka = u
9388

94-
lorenz_jac!(D[:lorenz,:lorenz], lorenz, p.lorenz, t)
95-
lotka_jac!(D[:lotka,:lotka], lotka, p.lotka, t)
89+
lorenz_jac!(@view(D[:lorenz,:lorenz]), lorenz, p.lorenz, t)
90+
lotka_jac!(@view(D[:lotka,:lotka]), lotka, p.lotka, t)
9691

97-
D[:lorenz,:lotka][:y,:x] = -c
98-
D[:lotka,:lorenz][:x,:x] = c
92+
@view(D[:lorenz,:lotka])[:y,:x] = -c
93+
@view(D[:lotka,:lorenz])[:x,:x] = c
9994
return nothing
10095
end
10196

docs/src/indexing_behavior.md

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Indexing Behavior
2+
3+
## Views and slices
4+
`ComponentArray`s slice, rather than view, when indexing. This catches some people by surprise when they are trying to use indexing on `ComponentVector`s for dynamic field access. Let's look at an example. We'll make a `ComponentVector` with a nested structure.
5+
```julia
6+
julia> using ComponentArrays
7+
8+
julia> ca = ComponentArray(a=5, b=[4, 1])
9+
ComponentVector{Int64}(a = 5, b = [4, 1])
10+
```
11+
Using dot notation, we can access and change properties as if `ca` was a regular `struct` or `NamedTuple`.
12+
```julia
13+
julia> ca.b[1] = 99;
14+
15+
julia> ca.a = 22;
16+
17+
julia> ca
18+
ComponentVector{Int64}(a = 22, b = [99, 1])
19+
```
20+
Now let's try with indexing:
21+
```julia
22+
julia> ca[:b][1] = 0
23+
0
24+
25+
julia> ca[:a] = 0
26+
27+
julia> ca
28+
ComponentVector{Int64}(a = 0, b = [99, 1])
29+
```
30+
We see that the `a` field changed but the `b` field didn't. When we did `ca[:b]`, it sliced into `ca`, thus creating a copy that would not update the original when we went to set the first element to `0`. On the other hand, since the update of the `a` field calls `setindex!` which updates in-place.
31+
32+
If viewing, rather than slicing, is the desired behavior, use the `@view` macro or `view` function:
33+
```julia
34+
julia> @view(ca[:b])[1] = 0
35+
36+
julia> ca
37+
ComponentVector{Int64}(a = 0, b = [0, 1])
38+
```
39+
40+
## Retaining component labels through index operations
41+
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:
42+
```julia
43+
julia> ca = ComponentArray(a=5, b=[4, 1], c=(a=2, b=[6, 30]))
44+
ComponentVector{Int64}(a = 5, b = [4, 1], c = (a = 2, b = [6, 30]))
45+
```
46+
If we wanted to get the `b` component while keeping the name, we can use the `KeepIndex` wrapper around our index:
47+
```julia
48+
julia> ca[KeepIndex(:b)]
49+
ComponentVector{Int64}(b = [4, 1])
50+
```
51+
Now instead of just returning a plain `Vector`, this returns a `ComponentVector` that keeps the `b` name. Of course, this is still compatible with `view`s, so we could have done `@view ca[KeepIndex(:b)]` if we wanted to retain the view into the origianl.
52+
53+
Similarly, we can use plain indexes like ranges or integers and they will keep the names of any components they capture:
54+
```julia
55+
julia> ca[KeepIndex(1)]
56+
ComponentVector{Int64}(a = 5)
57+
58+
julia> ca[KeepIndex(2:3)]
59+
ComponentVector{Int64}(b = [4, 1])
60+
61+
julia> ca[KeepIndex(1:3)]
62+
ComponentVector{Int64}(a = 5, b = [4, 1])
63+
64+
julia> ca[KeepIndex(2:end)]
65+
ComponentVector{Int64}(b = [4, 1], c = (a = 2, b = [6, 30]))
66+
```
67+
But what if our range doesn't capture a full component? We can see below that using `KeepIndex` on the first five elements returns a `ComponentVector` with those elements but only the `a` and `b` names, since the `c` component wasn't fully captured.
68+
```julia
69+
julia> ca[KeepIndex(1:5)]
70+
5-element ComponentVector{Int64} with axis Axis(a = 1, b = 2:3):
71+
5
72+
4
73+
1
74+
2
75+
6
76+
```

src/ComponentArrays.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ include("axis.jl")
1717
export AbstractAxis, Axis, PartitionedAxis, ShapedAxis, ViewAxis, FlatAxis
1818

1919
include("componentindex.jl")
20+
export KeepIndex
2021

2122
include("componentarray.jl")
2223
export ComponentArray, ComponentVector, ComponentMatrix, getaxes, getdata, valkeys

src/componentindex.jl

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,38 @@ ComponentIndex(vax::ViewAxis{Inds,IdxMap,Ax}) where {Inds,IdxMap,Ax} = Component
88

99
const FlatComponentIndex{Idx} = ComponentIndex{Idx, FlatAxis}
1010
const NullComponentIndex{Idx} = ComponentIndex{Idx, NullAxis}
11+
12+
Base.getindex(A::AbstractArray, inds::ComponentIndex...) = ComponentArray(A[(i.idx for i in inds)...], Tuple(i.ax for i in inds))
13+
14+
"""
15+
KeepIndex(idx)
16+
17+
Tag an index of a `ComponentArray` to retain it's `Axis` through indexing
18+
"""
19+
struct KeepIndex{Idx} end
20+
KeepIndex(idx) = KeepIndex{idx}()
21+
KeepIndex(idx::Integer) = KeepIndex(idx:idx)
22+
23+
Base.getindex(ax::AbstractAxis, i::KeepIndex{Idx}) where {Idx} = _getindex_keep(ax, Idx)
24+
25+
_getindex_keep(ax::AbstractAxis, ::Colon) = ComponentIndex(:, ax)
26+
function _getindex_keep(ax::AbstractAxis, idx::AbstractRange)
27+
idx_map = indexmap(ax)
28+
keeps = NamedTuple(s=>x for (s,x) in pairs(idx_map) if first(viewindex(x)) in idx && last(viewindex(x)) in idx)
29+
new_ax = reindex(Axis(keeps), -first(idx)+1)
30+
return ComponentIndex(idx, new_ax)
31+
end
32+
function _getindex_keep(ax::AbstractAxis, sym::Symbol)
33+
ci = ax[sym]
34+
idx = ci.idx
35+
if idx isa Integer
36+
idx = idx:idx
37+
end
38+
if ci.ax isa NullAxis || ci.ax isa FlatAxis
39+
new_ax = Axis(NamedTuple{(sym,)}((ci.idx,)))
40+
else
41+
new_ax = Axis(NamedTuple{(sym,)}((ViewAxis(idx, ci.ax),)))
42+
end
43+
new_ax = reindex(new_ax, -first(idx)+1)
44+
return ComponentIndex(idx, new_ax)
45+
end

src/show.jl

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,9 +53,18 @@ function Base.show(io::IO, x::ComponentVector)
5353
return nothing
5454
end
5555

56-
function Base.show(io::IO, ::MIME"text/plain", x::ComponentVector)
57-
_print_type_short(io, x)
58-
show(io, x)
56+
function Base.show(io::IO, mime::MIME"text/plain", x::ComponentVector)
57+
len = length(x)
58+
ax = getaxes(x)[1]
59+
if last_index(ax) == len
60+
_print_type_short(io, x)
61+
show(io, x)
62+
else
63+
print(io, "$len-element ")
64+
_print_type_short(io, x)
65+
println(io, " with axis $ax:")
66+
Base.print_array(io, getdata(x))
67+
end
5968
return nothing
6069
end
6170

test/runtests.jl

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,31 @@ end
266266
@test A[1,:] == ComponentVector(x=ones(Int,4))
267267
end
268268

269+
@testset "KeepIndex" begin
270+
let
271+
ca = ComponentArray(a=1, b=2, c=[3, 4], d=(a=[5, 6, 7], b=8))
272+
cmat = ca * ca'
273+
274+
@test ca[KeepIndex(:a)] == ca[KeepIndex(1)] == ComponentArray(a=1)
275+
@test ca[KeepIndex(:b)] == ca[KeepIndex(2)] == ComponentArray(b=2)
276+
@test ca[KeepIndex(:c)] == ca[KeepIndex(3:4)] == ComponentArray(c=[3, 4])
277+
@test ca[KeepIndex(:d)] == ca[KeepIndex(5:8)] == ComponentArray(d=(a=[5, 6, 7], b=8))
278+
279+
@test ca[KeepIndex(1:2)] == ComponentArray(a=1, b=2)
280+
@test ca[KeepIndex(1:3)] == ComponentArray([1, 2, 3], Axis(a=1, b=2)) # Drops c axis
281+
@test ca[KeepIndex(2:5)] == ComponentArray([2, 3, 4, 5], Axis(b=1, c=2:3))
282+
@test ca[KeepIndex(3:end)] == ComponentArray(c=[3, 4], d=(a=[5, 6, 7], b=8))
283+
284+
@test ca[KeepIndex(:)] == ca
285+
286+
@test cmat[KeepIndex(:a), KeepIndex(:b)] == ComponentArray(fill(2, 1, 1), Axis(a=1), Axis(b=1))
287+
@test cmat[KeepIndex(:), KeepIndex(:c)] == ComponentArray((1:8)*(3:4)', getaxes(ca)[1], Axis(c=1:2))
288+
@test cmat[KeepIndex(2:5), 1:2] == ComponentArray((2:5)*(1:2)', Axis(b=1, c=2:3), FlatAxis())
289+
@test cmat[KeepIndex(2), KeepIndex(3)] == ComponentArray(fill(2*3, 1, 1), Axis(b=1), FlatAxis())
290+
@test cmat[KeepIndex(2), 3] == ComponentArray(b=2*3)
291+
end
292+
end
293+
269294
@testset "Similar" begin
270295
@test typeof(similar(ca)) == typeof(ca)
271296
@test typeof(similar(ca2)) == typeof(ca2)

0 commit comments

Comments
 (0)