Setting all tracer Fields to be views of a 4D Array

[ctowery]

Hi developers!

I'm working on adding a finite-rate stiff chemical kinetics capability to the code that uses Strang splitting and the SciML ODE solving capabilities. After extensive testing of various possible OrdinaryDiffEq and DiffEqGPU methods, the only way to ensure I can solve a Strang-split chemical ODE system on GPU is to give an Ordinary solver all the chemical species as a single contiguous 4D Array or GPUArray (so halo cells and all, probably with no offsets). Therefore, what I'd like to do is allocate this 4D array on device first, then build the tracer fields that the NonhydrostaticModel uses by passing in 3D views to the base-level outer constructor for Fields. These views will have been converted to OffsetArray by the offset_data function. The constructor seems like it will accept this just fine, but I wanted to ask if I might trigger any previously known edge-case bugs by having the tracers not own their own underlying memory, or having OffsetArrays that index into a larger 4D base-one Array (i.e., any potential out-of-bounds issues inside of a compute kernel)?

[simone-silvestri]

Hi @ctowery. The strategy you have outlined should work just fine, provided you also include the space for the halos in the 4D array you preallocate.

[ctowery]

Excellent. Thanks guys!

I do think this kind of functionality would be broadly useful and would be happy to talk about it more. My timeline is to get something working as quickly as possible for AGU (my student @anpa5279 will be there with a WIP poster on our efforts up to the print deadline), and then as @anpa5279 starts doing some initial science work, I'll try and clean it up into a module with a similar feel to both buoyancy (since that's where the thermal EOS lives) and bgc (it may be possible to wrangle finite-rate chemical kinetics to match the bgc interface, but with the added capability of a split solver). I've been thinking of it as "SplitChemistry" or "ImplicitChemistry", as a way of marking the difference from bgc.

Very long term, this work might help point the way as to how to use SciML solvers as replacements for the custom timestepper and callback modules, since the SciML ecosystem duplicates all the current callback functionality, but with a lot more as a bonus. I know this was discussed in a thread somewhere a long time ago, and it seems like @simone-silvestri is moving closer to SciML compatibility with #4906

[glwagner]

As far as I know there aren't useful SciML solvers for fractional step / pressure correction methods or split-explicit free surface time stepping for hydrostatic models. The time-steppers we use (and the ones we want to develop in the future to make the model better) are specific to geophysical fluid dynamics. @simone-silvestri 's changes are part of a paper he is writing about a novel refined time-stepping methods for hydrostatic ocean dynamics.

It makes sense that they might be useful for fully explicit chemistry stepping though.

Hopefully "chemistry" doesn't need to be different from "biogeochemistry". If there is some aspect of the bgc interface that is not compatible with the splitting method then we should simply generalize the bgc interface I think. However note there is a fairly generic callback functionality (eg update_biogeochemical_fields) and it seems like substepping is covered by that pattern?

[ctowery]

For now I'm going to do Strang splitting ("because that's how it's done"), which means calling the chemistry before stage 1 of RK3, and then after stage 3, in a pattern of "chem for dt/2, advection for dt, chem for dt/2". For whatever reason, I've never seen a publication that operator splits chemistry solves inside RK stages without using a proper IMEX scheme (i.e., not just sub-stepping the chemistry). At the same time, the Strang splitting eliminates 1 expensive chemistry solve relative to operator-splitting inside the stages.

I point that out just to say I am modifying the RK3 step function and won't use either the tendency or update calls for biogeochemistry, which occur every stage. And then of course it initializes completely differently (aka, it needs the user to supply the memory, and it needs to perform all the ODE solver initializations) and must dispatch on the EOS fields (aka it takes T and S as additional arguments, so whether one or the other is a constant is important information).

[glwagner]

ah interesting! if you have code / a repo, it might help to share! I might be missing something, but before stage 1 and after stage 3 overlap, is that correct? Eg the loop will look like

update_state!(model) # during initialization

# Step 1 stage 1
rk3_substep! 
update_state!
# Step 1 stage 2
rk3_substep! 
update_state!
# Step 1 stage 3
rk3_substep! 
update_state!
# Step 2 stage 1
rk3_substep! 
update_state!

thus except for initialization, something in update_state! which is "after stage 3" is also "before stage 1"

model.clock.stage is available in update_biogeochemical_state! so it can be accessed if needed:

Oceananigans.jl/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl

Line 51 in db0c390

update_biogeochemical_state!(model.biogeochemistry, model)

You can build your custom biogeochemistry type to have the memory you need for ODE solver initializations and whatnot.

If the loop-logic doesn't generalize (or perhaps we need initialize_biogeochemistry!?) then we can add another hook that gets called before the tracer time-step, perhaps with the name time_step_biogeochemical_state! if appropriate

[ctowery]

The trick is that the splitting is only formally 2nd order accurate in between the half steps of chemistry, so if you just combine the "after stage 3" with "before stage 1", then any file outputs are formally first order accurate splittings. Now, how much this matters in practice is something that has to be measured on a case-by-case basis, since there can be different levels of coupling between velocity, pressure, and chemistry (e.g. for exo/endothermic chemistry, isothermal density change, or non-premixed reactions moderated by mixing). It's probably possible to intelligently apply update and callback functions with logic statements to not need to overload the RK3 method.