Prognostic and DiagnosticVariables array-agnostic

[milankl]

A draft superceding #493 to restructure the prognostic variables.

Includes device and ArrayType in SpectralGrid

julia> spectral_grid
SpectralGrid:
├ Spectral:   T63 LowerTriangularMatrix{Complex{Float32}}, radius = 6.371e6 m
├ Grid:       96-ring OctahedralGaussianGrid{Float32}, 10944 grid points
├ Resolution: 216km (average)
├ Vertical:   2-level SigmaCoordinates
└ Device:     CPU using Array

the variables 3D/4D are now directly in prognostic variables, not nested in .layers[1].timesteps[1] or .surface

julia> tree(progn)
PrognosticVariables{Float32, Vector{Float32}, OctahedralGaussianGrid{Float32}, Matrix{ComplexF32}, Array{ComplexF32, 3}, ShallowWater}
├ trunc
...
├ vor
├ div
├ temp
├ humid
├ pres
...

the type signature for PrognosticVariables however got a bit more complicated, maybe there's ways to simplify that

Base.@kwdef struct PrognosticVariables{
    NF<:AbstractFloat,
    ArrayType2D<:AbstractArray{NF, 1},                  # 2D real type (land + ocean)
    Grid<:AbstractGridArray{NF, 1, ArrayType2D},        # Grid for those
    ArrayTypeComplex3D<:AbstractArray{Complex{NF}, 2},  # 3D complex type (surface pressure)
    ArrayTypeComplex4D<:AbstractArray{Complex{NF}, 3},  # 4D complex type (layered variables)
    M<:ModelSetup,
} <: AbstractPrognosticVariables

[milankl]

Problem, while this works (LowerTriangularMatrix on layer 1, time step 1)

julia> progn.vor[:,1,1]
65×64 LowerTriangularMatrix{ComplexF32}:
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  …

this doesn't

julia> progn.vor[:,end,1]
ERROR: BoundsError: attempt to access 2144×2×2 Array{ComplexF32, 3} at index [1:2144, 64, 1]

the end gets translated to 64 which is size(progn.vor)[2] ...

[maximilian-gelbrecht]

This is a consequence of the LTA/LTM being a subtype of AbstractArray{T,4} in this case.

What we did with the two ways of indexing is a of a hack to be fair, the LTA can't be a 3-dim and 4-dim at the same time that easily. As far as I know we could maybe redefine the behaviour of end by defining lastindex(::LowerTriangularArray) but then it wouldn't work properly anymore for the 4-dim case. I guess we have to decide if we want the end to work as expected in the case of flat indexing (lm) or matrix indexing (l, m).

[maximilian-gelbrecht]

Does this actually appear in the dynamical core? Because the 3d indexing with leading colon does cause allocations, so it should be avoided anyway, or we can introduce a non-allocating view variarant. That's possible as well.

In the end, I think we can just avoid using the end there. We have the constant N_STEPS defined which can just use in its place there. Similarly we'll always have access to the number of horizontal levels as well easily.

[maximilian-gelbrecht]

Another alternative would be to do something a bit similar to EndpointRanges.jl and define a custom iend/lmend that works for the flat lm-indexing (and only there).

[milankl]

Another alternative would be to do something a bit similar to EndpointRanges.jl and define a custom iend/lmend that works for the flat lm-indexing (and only there).

Okay that's interesting, thanks for looking into it. No we wouldn't really need to use end inside the dynamical core or if we can easily replace it, it's more generally about using LowerTriangularArray intuitively. I guess if we've clearly defined what one can and cannot do with those then all good, but if I stumble across something that I expect to work and it doesn't then I'd rather flag this for discussion then to silently rewrite the code to just make it work

I am generally not super sure about the whole two ways of indexing LowerTriangularArray, while it seems possible to make this convenient I'm wondering whether in the end it just means we get hung up on cases where it doesn't work as expected or causes some small edge case issue that's it's not worth it overall. E.g. I see several possibilities

1. support flat and matrix indexing throughout but get hung up somewhat whenever this causes issues (what we have now)
2. support only flat indexing and use ij2k manually inside the dynamical core whenever necessary
3. support only matrix indexing by storing the upper triangle explicitly but disallow setindex! there
4. support matrix indexing only for LowerTriangularMatrix in addition to flat indexing for ...Array that way we have the more user facing LowerTriangularMatrix intuitive but are consitently indexing within the dynamical core

I obviously don't want to keep changing things but I see somewhat an appeal to make our life easier with 2-4. But I'm wondering whether in the end we should first reimplement the gradients and transforms for LowerTriangularArray to make a better informed decision also on what's important for performance and GPU

[maximilian-gelbrecht]

Let's not overthink it, I would say.

We defined LowerTriangularArray{T,N} to be a subtype of AbstractArray{T,N} but store an array AbstractArray{T,N-1}. All indexing with N indices (matrix indexing) will work flawlessly due to the abstract array interface.

Additionally we defined several functions so that the flat indexing with N-1 indices works in almost all cases. I think that's fine, and in some way that's the same as if we would only support matrix indexing and use ij2k always in the dynamical core. I think it's totally to keep it as it is, and just be a bit aware of flat indexing being a bit limited (the docs should mention that!).

The alternative would be to swap the behaviour by making it a subtype of AbstractArray{T,N-1} and have the matrix indexing to be somewhat limited.

And so far, the only limitation that I can think of are those connected with size and axes. And that is because in the end, it can't be a 3-d and 4-d array at the same time. That's not that bad.

Edit: Another way to look at is also that the way we have it currently implemented we gives us the matrix style indexing and some support for flat indexing. At the same time though we can also always directly index the .data object, when in doubt about the flat indexing.

[maximilian-gelbrecht]

Another alternative would be to do something a bit similar to EndpointRanges.jl and define a custom iend/lmend that works for the flat lm-indexing (and only there).

Okay that's interesting, thanks for looking into it.

As far as I can see it, EndpointRanges does even directly allow to define those as subtypes of their types. It's a pretty small package, so it wouldn't be too bad to depend on it. I could have a look at it some day.

[maximilian-gelbrecht]

Sounds good, I could do a code review next week.

I think the set_var! and get_var methods might also just not be necessary anymore now that we have N-dim arrays instead of vectors of arrays.