A functional tendency specification layer

In the balance law (mutating) functions, where we specify fluxes and sources,

• flux_first_order!
• flux_second_order!
• and source!,

an additional (functional) tendency specification layer can be placed on-top that has several nice properties. The functional layer:

• Separates tendency definitions from which tendencies are included in a particular model.
• Reduces duplicate implementations of tendency definitions (e.g., in optional submodel variants)
• Allows a more flexible combination of tendencies
• Allows a simple way to loop over all tendencies for all prognostic variables and recover each flux / source term. This will allow us a simple way to evaluate, for example, the energy budget.

Used modules / imports

Make running locally easier from ClimateMachine.jl/:

if !("." in LOAD_PATH)
    push!(LOAD_PATH, ".")
    nothing
end

First, using necessary modules:

using ClimateMachine.BalanceLaws
using ClimateMachine.VariableTemplates
using StaticArrays, Test

Import methods to overload

import ClimateMachine.BalanceLaws: prognostic_vars, eq_tends, flux

Define a balance law

Here, we define a simple balance law:

struct MyBalanceLaw <: BalanceLaw end

Define prognostic variable types

Here, we'll define some prognostic variable types, by sub-typing AbstractPrognosticVariable, for mass and energy:

struct Mass <: AbstractPrognosticVariable end
struct Energy <: AbstractPrognosticVariable end

Define prognostic_vars, which returns all prognostic variables

prognostic_vars(::MyBalanceLaw) = (Mass(), Energy());

Define tendency definition types

Tendency definitions types are made by subtyping TendencyDef. TendencyDef has one type parameters: the AbstractTendencyType, which can be either Flux{FirstOrder}, Flux{SecondOrder}, or Source.

struct Advection <: TendencyDef{Flux{FirstOrder}} end
struct Source1 <: TendencyDef{Source} end
struct Source2 <: TendencyDef{Source} end
struct Diffusion <: TendencyDef{Flux{SecondOrder}} end

Define eq_tends, which returns a tuple of tendency definitions (those sub-typed by TendencyDef), given

• the prognostic variable
• the model (balance law)
• the tendency type (Flux or Source)

eq_tends(::Mass, ::MyBalanceLaw, ::Flux{FirstOrder}) = (Advection(),);
eq_tends(::Energy, ::MyBalanceLaw, ::Flux{FirstOrder}) = (Advection(),);
eq_tends(::Mass, ::MyBalanceLaw, ::Flux{SecondOrder}) = ();
eq_tends(::Energy, ::MyBalanceLaw, ::Flux{SecondOrder}) = (Diffusion(),);
eq_tends(::Mass, ::MyBalanceLaw, ::Source) = (Source1(), Source2());
eq_tends(::Energy, ::MyBalanceLaw, ::Source) = (Source1(), Source2());

Testing prognostic_vars eq_tends

To test that prognostic_vars and eq_tends were implemented correctly, we'll create a balance law instance and call show_tendencies, to make sure that the tendency table is accurate.

bl = MyBalanceLaw()
show_tendencies(bl; table_complete = true)

PDE: ∂_t Y_i + (∇•F_1(Y))_i + (∇•F_2(Y,G)))_i = (S(Y,G))_i
┌──────────┬──────────────────┬───────────────────┬────────────────────┐
│ Equation │ Flux{FirstOrder} │ Flux{SecondOrder} │             Source │
│    (Y_i) │            (F_1) │             (F_2) │                (S) │
├──────────┼──────────────────┼───────────────────┼────────────────────┤
│     Mass │      (Advection) │                () │ (Source1, Source2) │
│   Energy │      (Advection) │       (Diffusion) │ (Source1, Source2) │
└──────────┴──────────────────┴───────────────────┴────────────────────┘

The table looks correct. Now we're ready to add the specification layer.

Adding the tendency specification layer

For the purpose of this tutorial, we'll only focus on adding the layer to the first order flux, since doing so for the second order flux and source functions follow the same exact pattern. In other words, we'll add a layer that tests the Flux{FirstOrder} column in the table above. First, we'll define individual flux kernels:

flux(::Mass, ::Advection, bl::MyBalanceLaw, args) =
    args.state.ρ * SVector(1, 1, 1);
flux(::Energy, ::Advection, bl::MyBalanceLaw, args) =
    args.state.ρe * SVector(1, 1, 1);

Define flux_first_order! and utilize eq_tends

function flux_first_order!(
    bl::MyBalanceLaw,
    flx::Grad,
    state::Vars,
    aux,
    t,
    direction,
)

    tend_type = Flux{FirstOrder}()
    args = (; state, aux, t, direction)

    # `Σfluxes(Mass(), eq_tends(Mass(), bl, tend_type), bl, args)` calls
    # `flux(::Mass, ::Advection, ...)` defined above:
    eqt_ρ = eq_tends(Mass(), bl, tend_type)
    flx.ρ = Σfluxes(Mass(), eqt_ρ, bl, args)

    # `Σfluxes(Energy(), eq_tends(Energy(), bl, tend_type), bl, args)` calls
    # `flux(::Energy, ::Advection, ...)` defined above:
    eqt_ρe = eq_tends(Energy(), bl, tend_type)
    flx.ρe = Σfluxes(Energy(), eqt_ρe, bl, args)
    return nothing
end;

Testing the tendency specification layer

Now, let's test flux_first_order! we need to initialize some dummy data to call it first:

FT = Float64; # float type
aux = (); # auxiliary fields
t = 0.0; # time
direction = nothing; # Direction

state = Vars{@vars(ρ::FT, ρe::FT)}([1, 2]);
flx = Grad{@vars(ρ::FT, ρe::FT)}(zeros(MArray{Tuple{3, 2}, FT}));

call flux_first_order!

flux_first_order!(bl, flx, state, aux, t, direction);

Test that flx has been properly mutated:

@testset "Test results" begin
    @test flx.ρ == [1, 1, 1]
    @test flx.ρe == [2, 2, 2]
end

nothing

Test Summary: | Pass  Total
Test results  |    2      2

This page was generated using Literate.jl.