Indexing `ComponentMatrix` with `FlatAxis` components

[dingraha]

When i try the following code:

using ComponentArrays: ComponentVector

function try1()
    ni = 4
    nj = 2
    nk = 3

    X = ComponentVector(a=0.0, b=zeros(Float64, nk), c=zeros(Float64, ni, nj))
    Y = ComponentVector(d=zeros(Float64, ni, nj))
    J = Y.*X'

    @show J[:d, :a] # works!
    @show J[:d, :c] # works!
    @show J[:d, :b] # doesn't work :-(

    return
end

I see this error:

J[:d, :a] = [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
J[:d, :c] = [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0
 0.0;;; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
ERROR: DimensionMismatch: new dimensions (4, 2) must be consistent with array size 24
Stacktrace:
  [1] (::Base.var"#throw_dmrsa#328")(dims::Tuple{Int64, Int64}, len::Int64)
    @ Base ./reshapedarray.jl:41
  [2] reshape
    @ Base ./reshapedarray.jl:45 [inlined]
  [3] maybe_reshape
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/componentarray.jl:237 [inlined]
  [4] ComponentArray
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/componentarray.jl:52 [inlined]
  [5] macro expansion
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/array_interface.jl:0 [inlined]
  [6] _getindex(::typeof(getindex), ::ComponentMatrix{Float64, Matrix{Float64}, Tuple{Axis{(d = ViewAxis(1:8, ShapedAxis((4, 2))),)}, Axis{(a = 1, b = 2:4, c = ViewAxis(5:12, ShapedAxis((4, 2))))}}}, ::Val{:d}, ::Val{:b})
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/array_interface.jl:119
  [7] getindex
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/array_interface.jl:103 [inlined]
  [8] getindex(::ComponentMatrix{Float64, Matrix{Float64}, Tuple{Axis{(d = ViewAxis(1:8, ShapedAxis((4, 2))),)}, Axis{(a = 1, b = 2:4, c = ViewAxis(5:12, ShapedAxis((4, 2))))}}}, ::Symbol, ::Symbol)
    @ ComponentArrays ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/dev/ComponentArrays/src/array_interface.jl:102
  [9] macro expansion
    @ ./show.jl:1181 [inlined]
 [10] try1()
    @ Main.CABug ~/projects/sr20_v2/dev/electric_sr20_acoustics/electric_sr20_acoustics/models/prop/ElectricSR20OMComponents/scripts/04_ca_bug.jl:16
 [11] top-level scope
    @ REPL[55]:1

It looks like maybe_reshape is ignoring the length of the b component in J, I guess because it's a FlatAxis, not a ShapedAxis. If I make b have a size of, say (1, nk) it works as expected.

Any ideas?

[dingraha]

I can actually fix this by commenting out this line, which creates a FlatAxis when constructing a ShapedAxis for a 1D array. That appears to preserve the necessary shape information that's needed by maybe_reshape. But that change breaks other things, unfortunately. For example, indexing a matrix with nested component axis like here no longer works, since the presence of the ShapedAxis triggers calling maybe_reshape, which can't figure out the necessary size of the nested component axis.

[dingraha]

OK, I attempted to fix this issue by introducing a Shaped1DAxis type in this PR: https://github.com/jonniedie/ComponentArrays.jl/pull/249. It fixes the issue described here, but breaks some relatively minor things (IMHO). The most serious issue appears to be indexing a ComponentIndex with multiple symbols: https://github.com/dingraha/ComponentArrays.jl/blob/0db92aec69fac664524cc95a8f5fa1d3a605ce93/test/runtests.jl#L362, but that already didn't work with multi-dimensional ShapedAxis.

[dingraha]

Now almost everything works with https://github.com/jonniedie/ComponentArrays.jl/pull/249. The things that are broken:

• Creating a ComponentArray with an empty NamedTuple is not identical to creating one with an empty Array, like this, as seen here:

ComponentArray(a = T[], b = (;)) == ComponentVector{T}(a = T[], b = T[])

• As mentioned above, indexing a ComponentIndex with multiple symbols doesn't work when one of the indices refers to a ShapedAxis, as seen here. IIUC that used to work when the index was one-dimensional (and was thus converted into a FlatAxis, but did not work for multi-dimensional ShapedAxis, and still doesn't.

[dingraha]

Managed to fix the test for indexing a ComponentIndex with multiple symbols. Now the only broken test in https://github.com/jonniedie/ComponentArrays.jl/pull/249 is the "empty NamedTuple vs empty Array one, which is relatively minor IMHO.