State variables

Overview

Models and processes must define dispatches of the variables method that return a Tuple of variables subtyping AbstractVariable:

struct PrognosticVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units, Var<:Terrarium.Variable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}, CL<:Union{Nothing, Terrarium.AbstractClosureRelation}, TV<:Union{Nothing, AuxiliaryVariable}, DT<:IntervalSets.AbstractInterval} <: Terrarium.AbstractProcessVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}

struct AuxiliaryVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units, Var<:Terrarium.Variable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}, DT<:IntervalSets.AbstractInterval, FC<:Union{Nothing, Function}} <: Terrarium.AbstractProcessVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}

struct InputVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units, Var<:Terrarium.Variable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}, DT<:IntervalSets.AbstractInterval, Def<:Union{Nothing, Function, Number}} <: Terrarium.AbstractProcessVariable{name, VD<:Terrarium.VarDims, UT<:Unitful.Units}

Variables are typically constructed using one of the convenience functions prognostic, auxiliary, or input.

These variable definitions are purely symbolic; they do not hold any data and cannot be used for computation. Calling initialize on a model, process, or Variables container (see following sections) results in corresponding Fields being allocated for each variable.

A default implementation of variables is provided for all AbstractModel and AbstractCoupledProcesses types that automatically collects variables from all AbstractProcess types defined therein:

variables(obj::Union{AbstractCoupledProcesses, AbstractModel})

Most state variables will thus be defined by implementation of AbstractProcess. As an example, suppose we are implementing a new process MyProcess and we want to define the necessary state variables. We do this by defining a new dispatch of the variables method:

variables(::MyProcess) = (
    prognostic(:progvar, XYZ()),
    auxiliary(:auxvar, XYZ()),
    auxiliary(:bc, XY()),
    input(:input, XY())
)

This will result in a total of five state variables being allocated upon initialization: two auxiliary variables named auxvar and bc as well as one prognostic variable named progvar along with its corresponding tendency variable which is created automatically. The second argument to the variable metadata constructors prognostic and auxiliary is a subtype of VarDims which specifies on which spatial dimensions the state variable should be defined. XYZ() corresponds to a 3D Field which varies both laterally and with depth. XY() corresponds to a 2D field which is discretized along the lateral X and Y dimensions only.

Merging and promotion rules

Models rarely consist of only a single process, and many processes inevitably need access to the same physical variables. Terrarium handles this by automatically merging variables that are duplicated between model components. Duplicate variables must have the same dimensions and physical units in order to be merged, otherwise an error will be raised during initialization. Variables with identical names, dimensions, and units but different types (i.e. prognostic vs. auxiliary vs. input) are merged according to two simple promotion rules:

• input variables matching a corresponding prognostic or auxiliary variable are "promoted" to the corresponding prognostic/auxiliary type
• Clashing definitions of prognostic and auxiliary definitions are disallowed and result in an error

The reason for the latter rule is that variables which are treated as prognostic (integrated by the time stepper) cannot simultaneously be treated as auxiliary (i.e. derived) from the prognostic state.

These merging and promotion rules are defined and applied by the Variables container:

struct Variables{ProgVars, TendVars, AuxVars, InputVars, Namespaces}

Variables also provides helpful dispatches of show which concisely summarize the variables in the container after merging/promotion rules have been applied.

Auxiliary variables with derived Fields

As briefly discussed in the doc section on Fields, Oceananigans features a powerful system of defining Fields that are lazily computed as operators defined over the model grid. Terrarium takes advantage of this by permitting custom Field constructors in the definition of auxiliary variables:

auxiliary(name, dims; ...)
auxiliary(name, dims, ctor; ...)
auxiliary(name, dims, ctor, params; units, domain, desc)

Here the ctor argument should be a function with signature (process::ProcessType, grid, clock, fields) that returns an Oceananigans operator-derived Field.

As a simple example, suppose we want to define an auxiliary variable C for the hypothetical process type Pythagoras which is derived from two other state variables (Fields) length and width via the relation $C = \sqrt{A^2 + B^2}$. This could be accomplished as follows:

struct Pythagoras{NF} <: Terrarium.AbstractProcess{NF} end

Terrarium.variables(pythag::Pythagoras) = (
    auxiliary(:hypotenuse, XY(), hypotenuse, pythag),
    input(:length, XY()),
    input(:width, XY())
)

function hypotenuse(::Pythagoras, grid, clock, fields)
    hypotenuse = sqrt(fields.length^2 + fields.width^2)
    return hypotenuse
end

grid = ColumnGrid(CPU(), Float64, UniformSpacing(N = 1))
state = initialize(Pythagoras{Float64}(), grid)
state.hypotenuse

UnaryOperation at (Center, Center, ⋅)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 1×0×1 halo
└── tree: 
    sqrt at (Center, Center, ⋅) via identity
    └── + at (Center, Center, ⋅)
        ├── ^ at (Center, Center, ⋅)
        │   ├── 1×1×1 Field{Center, Center, Nothing} reduced over dims = (3,) on Oceananigans.Grids.RectilinearGrid on CPU
        │   └── 2
        └── ^ at (Center, Center, ⋅)
            ├── 1×1×1 Field{Center, Center, Nothing} reduced over dims = (3,) on Oceananigans.Grids.RectilinearGrid on CPU
            └── 2

The benefit of using this pattern is that we save memory overhead by avoiding the allocation of new Field memory for hypotenuse. Instead, the expression in the above function is evaluated each time hypotenuse is indexed:

set!(state.length, 2)
set!(state.width, 3)
state.hypotenuse[1,1,1]

3.605551275463989

As another more concrete example, consider the soil_moisture_limiting_factor Field constructor for FieldCapacityLimitedPAW:

function soil_moisture_limiting_factor(::FieldCapacityLimitedPAW, grid, clock, fields)
    Δz = zspacings(get_field_grid(grid), Center(), Center(), Center())
    β = Integral(fields.plant_available_water * fields.root_fraction / Δz, dims = 3)
    return Field(β)
end

The resulting factor β is an integral over the enclosed function of plant_available_water and root_fraction. Since Integral is a so-called "reduction" operator, it cannot be lazily computed and instead compute! must be called on the Field in each time step, e.g:

compute!(state.soil_moisture_limiting_factor)