TendencyTermField (or something like it) for diagnosing exact tendency terms and fluxes

[glwagner]

As being discussed over on CliMA/LESbrary.jl#25, and also as discussed with @christophernhill and @ali-ramadhan separately, we need a way to diagnose the terms that arise in the velocity and tracer tendency equations, as well as the fluxes that give rise to the flux divergences that appear there.

One way to achieve this purpose is to define new types of operands for ComputedFields. Currently, operand is always an AbstractOperation. However, we can come up with a new type that looks something like

struct FunctionOperand{F, G, A}
    func :: F
    grid :: G
    args :: A
end

@inline Base.getindex(o::FunctionOperand, i, j, k) = o.func(i, j, k, o.grid, o.args...)

This works because the kernel that computes a ComputedField's data is

Oceananigans.jl/src/Fields/computed_field.jl

Lines 85 to 88 in c5f47e0

	@kernel function _compute!(data, operand)
	i, j, k = @index(Global, NTuple)
	@inbounds data[i, j, k] = operand[i, j, k]
	end

.

Then with a bit of boilerplate we can define constructors for all the terms we might need, eg

TendencyTermField(X, Y, Z, term_func, grid; args, data=nothing) =
    ComputedField{X, Y, Z}(FunctionOperand(term_func, grid, args), data=data)

using Oceananigans.Advection: momentum_flux_uu, momentum_flux_uw

uu = TendencyTermField(Cell, Cell, Cell, momentum_flux_uu, grid, args=(model.advection, model.velocities.u, model.velocities.u))
wu = TendencyTermField(Cell, Cell, Face, momentum_flux_uw, grid, args=(model.advection, model.velocities.w, model.velocities.u))

We probably want to define aliases for all the terms that appear in our tendency equations, as well as the advective and diffusive fluxes, so that we can ensure they are correct and correctly located on the staggered grid. Unfortunately this does involve a lot of boiler plate and introduces a maintenance and testing burden. If we can push responsibility more to users I am open to that, but I'm not 100% how to make this process more programmatic. Ideas very welcome.

If functions like momentum_flux_uu are going to emerge from darkness into users' scripts, we may want to have a discussion about whether our names / naming convention is sensible.

By the way, ComputedField seems to assume that operand has a property called grid:

Oceananigans.jl/src/Fields/computed_field.jl

Lines 43 to 47 in c5f47e0

	function ComputedField(operand; data=nothing, recompute_safely=true)
	loc = location(operand)
	arch = architecture(operand)
	grid = operand.grid

We could create a type called AbstractOperand, of which AbstractOperation (and other concrete operands) are subtypes.

cc @qingli411, @BrodiePearson

[glwagner]

Resolving this issue requires first resolving #1063 and #1069 .

[ali-ramadhan]

We probably want to define aliases for all the terms that appear in our tendency equations

If diagnosing all the different terms for various budgets (momentum, tracer, TKE, etc.) is an important use case then I don't think we can escape the boilerplate associated with that: https://mitgcm.readthedocs.io/en/latest/outp_pkgs/outp_pkgs.html#mitgcm-kernel-available-diagnostics-list

If we can push responsibility more to users I am open to that, but I'm not 100% how to make this process more programmatic. Ideas very welcome.

We might be able to get away with clever @eval loops but that could significantly degrade the readability of the code. We might just have to embrace some boilerplate...

Maybe we can start with a barebones list of TendencyTermFields. Then we can decide whether we keep growing the list inside Oceananigans.jl as users need to diagnose new terms or just teach them to define their own TendencyTermFields.

I'd probably advocate for growing the list in Oceananigans.jl if we think the list contains useful terms that other users would be interested in using. But since it is always possible for power users (or any user with some help) to define their own TendencyTermFields we should be selective in which terms we include in Oceananigans as it does unfortunately come with a maintenance burden.

Growing the list inside Oceananigans.jl (not relying on users to define all terms) has the added benefit that we can keep functions like momentum_flux_uv unexported and internal (so we can change them easily).

[glwagner]

Perhaps to start we can implement

1. Momentum fluxes (9 of them)
2. Tracer fluxes (3 of them)
3. Every term in the tracer budget (since it seems tracer budgets are more important than momentum budgets.) There are 5 terms there:

Oceananigans.jl/src/TimeSteppers/velocity_and_tracer_tendencies.jl

Lines 233 to 238 in c5f47e0

	return ( - div_Uc(i, j, k, grid, advection, velocities, c)
	- div_Uc(i, j, k, grid, advection, background_fields.velocities, c)
	- div_Uc(i, j, k, grid, advection, velocities, background_fields_c)
	+ ∇_κ_∇c(i, j, k, grid, clock, closure, c, val_tracer_index, diffusivities, tracers, buoyancy)
	+ forcing(i, j, k, grid, clock,
	merge(velocities, tracers, regularize_diffusivities_tuple(diffusivities))))

That's a total of 17 TendencyTermFields.

This still does omit a large number of terms that are part of the momentum budgets, so we do save something by not implementing every term.

[glwagner]

A related feature would be fields like KineticEnergyDissipation and TracerVarianceDissipation that dispatch on the closure type.

Ideally these terms could be written out using AbstractOperations, but it appears that complex operations do not compile on the GPU.

[ali-ramadhan]

I wonder if it makes sense to spend some effort to try and get more complex operations to compile on the GPU (e.g. so we can see if it's even performant: #870) so if it's performant we can start using them to work on issues like this.

[glwagner]

We would use KernelFunctionOperation for this now.