jonniedie / ComponentArrays.jl

Arrays with arbitrarily nested named components.
MIT License
291 stars 35 forks source link

How to efficiently access first element of all subarrays, i.e. [:,1] or explicitly [(:SWATI, :GWAT),:amt] #165

Open fabern opened 2 years ago

fabern commented 2 years ago

Hi, I was wondering if there ist an efficient way to access the first element of all the subarrays in a ComponentVector of ComponentVectors?

using BenchmarkTools
using ComponentArrays #, LabelledArrays, RecursiveArrayTools

SWATI=ComponentArray(amt = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], conc = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.10])
GWAT =ComponentArray(amt = 1.0, conc = 3.0)
u0 = ComponentArray(SWATI = SWATI, GWAT = GWAT)

@btime sum(reduce(vcat,[u0.SWATI.amt, u0.GWAT.amt])) # 5.175 μs (54 allocations: 1.98 KiB)
@btime sum(vcat(u0.SWATI.amt, u0.GWAT.amt))          # 3.912 μs (48 allocations: 1.72 KiB)

Or would this require defining the ComponentVector with permuted dimensions? Further, I saw higher dimensional ComponentArrays mentioned in the README.md, but I have not seen any example how to generate them in the documentation. Could you provide an example of this, please?

jonniedie commented 1 year ago

There's a function called valkeys that can type-stabily create the keys of a ComponentArray wrapped in Vals for efficient iteration. Then you can just map over the keys and index into the ComponentArray. This will get you a Tuple of the values and you can sum over the inner levels then outer, if you want.

using BenchmarkTools, ComponentArrays

SWATI = ComponentArray(amt = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], conc = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.10])
GWAT = ComponentArray(amt = 1.0, conc = 3.0)
u0 = ComponentArray(SWATI = SWATI, GWAT = GWAT)
julia> map(k -> view(u0, k).amt, valkeys(u0))
([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], 1.0)

julia> @btime sum(map(k -> sum(view($u0, k).amt), valkeys($u0)))
  4.333 ns (0 allocations: 0 bytes)
56.0

Note also that I'm using view(u0, k) here instead of u0[k] to avoid the allocation (getindex on ComponentArrays slices, unlike getproperty, which views).

jonniedie commented 1 year ago

As for the multidimensional ComponentArrays, that wouldn't work here because the number of elements of each of the subcomponents are different. If they weren't jagged like that, you could make a ComponentMatrix by just concatenating the ComponentVectors together

julia> SWATI = ComponentArray(amt = [1.0, 2.0], conc = [0.1, 0.2]);

julia> GWAT = ComponentArray(amt = [3.0, 4.0], conc = [0.3, 0.4]);

julia> u0 = [SWATI GWAT]
4×2 ComponentMatrix{Float64} with axes Axis(amt = 1:2, conc = 3:4) × FlatAxis()
 1.0  3.0
 2.0  4.0
 0.1  0.3
 0.2  0.4

julia> u0[:amt, :]
2×2 Matrix{Float64}:
 1.0  3.0
 2.0  4.0

julia> @btime sum(@view $u0[:amt, :]) # Also using @view here
  6.166 ns (0 allocations: 0 bytes)
10.0
fabern commented 1 year ago

Thanks for this helpful answer, Jonnie!

The first part solves my problem perfectly well.

For the second part regarding multidimensional ComponentMatrix: in your example the second axis is simply FlatAxis(). Is there a possibility to have both axes with named entires, e.g. allowing to access u0[:amt, :SWATI] ?