Open milankl opened 2 months ago
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]
...
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).
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.
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).
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
ij2k
manually inside the dynamical core whenever necessarysetindex!
thereLowerTriangularMatrix
in addition to flat indexing for ...Array
that way we have the more user facing LowerTriangularMatrix
intuitive but are consitently indexing within the dynamical coreI 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
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.
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 with the last commits I've redefined all prognostic and diagnostic variables similar to
@kwdef struct PrognosticVariables{
NF, # <: AbstractFloat
ArrayType, # Array, CuArray, ...
SpectralVariable3D, # <: LowerTriangularArray
SpectralVariable4D, # <: LowerTriangularArray
GridVariable2D, # <: AbstractGridArray
ParticleVector, # <: AbstractVector{Particle{NF}}
} <: AbstractPrognosticVariables
...
vor::SpectralVariable4D = zeros(SpectralVariable4D, trunc+2, trunc+1, nlayers, nsteps)
pres::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nsteps)
...
end
That means all variable structs have NF
as the first parameter and the array type (Array
, JLArray
, CuArray
, ...) as the second for dispatch. The other parameters are then just the types of the n-dimensional spectral/grid arrays within to have concrete types throughout.
I've removed all the layers/timesteps and they are now structured as follows
julia> progn = PrognosticVariables(spectral_grid)
PrognosticVariables{Float32, Array}
├ vor: T31, 8-layer, 2-steps LowerTriangularArray{Float32}
├ div: T31, 8-layer, 2-steps LowerTriangularArray{Float32}
├ temp: T31, 8-layer, 2-steps LowerTriangularArray{Float32}
├ humid: T31, 8-layer, 2-steps LowerTriangularArray{Float32}
├ pres: T31, 1-layer, 2-steps LowerTriangularArray{Float32}
├┐ocean: PrognosticVariablesOcean{Float32}
│├ sea_surface_temperature: 48-ring OctahedralGaussianGrid{Float32}
│└ sea_ice_concentration: 48-ring OctahedralGaussianGrid{Float32}
├┐land: PrognosticVariablesLand{Float32}
│├ land_surface_temperature: 48-ring OctahedralGaussianGrid{Float32}
│├ soil_moisture_layer1: 48-ring OctahedralGaussianGrid{Float32}
│└ soil_moisture_layer2: 48-ring OctahedralGaussianGrid{Float32}
├ particles: 0-element Vector{Particle{Float32}}
├ scale: 1.0
└ clock: 2000-01-01T00:00:00
julia> diagn = DiagnosticVariables(spectral_grid)
DiagnosticVariables{Float32, Array, OctahedralGaussianGrid}
├ resolution: T31, 48 rings, 8 layers, 0 particles
├ tendencies::Tendencies
├ grid::GridVariables
├ dynamics::DynamicsVariables
├ physics::PhysicsVariables
├ particles::ParticleVariables
├ columns::Vector{ColumnVariables}
└ scale: 1.0
julia> diagn.tendencies
Tendencies
├ vor_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ div_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ temp_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ humid_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ u_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ v_tend: T31, 8-layer LowerTriangularArray{ComplexF32}
├ pres_tend: T31, 1-layer LowerTriangularArray{ComplexF32}
├ u_tend_grid: 48-ring, 8-layer OctahedralGaussianArray{Float32}
├ v_tend_grid: 48-ring, 8-layer OctahedralGaussianArray{Float32}
├ temp_tend_grid: 48-ring, 8-layer OctahedralGaussianArray{Float32}
├ humid_tend_grid: 48-ring, 8-layer OctahedralGaussianArray{Float32}
└ pres_tend_grid: 48-ring, 1-layer OctahedralGaussianArray{Float32}
I believe this is what we want, but let me know if you have any remarks/requests/comments! Unfortunately, this also means that your set_var!
functionality has to be completely overhauled, maybe we can do eventually something similar to Oceananigans
vor₀(φ,θ,σ) =
set!(prognostic_variables, vor=vor₀)
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.
A draft superceding #493 to restructure the prognostic variables.
Includes
device
andArrayType
inSpectralGrid
the variables 3D/4D are now directly in prognostic variables, not nested in
.layers[1].timesteps[1]
or.surface
the type signature for
PrognosticVariables
however got a bit more complicated, maybe there's ways to simplify that