Adds split timestepping for physics and BGC

[jagoosw]

This is a WIP that I started doing, but then ran out of time but thought it might be of interest to others at some point.

This PR would add a "Strange splitting" time stepper where the biogeochemical components of the tendencies are computed more frequently than the transport since biogeochemistry can often be much stiffer and so it is relatively common practice to step the bgc at different frequencies to the physics (e.g. NEMO-PISCES allows multiple euler substeps to be taken between each physics step).

Strange splitting assumes the tendency can be written as a stiff and non-stiff part:
$\frac{\partial C}{\partial t} = \mathcal{A} + \mathcal{B}$,
where $\mathcal{A}$ is the advection (less stiff) and $\mathcal{B}$ is the BGC (more stiff) components.
And we basically take half a step with just the stiff component, then a full step with just the non-stiff component, and finally another step with the stiff component.
i.e. we have $C^n$ then:
$C^{n+1/2} = C^n + \int_{t_n}^{t_n+\Delta t/2}\mathcal{B}\left(C^n\right)dt,$
$C^* = C^{n+1/2} + \int_{t_n}^{t_n+\Delta t}\mathcal{A}\left(C^{n+1/2}, \vec{u^n}\right) dt, $
$C^{n+1} = C^* + \int_{t_n +\Delta t / 2}^{t_n+\Delta t}\mathcal{B}\left(C^*\right)dt.$

This is supposedly $\mathcal{O}(\Delta t ^2)$ from the splitting, so you can take the substeps with $\mathcal{O}(\Delta t ^2)$ schemes.

To do this I had to implement quite a few changes so that you can optionally turn off bgc transitions in the normal tendency calculation, and then add some new functions to the time steppers to allow them to just compute the tendencies for, and step the bgc.

This currently does not work, but I ran out of time to debug it. If there is interest in using this I would be happy to have another look at it.

Another thought I had was that we could allow the biogeochemical sub-stepping to be performed by e.g. DifferentialEquations.jl timesteppers, but then I realised it wouldn't be that straight forward to make a wrapper for them to work in oceananigans so decided to leave it for another time.

[glwagner]

Note there is already substepping implemented for CATKEVerticalDiffusivity and TKEDissipationVerticalDiffusivity.

What are the challenges? For the closures this feature was relatively straightforward to implement. But this PR seems quite large.

For many reasons it is often better to open the PR first (better yet, an issue that defines the problem, so we can discuss designs), rather than at a late stage where feedback is difficult to manifest.

Also, rather than supporting this generally for all time-steppers, I would argue that the correct approach is to implement this for just one model and one time-stepper. Once the proof of concept is well developed and tested, it can be applied more broadly in a separate PR.

[jagoosw]

Note there is already substepping implemented for CATKEVerticalDiffusivity and TKEDissipationVerticalDiffusivity.

I hadn't realised that, I'll look into it before I work on this again

What are the challenges? For the closures this feature was relatively straightforward to implement. But this PR seems quite large.

Most of the changes are separating the bgc transitions from the rest of the tendencies which isn't that hard but is just quite a lot of lines.

For many reasons it is often better to open the PR first (better yet, an issue that defines the problem, so we can discuss designs), rather than at a late stage where feedback is difficult to manifest.

Also, rather than supporting this generally for all time-steppers, I would argue that the correct approach is to implement this for just one model and one time-stepper. Once the proof of concept is well developed and tested, it can be applied more broadly in a separate PR.

Yeah, this makes sense, when I started doing this I thought it would be relatively simple, but then realised it's not. I would probably advocate for us to take this as a first draft and start again after discussion if this is something we go forward with.

[glwagner]

Most of the changes are separating the bgc transitions from the rest of the tendencies which isn't that hard but is just quite a lot of lines.

Can you elaborate for the benefit of future generations?

[jagoosw]

I didn't mean to commit any of this file!

[jagoosw]

Most of the changes are separating the bgc transitions from the rest of the tendencies which isn't that hard but is just quite a lot of lines.

Can you elaborate for the benefit of future generations?

For sure, the problems were:

• to split the bgc tendencies from the rest I needed a way to turn them off in the normal tendency computation, I did this by passing a keyword all the way through to biogeochemical_transitions - when I did it this way the intention was to touch as little of the existing code as possible, but it might not be the cleanest/most maintainable
• I then needed to be able to compute the bgc tendencies only so I added a function to the bgc file to launch a kernel that just computes biogeochemical_transitions
• Then I made the new time stepper which is just a container for two different timesteppers called physics and biogeochemistry (which are normal timesteppers so each have a set of tendencies etc.)
• this is where I had to make some more changes to the existing code (and probably the messiest part of what I did), because the code often expects the tendencies to live at timestepper.G / $G^-$, so I made functions that dispatched on the timestepper and usually returned that, but for the new timestepepr returned timestepper.physics.G etc.
• Then I had to add time_step_biogeochemistry! functions which do the same as time_step! but only comput the bgc tendencies

So to summarise: 1) separating the bgc tendency calculations from the rest, 2) make somewhere for the bgc tendencies to live and return the physics tendencies when expected, 3) step the bgc on its own

[glwagner]

Thank you for the written explanation! Here are my comments:

to split the bgc tendencies from the rest I needed a way to turn them off in the normal tendency computation, I did this by passing a keyword all the way through to biogeochemical_transitions - when I did it this way the intention was to touch as little of the existing code as possible, but it might not be the cleanest/most maintainable

I don't think we need to change anything in the existing kernels code. Instead, we can design an interface that allows biogeochemical models (like the models implemented in OceanBioME) to "opt-in" to a substepping scheme. In this design, there's no need to change the existing "slow" kernels I don't think. However, we could consider changing the name of them. For example, all that's needed is for the implementation to know when to return a slow source term vs fast source term.

this is where I had to make some more changes to the existing code (and probably the messiest part of what I did), because the code often expects the tendencies to live at timestepper.G / $G^-$, so I made functions that dispatched on the timestepper and usually returned that, but for the new timestepepr returned timestepper.physics.G etc.

I'm not sure that new tendencies are needed for substepping. The substepping scheme for CATKE manages to avoid allocating any additional tendencies by preserving the "slow" source term:

Oceananigans.jl/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl

Lines 167 to 171 in 6c40d7e

	@inbounds begin
	total_Gⁿe = slow_Gⁿe[i, j, k] + fast_Gⁿe
	e[i, j, k] += Δτ * (α * total_Gⁿe - β * G⁻e[i, j, k])
	G⁻e[i, j, k] = total_Gⁿe
	end

The only change that is needed within Oceananigans (in principle) --- as far as I can tell --- is to skip the tracer update for certain tracers (like we do for CATKE and TKEDissipation):

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

Lines 78 to 85 in 6c40d7e

	# TODO: do better than this silly criteria, also need to check closure tuples
	if closure isa FlavorOfCATKE && tracer_name == :e
	@debug "Skipping AB2 step for e"
	elseif closure isa FlavorOfTD && tracer_name == :ϵ
	@debug "Skipping AB2 step for ϵ"
	elseif closure isa FlavorOfTD && tracer_name == :e
	@debug "Skipping AB2 step for e"
	else

Then the implementer of the BGC model has to perform the substepping inside update_biogeochemical_state!:

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

Line 49 in 6c40d7e

update_biogeochemical_state!(model.biogeochemistry, model)

Possibly we can go further and define an interface that does the substepping automatically though. I think that an effort like this would be good not to go so far, and first test ideas in a "minimal" implementation that simply uses update_biogeochemical_state!.

[glwagner]

So in summary, a future effort should take these steps:

1. Write a function like substepped_tracers(bgc::AbstractBiogeochemistry) that can be extended by a hypothetical substepping BGC implementation.
2. Prototype this concept by implementing a new (hopefully simple, one tracer) BGC model, and prototype the time-stepping algorithm that substeps the tracer forward during update_biogeocemical_state!.

This should be easy to merge since it only requires defining the one new function substepped_tracers for skipping the tracer update. Furthermore, the prototype will demonstrate that the substepping is useful for a stiff problem, and will also develop a substepping scheme.

Next, we can consider building an interface for doing the substepping itself, much like we have an interface for doing ordinary time-stepping. That will require a bit more design, but I think the initial prototype will give us a lot of information about the best way to go about it.