Adds an interface for inserting biogeochemistry models into Oceananigans models

[glwagner]

Plus a simple test with a plankton growth/death model similar to examples/convecting_plankton.jl.

TODO:

• Add the interface to HydrostaticFreeSurfaceModel
• Better interface for defining the biogeochemical source function
• Better interface for automagically adding biogeochemical tracers?

With @jagoosw

[jagoosw]

I've had a play implementing an NPD model in this framework now and think that the other way to define the source functions is preferential:

Oceananigans.jl/test/test_biogeochemistry.jl

Lines 33 to 37 in 45971e4

	@inline function growing_and_grazing(i, j, k, grid, ::Val{:P}, clock, bgc::SimplePlanktonGrowthDeath, fields)
	z = znode(Center(), k, grid)
	P = @inbounds fields.P[i, j, k]
	return (bgc.μ₀ * exp(z / bgc.λ) - bgc.m) * P
	end

We could even go one step further and define a model like this:

struct NPD
    Kₙ :: Float64
    m :: Float64
    nitrif :: Float64
end

validate_biogeochemistry(::NPD, tracernames) = all([T ∈ tracernames for T in [:N, :P, :D]])

@inline function (model::NPD)(i, j, k, grid, ::Val{:N}, clock, fields)
    P = @inbounds fields.P[i, j, k]
    N = @inbounds fields.N[i, j, k]
    D = @inbounds fields.D[i, j, k]

    return model.nitrif*D - P*N/(N+model.Kₙ) 
end

@inline function (model::NPD)(i, j, k, grid, ::Val{:P}, clock, fields)
    P = @inbounds fields.P[i, j, k]
    N = @inbounds fields.N[i, j, k]
    return P*N/(N+model.Kₙ) - model.m*P
end

@inline function (model::NPD)(i, j, k, grid, ::Val{:D}, clock, fields)
    P = @inbounds fields.P[i, j, k]
    D = @inbounds fields.D[i, j, k]

    return model.m*P - model.nitrif*D
end

@inline (model::NPD)(args...) = 0.0

This negates the need to define a get_biogeochemial_forcing function, but you do have to define the zero function (last line) and I'm not sure how clear and usable this API is to most users?

[navidcy]

related to #2512

[jagoosw]

@johnryantaylor do you have any preferences for this API?

[glwagner]

We can use an abstract type to obviate the need for "zero functions":

abstract type AbstractBiogeochemistry end

struct NutrientsPlanktonDetritus{FT} <: AbstractBiogeochemistry
    background_nutrients :: FT
    mortality_rate :: FT
    nitrification :: FT
end

then with

@inline (::AbstractBiogeochemistry)(i, j, k, grid, val_tracer_name, clock, fields) = zero(grid)

users don't need to define the "netural biogeochemical forcing" themselves.

Very Important: always use verbose names! I don't know how to enforce that within the API 😂

I like this interface. Let's figure out if Val{symbol} is GPU-friendly.

[glwagner]

Instead of validate_biogeochemistry, maybe

required_biogeochemical_tracers(::NutrientsPlanktonDetritus) = (:N, :P, :D)

is a better syntax. Then users don't have to write their own error messages (we'll handle that in Oceananigans). It also gives us flexibility regarding the choice between 1) automagically adding the biogeochemical tracers or 2) requiring the user to add tracers themselves.

We starting to use tracers for a few important things (TKE, biogeochemistry, buoyancy models). We may also need to come up with a system for handling "name clashes" gracefully. For example, someone might introduce a biogeochemistry model with tracer e, not realizing that this would prevent them from using the biogeochemistry model with CATKE. There's a tension between simplicity and readability (which we get with short, intuitive tracer names like e, T, S), and catastrophic "name clashes" that will inevitably occur as model complexity increases.

[jagoosw]

Ah I thought we could do something with abstract types for that but couldn't get it to work before. I'll try this on a GPU now if I can get one.

And that makes sense for the validation.

[glwagner]

A lot of tests are failing too, we need to fix those

[jagoosw]

test/test_biogeochemistry.jl as it currently is in this branch runs on GPU so it looks like Val{Symbol} does work!

[johnryantaylor]

This API looks good to me. I agree with all of @glwagner's suggestions. If we don't use explicit function names, then I think that it will be important to help the users by adding a comment at the start of the function to explain what that definition is doing and how to call the function.

[jagoosw]

I've had a go at implementing a proper NPZD model (rather than one I just made up on the fly) and have some thoughts on how we should modify the API:

• I think we need an required_biogeochemical_auxiliary_fields like required_biogeochemical_tracers because for most models we're going to want the user to at least specify a PAR field (I suppose we may want this to also check the shape of the field because some models may have a pre defined depth dependence of PAR so we might want the user to specify a 2D PAR field rather than doing it properly by integrating a 3D field)
• Given what you said the other day about callbacks only being used for features that should be built into Oceananigans we might want to have a think about how a BGC model can specify the attenuation of PAR. You've mentioned that we could define some kind of integrated field?
• It might be helpful to have a simpler interface for advection in biogeochemical models. Although a user could just add another forcing, I think the only way for a model to automatically add an advective forcing is how I've implimented it in the below example. I Think this works quite well since a lot of BGC models write the sinking terms with the other forcing terms, but its a little cumbersome to write e.g. sinking = div_Uc(i, j, k, grid, bgc.adv_scheme, bgc.u⃗ᵖ, fields.P), and model makers will need to do the setup stuff I've done to make the advective velocity fields.

You can see my implementation here and a script using it here since I thought it was probably too complicated for the test (and will change that back to a 1 variable model later). Not finished making it work but yet but will be done soon.

[glwagner]

I think we need an required_biogeochemical_auxiliary_fields like required_biogeochemical_tracers because for most models we're going to want the user to at least specify a PAR field (I suppose we may want this to also check the shape of the field because some models may have a pre defined depth dependence of PAR so we might want the user to specify a 2D PAR field rather than doing it properly by integrating a 3D field)

Should the user specify this, or should the biogeochemical model add it its struct (taking in grid for this purpose) and evaluate it during update_biogeochemical_state?

[glwagner]

It might be helpful to have a simpler interface for advection in biogeochemical models. Although a user could just add another forcing, I think the only way for a model to automatically add an advective forcing is how I've implimented it in the below example. I Think this works quite well since a lot of BGC models write the sinking terms with the other forcing terms, but its a little cumbersome to write e.g. sinking = div_Uc(i, j, k, grid, bgc.adv_scheme, bgc.u⃗ᵖ, fields.P), and model makers will need to do the setup stuff I've done to make the advective velocity fields.

Can't we add advective terms via the biogeochemical forcing term?

edit: I see you mentioned that, so I might be missing something... I'll take a look at your examples.

PS we should add those scripts to validation/biogeochemistry/

[glwagner]

I think we can impose a little more structure that eases biogeochemical model development.

One route is to build out a layer on top of AbstractBiogeochemistry for models with a common form, something like

struct TracerBasedBiogeochemistry
    biogeochemical_tracers
    drift_advection_schemes
    drift_velocities
    transitions
    auxiliary_fields
end

and perhaps more properties. This is similar to how SeawaterBuoyancy works, in that it provides a concrete structure with a slot equation_of_state whereby "external" packages like SeawaterPolynomials can insert custom behavior.

Another possibility is to build out a new abstract type below AbstractBiogeochemistry with a function-based interface (somehow).

I was also thinking it would be nice to supply a "continuous form" interface, so that model developers can implement functions that look something like

biogeochemical_transition(x, y, z, t, N, P, Z, D, parameters)

rather than having to use the "discrete form".

This sort of structure could also be provided by an external package. The advantage of including it here is that we get tighter coupling with Oceananigans development.

[jagoosw]

I think we need an required_biogeochemical_auxiliary_fields like required_biogeochemical_tracers because for most models we're going to want the user to at least specify a PAR field (I suppose we may want this to also check the shape of the field because some models may have a pre defined depth dependence of PAR so we might want the user to specify a 2D PAR field rather than doing it properly by integrating a 3D field)

Should the user specify this, or should the biogeochemical model add it its struct (taking in grid for this purpose) and evaluate it during update_biogeochemical_state?

This is a good point, I think doing it as part of the model in update state is a much better solution

[jagoosw]

Can't we add advective terms via the biogeochemical forcing term?

edit: I see you mentioned that, so I might be missing something... I'll take a look at your examples.

PS we should add those scripts to validation/biogeochemistry/

I think we can impose a little more structure that eases biogeochemical model development.

One route is to build out a layer on top of AbstractBiogeochemistry for models with a common form, something like

struct TracerBasedBiogeochemistry
    biogeochemical_tracers
    drift_advection_schemes
    drift_velocities
    transitions
    auxiliary_fields
end

and perhaps more properties. This is similar to how SeawaterBuoyancy works, in that it provides a concrete structure with a slot equation_of_state whereby "external" packages like SeawaterPolynomials can insert custom behavior.

Another possibility is to build out a new abstract type below AbstractBiogeochemistry with a function-based interface (somehow).

I was also thinking it would be nice to supply a "continuous form" interface, so that model developers can implement functions that look something like

biogeochemical_transition(x, y, z, t, N, P, Z, D, parameters)

rather than having to use the "discrete form".

This sort of structure could also be provided by an external package. The advantage of including it here is that we get tighter coupling with Oceananigans development.

This seems like a good idea, I will have a think and try and come up with something

[jagoosw]

Sorry meant to put [skip ci]

[jagoosw]

Anything else needed now?

[navidcy]

2 spaces for markdown

[navidcy]

This looks good to me! The docstring for AbstractBiochemistry should evolve into a documentation page in the Docs soon! That can be a new PR. :)

[glwagner]

@seamanticscience docs for Biogeochemistry could be something to work on once this is merged

[jagoosw]

Is this ready to merge now?

[navidcy]

I think so. @glwagner?

[glwagner]

yay!