Additional LowerTriangularArray functionality

[milankl]

Spillover functionality from #543 that we could consider implementing or, alternatively, this is an issue to document the stuff that doesn't work but that user might expect to work

1. Indexing with trailing colon :

julia> L[1, :]
ERROR: MethodError: no method matching isless(::Int64, ::Colon)

2. Indexing with end

julia> L[:, end]
3-element Vector{Float64}:
 1.0
 1.0
 1.0

should be 0.0, 0.0, ..., 1.0 (= zeros for the upper triangle)

3. Broadcasting of vector with LowerTriangularMatrix

k = 1:32
A = randn(LowerTriangularMatrix{Complex{Float32}}, 32, 32)
A .*= k.^-2

doesn't work because A and k are technically both vectors of different sizes (560 for A vs 32 for k)

[milankl]

4. LowerTriangularMatrix as ones, zeros, rand, randn generator for any size

Currently this fails

julia> L = zeros(LowerTriangularMatrix{ComplexF32}, 5, 5, 1);
ERROR: MethodError: no method matching Vector{ComplexF32}(::Matrix{ComplexF32})

Although zeros(LowerTriangularMatrix{ComplexF32}, 5, 5) works fine. One needs to use Array for any higher dimensions, which I find confusing.

[milankl]

5. .= between Array and LowerTriangularArray

julia> L = zeros(LowerTriangularArray{ComplexF32}, 5, 5)
15-element, 5x5 LowerTriangularMatrix{ComplexF32}
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im

julia> M = randn(ComplexF32, size(L)...)
15-element Vector{ComplexF32}:
 -0.020682817f0 + 0.41503236f0im
    0.7638009f0 - 0.2483163f0im
   0.25317496f0 - 0.17696327f0im
   -1.3516215f0 + 0.16562235f0im
   0.31752998f0 + 0.42242694f0im
  -0.22476453f0 - 0.48526806f0im
    0.8692698f0 - 0.13144933f0im
   0.28967917f0 + 0.692557f0im
    0.4887116f0 - 0.14684263f0im
   -0.5528534f0 + 0.6710106f0im
   0.23252673f0 - 0.54855406f0im
 -0.066379264f0 - 0.30891386f0im
  -0.12736541f0 - 0.067103244f0im
    1.1546654f0 + 1.010668f0im
  -0.83623916f0 - 0.08629082f0im

julia> L .= M
ERROR: BoundsError
Stacktrace:
 [1] copyto!(L::LowerTriangularMatrix{ComplexF32}, M::Vector{ComplexF32})
   @ Main.SpeedyWeather.LowerTriangularMatrices ~/git/SpeedyWeather.jl/src/LowerTriangularMatrices/lower_triangular_matrix.jl:481

[maximilian-gelbrecht]

#561 adresses point 5

I don't think point 2 is realistically doable. Making end work simultaneously for vector and matrix-style indexing would require a lot of work, and I am not sure it's even possible.

[milankl]

I don't think point 2 is realistically doable. Making end work simultaneously for vector and matrix-style indexing would require a lot of work, and I am not sure it's even possible.

Yes that one is more to list the issue not to resolve it.

[maximilian-gelbrecht]

4. LowerTriangularMatrix as ones, zeros, rand, randn generator for any size

Currently this fails

julia> L = zeros(LowerTriangularMatrix{ComplexF32}, 5, 5, 1);
ERROR: MethodError: no method matching Vector{ComplexF32}(::Matrix{ComplexF32})

Although zeros(LowerTriangularMatrix{ComplexF32}, 5, 5) works fine. One needs to use Array for any higher dimensions, which I find confusing.

Do you just want the generators to be the same for LowerTriangularMatrix{ComplexF32} and LowerTriangularArray{ComplexF32}?

[maximilian-gelbrecht]

3. Broadcasting of vector with LowerTriangularMatrix

k = 1:32
A = randn(LowerTriangularMatrix{Complex{Float32}}, 32, 32)
A .*= k.^-2

doesn't work because A and k are technically both vectors of different sizes (560 for A vs 32 for k)

I am surprised this even works for regular arrays to be honest. For me it's not a well defined operation, but yeah, it does it row-wise multiplication.

I had a brief look but I couldn't find where this is implemented for Array or CuArray

[maximilian-gelbrecht]

I don't think point 2 is realistically doable. Making end work simultaneously for vector and matrix-style indexing would require a lot of work, and I am not sure it's even possible.

Yes that one is more to list the issue not to resolve it.

I added a warn note to the documentation about this in #561

[milankl]

Do you just want the generators to be the same for LowerTriangularMatrix{ComplexF32} and LowerTriangularArray{ComplexF32}?

I think so, the size should just be given my arguments. I think we have the same for <:AbstractGrid vs <:AbstractGridArray
You agree?

[maximilian-gelbrecht]

Ok, done. Also added to #561

[maximilian-gelbrecht]

I accidentally closed this with #562 being merged.

We can keep this issue, in case we still want to tackle the broadcasting issue 3)

[milankl]

Yes, or just to collect similar issues in the future, I bet there might be a few more coming in the next weeks.

[maximilian-gelbrecht]

When we have some spare time, it would great to implement more support for mixed indexing with ranges and colon. Like given a LxMxN LTA, currently you can't do L[1:10,1,:].

[milankl]

Yes, also this one is not intuitive

julia> L = randn(LowerTriangularMatrix, 5, 5)
15-element, 5x5 LowerTriangularMatrix{Float64}
 -1.27161     0.0        0.0       0.0      0.0
  0.666496   -0.788112   0.0       0.0      0.0
 -0.800069   -1.83198    0.928561  0.0      0.0
  0.0481627  -1.80951    0.953103  0.32636  0.0
  0.166488   -0.172987  -3.3219    1.77804  0.219669

julia> L[2:4, :] .= 0
3×1 view(reshape(::LowerTriangularMatrix{Float64}, 15, 1), 2:4, :) with eltype Float64:
 0.0
 0.0
 0.0

julia> L
15-element, 5x5 LowerTriangularMatrix{Float64}
 -1.27161    0.0        0.0       0.0      0.0
  0.0       -0.788112   0.0       0.0      0.0
  0.0       -1.83198    0.928561  0.0      0.0
  0.0       -1.80951    0.953103  0.32636  0.0
  0.166488  -0.172987  -3.3219    1.77804  0.219669

and actually led to a bug in our definition or the RandomWaves initial conditions

[milankl]

Note to self to write :, 2 as

@inline function Base.getindex(L::LowerTriangularMatrix, ::Colon, j::Integer)
    k1 = ij2k(j, j, L.m)
    kend = ij2k(L.m, j, L.m)
    return vcat(zeros(eltype(L), j-1), L.data[k1:kend])
end

to avoid the matrix allocation

[milankl]

#618 adds getindex for [l, m, k] with l, m integers but k::CartesianIndex

[milankl]

[:, end:-1:1] also escapes for RingGrids, see here #630 (comment)

[maximilian-gelbrecht]

[:, end:-1:1] also escapes for RingGrids, see here #630 (comment)

Addressed by #637

[milankl]

@jackleland just found this one, which doesn't work

julia> scratch = rand(LowerTriangularArray{ComplexF32}, 3, 3, 4, 2)
6×4×2 LowerTriangularArray{ComplexF32, 3, Array{ComplexF32, 3}}:
[:, :, 1] =
  0.220453+0.0892158im   0.899934+0.813884im  0.878469+0.32872im   0.413634+0.510188im
  0.599904+0.552259im    0.388676+0.103372im  0.807228+0.361473im  0.624458+0.888527im
  0.149992+0.435118im   0.0621975+0.02527im   0.381491+0.462077im  0.440262+0.91658im
  0.819786+0.249031im      0.4891+0.875671im  0.908225+0.916628im  0.659345+0.760382im
 0.0983859+0.626204im    0.992493+0.314848im  0.281608+0.631548im   0.84138+0.805144im
  0.886379+0.248672im    0.183932+0.701879im  0.653455+0.659489im  0.184522+0.145135im

[:, :, 2] =
  0.602306+0.0472334im    0.397111+0.742593im  0.978268+0.459059im  0.624498+0.93066im
  0.696847+0.487891im   0.00316775+0.511647im  0.752446+0.765074im  0.816116+0.434923im
  0.696755+0.384681im    0.0559902+0.276218im  0.537321+0.978618im  0.328573+0.843065im
  0.745481+0.840772im     0.813045+0.274017im  0.742856+0.262055im  0.691826+0.465365im
 0.0240027+0.886574im     0.817061+0.813867im   0.60937+0.660386im  0.333955+0.560799im
  0.545956+0.91619im      0.965738+0.493802im  0.315271+0.662712im  0.923401+0.808206im

julia> sum(scratch, dims=2)
ERROR: MethodError: no method matching ij2k(::Int64, ::CartesianIndex{2}, ::Int64)
The function `ij2k` exists, but no method is defined for this combination of argument types.

with rand(LowerTriangularArray{ComplexF32}, 3, 3, 4, 2) (i.e. one dimension less) it works, but also escapes

julia> scratch = rand(LowerTriangularArray{ComplexF32}, 3, 3, 4)
6×4 LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}:
  0.278973+0.446994im   0.275617+0.595798im   0.693215+0.63918im    0.458617+0.879801im
  0.995007+0.790848im   0.156811+0.5418im     0.539534+0.913675im   0.688073+0.144498im
 0.0767311+0.320612im   0.385006+0.36095im    0.804902+0.0553126im   0.31757+0.344941im
   0.12387+0.521824im   0.910962+0.799601im   0.364722+0.577429im   0.236999+0.580829im
 0.0170264+0.442208im   0.338848+0.791972im  0.0579183+0.700984im   0.561653+0.0778375im
   0.44149+0.0136747im  0.791868+0.576559im   0.303011+0.444414im   0.963527+0.795173im

julia> sum(scratch, dims=2)
6×1 Matrix{ComplexF32}:
  1.7064213f0 + 2.5617726f0im
   2.379425f0 + 2.3908212f0im
   1.584209f0 + 1.0818163f0im
  1.6365528f0 + 2.4796834f0im
 0.97544634f0 + 2.0130017f0im
  2.4998963f0 + 1.8298211f0im

[maximilian-gelbrecht]

Yeah, sum, cat etc either don't work directly or escape. I was thinking about adding cat a few days ago.

Do we want to add them, or just use sum(scratch.data, ...)?

The reason, I didn't add it so far is that

1. What dimension logic to use in this case? It could be confusing. Is the dims referring to the flat or matrix[*] indexing?
2. It might be type unstable. If one allows cat/sum along dims=1 to return a regular ArrayType and for other dims a LowerTriangularArray then we have a type instability. Solution would be to just return an error when attempting to do it over dims=1. Is this too restrictive?

[*] Doing a lot of these operation along the l/m dimension separately would also require a lot of additional logic to index everything correctly

[maximilian-gelbrecht]

#668 fixes the sum issue.