Abstraction for "equation" for performance, code clarity, memory footprint reduction, and powerful user interface

[glwagner]

I'm thinking about the future and hoping to start a discussion about the future of our equation abstraction system.

The problem

We need an abstraction for the concept of an equation, so that we can make the model as performant and lightweight as possible for a given use case. For example, I think we should require that

1. The memory footprint of our model is no larger than it needs to be for a given problem (no 'extra' allocation of memory for unused tracers, for hydrostatic pressure when running in non hydrostatic mode, etc).

2. We do not perform unnecessary floating-point computations or indexing into arrays (the latter is especially important in GPU code) for unused tracers or hydrostatic pressure fields.

3. Equations are constructed / specified clearly and concisely (both in source code and user scripts).

4. Users can specify arbitrary types of forcing, including numbers, arrays, or functions (solving Allow forcing to be an array. #110).

5. We can support arbitrary tracers with various features, such as sinking/rising velocities, or reaction systems for biological/chemical tracer systems.

When I have talked to various people about this, there was a concern that this system would be 'inelegant' or 'complex'. However I believe an equation abstraction system provides the opposite: with an abstraction system, equations are 'written down' in some logical place (like a file equations.jl in the src directory where they can be easily read and modified, rather than buried inside a time-stepping loop. Correspondly, our time-stepping code becomes shorter and more concise. Using multiple dispatch correctly, we avoid the infinite if-statement problem. This abstraction may also make the code more modular such that we move closer to supporting multiple time-steppers.

Below I provide one example of an implementation that would solve some of the problems I listed. However, this is not the only solution, and I think we should expend some intellectual effort and have a discussion about what the best solution might be, so that we design something that is nice, easy to extend, performant, and powerful.

A list of kernel equations in a named tuple

The simplest solution for this abstraction is probably just to add new fields to Model (model.equations.velocities and model.equations.tracers) that are named tuples of kernel equations. An example of how this might work while demonstrating hierarchical multiple dispatch is:

forcing(i, j, k, grid, F::Function, u, v, w, T, S) = F(grid, u, v, w, T, S, i, j, k)
forcing(i, j, k, grid, F::AbstractArray, u, v, w, T, S) = F[i, j, k]

u_eqn(i, j, k, grid, etc...) = (-u∇u(grid, u, v, w, i, j, k)
    + fv(grid, v, fCor, i, j, k)
    - δx_c2f(grid, pHY′, i, j, k) / (Δx * ρ₀)
    + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, eos, grav, u, v, w, T, S)
    + forcing(i, j, k, grid, F, u, v, w, T, S)
)

# Note omission of pressure term here
u_eqn(i, j, k, grid, pHY::Nothing, etc...) = (-u∇u(grid, u, v, w, i, j, k)
    + fv(grid, v, fCor, i, j, k)
    + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, eos, grav, u, v, w, T, S)
    + forcing(i, j, k, grid, F, u, v, w T, S)
)

linear_u_eqn(i, j, k, grid, etc....) =  (fv(grid, v, fCor, i, j, k)
    - δx_c2f(grid, pHY′, i, j, k) / (Δx * ρ₀)
    + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, eos, grav, u, v, w, T, S)
    + forcing(i, j, k, grid, F, u, v, w T, S)
)

v_eqn(...) = ...

# in the model constructor:

if eqntype = :linear
    velocities_equation = (u=linear_u_eqn, v=linear_v_eqn, w=linear_w_eqn)
else
    velocities_equation = (u=u_eqn, v=v_eqn, w=w_eqn)
end

model.equation = (velocities=velocities_equation, tracers=...)

# Inside calculate_interior_source_terms!

G.u[i, j, k] = eqn.velocities.u(i, j, k, grid, etc...)

ntuple(Val(length(tracers))) do tr
    Gtracer = getfield(G, keys(tracers)[tr])
    tracer_equation = eqn.tracers[tr]
    Gtracer[i, j, k] = tracer_equation(i, j, k, grid, etc...)
end

Equations as new types

Even more flexibility / clarity might be provided by a callable Equation type implementation with appropriate type parameters (Linear, Hydrostatic, etc) which would avoid the need for if statements as on the last line of the code above. I'm not sure exactly what that would look like. The design problem for a new Equation type does seem a bit more difficult and would require care --- thinking about how we would implement a nonlinear biological model might provide a nice use case to help use make our ideas concrete.

Hopefully this provides food for thought and also motivates us to think carefully about our design decisions.

[glwagner]

A simpler (but less general) solution than what is proposed above is simply to introduce a new struct:

struct BoussinesqEquations{A, B, C, D, F, P} <: AbstractEquation
    advection :: A
    buoyancy :: B
    coriolis :: C
    diffusivity :: D
    forcing :: F
    pressure :: P
end

This will reduce the number of fields in Model. We can also implement new equations 'types' that dispatch to different solvers if need be. This may also dramatically clean up the time-stepping functionality since we can pass eq::Equation in bulk rather than each parameter struct individually.

Less general than a tuple of tendency kernels, but probably sufficient for anything we want to do in the forseeable future.

We can then have a function eg required_halo_size(eq::AbstractEquation) that calculates the halo size according to the maximum halo size required by each of its 'terms'.

Could also have an independent struct for tracer equations (or tuple of structs for each tracer).

By the way, the more I think about it, the more I think that 'closure' should actually be called 'diffusivity', for clarity...

cc @ali-ramadhan

[ali-ramadhan]

I like this idea, it would definitely work towards having an equations abstraction.

Although I might be wary of doing this while we only support one equation set. Might make sense to consider an AbstractEquation type once we have e.g. BoussinesqEquations and CompressibleNavierStokesEquations to make sure the abstraction generalizes well beyond just the Boussinesq equations?

[glwagner]

Yesterday, @ali-ramadhan and I discussed a new idea for improving the abstraction of equations.

The idea is to abstract a RightHandSide consisting of a tuple of Fluxes and VolumeTerms or SourceTerms.

Each Flux or VolumeTerm type would define a getindex method, and carry around the references needed to execute that getindex method on CPU or GPU.

This would greatly simplify the time-stepping routines, which currently involve long function signatures. There'd be no need for 'unpacking', because each term in the equation would perform unpacking upon instantiation. It would also probably be easier for users to extend / add terms to an equation.

As an example, we can consider a simple implementation

struct AdvectiveFlux{S, D, C, G}
    u :: D
    v :: D
    w :: D
    c :: C
    grid :: G
    scheme :: S
end

getindex(adv::AdvectiveFlux{Centered}, i, j, k) = # centered advective flux calculation

struct IsotropicDiffusiveFlux{N, D, G}
    ν :: N
    ψ :: D
    grid :: G
end

getindex(diff::IsotropicDiffusiveFlux, i, j, k) = # calculates flux due to isotropic diffusion by ν

struct RightHandSide{F, V}
    fluxes :: F
    volume_terms :: V
end

advection = AdvectiveFlux(velocities..., tracers.c)
diffusion = IsotropicDiffusiveFlux(ν, tracers.c)

tracer_rhs = RightHandSide((advection, diffusion), nothing)

We'd have functions that look something like

function x_flux_divergence(i, j, k, grid, fluxes...)
    incoming_flux = add_fluxes(i, j, k, grid, fluxes...)
    outgoing_flux = add_fluxes(i+1, j, k, grid, fluxes...)
    return (incoming_flux - outgoing_flux) * grid.Ax / grid.V
end

... for example. Obviously questions of performance are paramount, though in the case that everything is inlined I think there is hope.

A downside of this approach is that we can't use shared memory stencils on the GPU. Shared memory stencils on the GPU require functions for all terms that avoid carrying around internal references to data (since we need to be able to pass them references to shared memory blocks, rather than references to global memory.

This is a fair rewrite of the code internals. For example, each term in the u_velocity_tendency:

https://github.com/climate-machine/Oceananigans.jl/blob/8e3c27504be68ca06bacc7502cd6095ae390f8c6/src/TimeSteppers/velocity_and_tracer_tendencies.jl#L24

would get it's own type.

It's worth brainstorming ways to implement such abstraction incrementally so we might avoid possibly time-consuming total-code-demolishment. This is mostly food for thought at this point. I don't think we should take action without substantial consideration.

[glwagner]

I don't think we want this anymore: we currently use IncompressibleModel as a flat container for both equation specification and numerical methods specification. Flat is nice.