GitHub

Implementing processes

This page is a step-by-step guide for implementing a new concrete process type in Terrarium. We assume familiarity with the Basic concepts and Core interfaces pages, both of which cover the AbstractProcess interface and variable system. For a concrete, self-contained example that follows the workflow below step by step, see the Implementing a process: 1D linear heat diffusion example.

Control flow

Every process in Terrarium is implemented across three levels of abstraction:

  • I. Top-level interface methods for AbstractProcess, i.e. variables, initialize!, compute_auxiliary!, and compute_tendencies! (see doc section on the AbstractProcess interface)
  • II. Kernel functions; this includes both the @kernel entry point and subsequently inlined functions with signatures of the form compute_*(i,j[,k], grid, fields, ::Process, args...)
  • III. Primitive methods defined directly by the subtype of AbstractProcess. These are scalar-valued, grid-independent functions corresponding to individual terms or expressions in the mathematical formulation of the process physics.

The flow of execution is I → II → III: top-level interface methods are invoked by the enclosing AbstractModel, these methods then launch! their corresponding @kernels which in turn call the inner kernel functions.

Development workflow

We recommend to start implementing new processes using the following workflow:

  1. Start by drafting an initial version of the process struct subtyping AbstractProcess; it's good to start simple by directly putting all relevant parameters into this struct. You can later consider grouping these parameters into logical parameterization types.
  2. Define variables(::Process) for the new process type. Identify which variables are prognostic, which are auxiliary, and which might be assumed to be inputs, provided either by external forcing data or by another process.
  3. Write down unit tests and initial implementations of the scalar-valued process methods (level III). These should correspond to the key equations or terms in the mathematical formulation of the physics. Note that these methods are invoked on each elementary volume and thus cannot be used to define equations involving spatial gradients or fluxes.
  4. Implement the compute_*(i, j[, k], grid, fields, ...) kernel functions with the corresponding @kernel entry points compute_auxiliary_kernel! and compute_tendencies_kernel! (level II). Functions which might be independently useful outside of the process interface (e.g. adjust_saturation_profile! in SoilHydrology) can also define their own @kernel entry points.
  5. Finally, implement the top-level process methods for initialize!, compute_auxiliary! and compute_tendencies! that launch! their respective @kernels (level I).

As mentioned above, see the example Implementing a process: 1D linear heat diffusion for a detailed walkthrough of this workflow.

Enabling implicit timestepping

By default, all prognostic variables are integrated explicitly. A diffusion-like operator can instead be integrated implicitly under an IMEX timestepper (see Implicit-explicit time stepping). To support this in a process:

  1. Accept a Timestepping argument. Give the process's compute_tendencies! a trailing positional argument timestepping::Timestepping = Explicit() and pass it through the kernel to the kernel functions, which should then dispatch on Explicit or Implicit. The model obtains the resolved class of a variable with timestepping(state, Val(:name)). See compute_tendencies! of SoilThermodynamics for an example.
  2. Omit the implicitly treated terms under Implicit(). The tendency under Implicit() must contain only the explicitly treated terms (e.g. sources and boundary fluxes). The omitted terms are instead applied via the corresponding implicit operator. The Explicit() implementation of compute_tendencies! should typically remain unchanged so that users also have the option to use explicit timesteppers.
  3. Declare prognostic variables that should be solved implicitly. Define timestepping(::AbstractVariable{:name}, ::YourModel, ::AbstractIMEX) = Implicit() for the model that owns the variable.
  4. Define the implicit operator. The implicit timestepper calls implicit_operator(model, state, names) once on the host with the names of all routed variables when its cache is allocated at initialization, and the model returns one AbstractLinearOperator for all of them. Every model with implicitly routed variables must define this method; the fallback throws an ArgumentError at initialization. To solve the variables one after the other, return a SplitOperator of newly allocated per-variable operators (e.g. a TridiagonalOperator for a vertical diffusion operator) in the order in which they should be solved; a model that couples its variables returns its own operator type instead.
  5. Implement the process hooks. In each iteration, the implicit solver calls the hooks linearize!, implicit_update!, and implicit_closure! for the variable, which the model forwards to the process that owns it (the linear solve, solve!, is generic over the operator). Implement them for the process; implicit_closure! can usually forward to the regular closure! of the process. Their signatures and what each must do are described in Process hooks. The linear system is solved for the variable of the prognostic variable's closure relation, so a closure implements varname if its variable can be integrated implicitly. See SoilThermodynamics and SoilHydrology with RichardsEq for examples.
Terrarium.implicit_operator — Function
implicit_operator(model::AbstractModel, state, names::Tuple)

Return a newly allocated AbstractLinearOperator for the implicitly routed prognostic variables names of model (e.g. a SplitOperator of per-variable operators in solve order). Called once on the host at initialize; methods throw an ArgumentError for unsupported configurations, as does the fallback. SoilModel and LandModel forward to implicit_operator(soil::AbstractSoil, state, grid, names).

source
implicit_operator(
    hydrology::Terrarium.AbstractSoilHydrology,
    state,
    grid
)

Fallback for hydrology processes without an implicit operator; throws an ArgumentError.

source
implicit_operator(
    hydrology::SoilHydrology{NF, RichardsEq, SaturationClosure, SoilHydraulics} where {SaturationClosure<:Terrarium.AbstractSoilWaterClosure, SoilHydraulics<:(Terrarium.AbstractSoilHydraulics{NF})},
    state,
    grid
) where NF

Allocate the TridiagonalOperator of saturation_water_ice on ground_domain(grid). Throws an ArgumentError unless the closure is the SoilSaturationPressureClosure and pressure_head has no value boundary conditions.

The operator is the backward Euler step of the Richards equation in the total pressure head Ψ, after the mass-conservative modified Picard iteration of [3]: with porosity φ, saturation ξ, and ξ⁺ = ξⁿ + Δt G_exp from the explicit stage, the water content is linearized about the iterate m, φ ξᵐ⁺¹ ≈ φ ξᵐ + Cᵐ (Ψᵐ⁺¹ - Ψᵐ) with the specific moisture capacity Cᵐ and the hydraulic conductivity Kᵐ lagged, giving the tridiagonal system Cᵐ Ψᵐ⁺¹ - Δt ∂z(Kᵐ ∂z Ψᵐ⁺¹) = Cᵐ Ψᵐ - φ (ξᵐ - ξ⁺) (see the soil hydrology docs for the derivation). Each Picard iteration m = 0, …, M-1, starting from ξ⁰ = ξⁿ, calls the hooks:

  1. linearize!: hydraulic conductivity, lagged Kᵐ (faces) and Cᵐ (centers, bounded below by minimum_moisture_capacity so that frozen saturated cells stay regular), and the right-hand side;
  2. solve!: the system above for Ψᵐ⁺¹;
  3. implicit_update!: ξᵐ⁺¹ = ξ⁺ + Δt ∂z(Kᵐ ∂z Ψᵐ⁺¹) / φ through the interior faces, which conserves water exactly, and water_flux ← -Kᵐ ∂z Ψᵐ⁺¹, the flux with which the heat step advects heat;
  4. implicit_closure!: redistribution of over- and undersaturation, water table, Ψᵐ⁺¹, and hydraulic conductivity. Since every iteration restarts from ξ⁺, only the final one routes surface excess water to the runoff pool and moves the heat of redistributed water.

The solver calls the hooks through the SoilModel/LandModel forwarding methods and the SoilEnergyWaterCarbon methods in soil_coupled.jl.

References

  • [3] Celia, Bouloutas, and Zarba, Water Resources Research 26, 1483–1496 (1990)
source
implicit_operator(
    energy::SoilThermodynamics,
    hydrology::Terrarium.AbstractSoilHydrology,
    state,
    grid
)

Return the implicit operator of the soil internal energy: a TridiagonalOperator on ground_domain(grid) for heat conduction in temperature, with a face velocity for upwind advection of heat by the water flux if the hydrology lets the water flow (SoilHydrology with RichardsEq). Validates on the host that the freeze curve is FreeWater (the only one with an implemented regularized apparent heat capacity) and that temperature has no value boundary conditions (not yet supported). Flux boundary conditions on internal_energy and gradient boundary conditions on temperature are supported; the latter prescribe the diffusive flux -κ ∂T/∂z through the boundary face, which enters the explicit tendency via compute_boundary_energy_tendency.

The operator is the backward Euler step of ∂U/∂t = ∂z(κ ∂z T) - ∂z(C_w T j_w) + G_exp in temperature, with first-order upwind advection by the vertical water flux j_w (T_up the temperature of the upwind cell of a face) and a conservative update of the energy (see the soil energy docs for the derivation). With the Picard solver, for iterations m = 0, …, M-1, starting from U⁰ = Uⁿ (plus any heat moved by the closure of the water step) and T⁰ ← closure(U⁰):

  1. linearize!: Tᵐ ← closure(Uᵐ), lagged κᵐ (faces), regularized apparent heat capacity C_appᵐ (centers), advective face velocity C_w j_w (faces), and rhs = C_appᵐ Tᵐ - Uᵐ + U⁺ with U⁺ = U⁰ + Δt G_exp;
  2. solve!: C_appᵐ Tᵐ⁺¹ - Δt ∂z(κᵐ ∂z Tᵐ⁺¹) + Δt ∂z(C_w j_w T_upᵐ⁺¹) = rhs;
  3. implicit_update!: Uᵐ⁺¹ = U⁺ + Δt [∂z(κᵐ ∂z Tᵐ⁺¹) - ∂z(C_w j_w T_upᵐ⁺¹)] through the interior faces only;
  4. implicit_closure!: Tᵐ⁺¹ and the liquid water fraction ← closure(Uᵐ⁺¹).

The solver calls the hooks through the SoilModel/LandModel forwarding methods and the SoilEnergyWaterCarbon methods in soil_coupled.jl.

source
implicit_operator(
    soil::Terrarium.AbstractSoil,
    state,
    grid,
    names::Tuple{Vararg{Symbol}}
) -> Terrarium.SplitOperator

Return the implicit operator of the soil prognostic variables names on grid: a SplitOperator of per-variable operators that solves saturation_water_ice before internal_energy, whatever the order of names. The heat step then advects heat with the water flux of the current step (water_flux, see advective_heat_flux) and sees the end-of-step saturation, while the Richards step sees the beginning-of-step temperature and liquid water fraction (the effect of frozen ground on the hydraulic conductivity is lagged by one step). Other names follow in their given order and are rejected by the per-variable fallback.

source
implicit_operator(soil::AbstractSoil, state, grid, ::Val{name})

Return the implicit operator of the soil prognostic variable name on grid. SoilEnergyWaterCarbon returns that of the energy process for internal_energy and of the hydrology process for saturation_water_ice; this fallback throws an ArgumentError.

source
implicit_operator(snow::AbstractSnow, state, grid, ::Val{name})

Return the implicit operator of the snow prognostic variable name: a DiagonalOperator for the meltwater drainage of snow_water_equivalent (see linearize!). The fallback throws an ArgumentError.

source
implicit_operator(
    snow::SingleLayerSnow,
    state,
    grid,
    ::Val{:snow_water_equivalent}
) -> DiagonalOperator

Allocate the DiagonalOperator of the implicit meltwater drainage of snow_water_equivalent, solved in the liquid fraction snow_liquid_fraction (see implicit_variable).

Drainage removes liquid water at the zero-enthalpy reference (liquid water at 0 °C), so it changes neither the snow energy Ū nor the ice content W_i = (1 − θ) W = −Ū / (ρ_w L_f), clamped to [0, W]. With Ū⁺ and W⁺ = Wⁿ + Δt G_exp from the explicit stage (snowfall, rain, and sublimation), the backward Euler step of the drainage M_r(θ, W) (compute_meltwater_outflow) is a scalar equation per column in the liquid fraction θ, via W = W_i / (1 − θ):

G(θ) = W_i / (1 − θ) + Δt M_r(θ, W_i / (1 − θ)) − W⁺ = 0.

Each Picard iteration takes one Newton step: linearizing G about the iterate θᵐ, G(θ) ≈ G(θᵐ) + G′(θᵐ) (θ − θᵐ), and setting this to zero gives the diagonal system G′(θᵐ) θᵐ⁺¹ = G′(θᵐ) θᵐ − G(θᵐ) (see the snow mass balance docs for the convergence argument). Each iteration calls the hooks:

  1. linearize!: the diagonal G′(θᵐ) and right-hand side G′(θᵐ) θᵐ − G(θᵐ), at the liquid fraction θᵐ of the current iterate Wᵐ, bounded to [0, θ⁺];
  2. solve!: θᵐ⁺¹ in snow_liquid_fraction;
  3. implicit_update!: Wᵐ⁺¹ = W_i / (1 − θᵐ⁺¹) with min(L_c, θ⁺) ≤ θᵐ⁺¹ ≤ θ⁺, since the iteration starts from Wⁿ, which may lie below the root, and meltwater_flux ← (W⁺ − Wᵐ⁺¹) / Δt, the drainage passed on to the surface hydrology in the next step;
  4. implicit_closure!: the snow closure at the new iterate.

A fully liquid pack (W_i = 0, θ = 1) drains at min(K_sat, (1 − L_c) W / τ_m), so its step has the exact solution max(W⁺ − Δt K_sat, W⁺ / (1 + Δt (1 − L_c) / τ_m)), independent of the iterate.

source
Terrarium.timestepping — Function
timestepping(timestepper::AbstractTimeStepper)::Timestepping

Return the Timestepping trait (Explicit or Implicit) of the given timestepper. Every concrete timestepper must define this trait (e.g. timestepping(::ForwardEuler) = Explicit()); there is no default so that a new scheme declares its class explicitly. It is used, among other things, to route each sub-stepper of an AbstractIMEX to its slice of the IMEXCache.

timestepping(var::AbstractVariable, model::AbstractModel, timestepper::AbstractTimeStepper)::Timestepping

Return the Timestepping class with which the prognostic variable var of model is integrated under timestepper. Defaults to Explicit() for all variables; specialize this method (typically on an AbstractIMEX timestepper together with particular variable and/or model types) to route selected variables to the implicit sub-stepper.

source
timestepping(state::StateVariables, ::Val{name})
timestepping(state::StateVariables, var::AbstractVariable{name})

Return the resolved Timestepping class (Explicit or Implicit) of the prognostic variable name in state. Under an AbstractIMEX timestepper the class is read from the IMEXCache type parameter in state.timestepper_cache; for any other cache every variable is Explicit(), as are names that are not top-level prognostic variables (including variables in namespaces, which an IMEX does not route).

The lookup is resolved at compile time (type stable), so the trait can be passed to process methods (e.g. compute_tendencies!) that dispatch on it.

source
timestepping(timestepper::Optional{AbstractTimeStepper}, state::StateVariables, ::Val{name})

Return the Timestepping class with which the prognostic variable name of state is integrated by timestepper, for process methods that receive the model's timestepper. A single timestepper uses its own class (see timestepping(timestepper)), an AbstractIMEX the per-variable class in its IMEXCache (see timestepping(state, ::Val)), and nothing (e.g. standalone process evaluation) gives Explicit(). Type stable.

source
timestepping(cache::IMEXCache)

Return the resolved timestepping classes (a tuple of Explicit/Implicit instances in prognostic-variable order) held in the IMEXCache type parameter.

source
Terrarium.linearize! — Function
linearize!(operator::AbstractLinearOperator, rhs, state, model::AbstractModel, explicit_stage, Δt, ::Iteration{name})

Process hook of the implicit solver. Fill the lagged coefficients of operator and the right-hand side rhs for variable name at the current iterate in state, given the explicit stage explicit_stage (u⁺ = uⁿ + Δt G_exp). The fallback throws an ArgumentError.

source
linearize!(
    operator::TridiagonalOperator,
    rhs,
    state,
    grid,
    hydrology::SoilHydrology{NF, RichardsEq, SaturationClosure, SoilHydraulics} where {SaturationClosure<:Terrarium.AbstractSoilWaterClosure, SoilHydraulics<:(Terrarium.AbstractSoilHydraulics{NF})},
    soil::Terrarium.AbstractSoil,
    constants::PhysicalConstants,
    explicit_saturation,
    Δt
) where NF

Update the hydraulic conductivity (compute_hydraulics!), fill the lagged capacity and face conductivity of operator (compute_implicit_water_flow_coefficients!), and assemble the right-hand side C Ψᵐ - φ (ξᵐ - ξ⁺) in rhs (compute_implicit_water_flow_rhs), with ξ⁺ = explicit_saturation. [3] recommend at least two iterations.

References

  • [3] Celia et al., Water Resources Research 26, 1483–1496 (1990)
source
linearize!(
    operator::TridiagonalOperator,
    rhs,
    state,
    grid,
    energy::SoilThermodynamics,
    soil::Terrarium.AbstractSoil,
    constants::PhysicalConstants,
    explicit_energy,
    Δt
)

Linearize the implicit heat operator of the soil internal energy at the current iterate (linearly implicit backward Euler in temperature, conservative update of the internal energy). On entry, explicit_energy holds U⁺ = U⁰ + Δt G_exp, where G_exp contains the flux boundary conditions (e.g. SoilHeatFlux and GeothermalHeatFlux), the diffusive fluxes implied by gradient boundary conditions on temperature, the advective heat fluxes through the boundary faces, and the heat carried by water crossing the boundaries (compute_boundary_advected_heat!), but not the fluxes through the interior faces. This hook

  1. re-evaluates the energy closure, so that temperature and liquid water fraction are consistent with the internal energy and the current saturation (changed by the water step since tⁿ),
  2. evaluates the face conductivity κ (compute_implicit_conductivity!), the regularized apparent heat capacity C_app (compute_implicit_capacity!), and, with water flow, the advective face velocity C_w j_w from the water_flux of the step (compute_implicit_advective_velocity!), storing them in the face_conductivity, diagonal_coefficient, and face_velocity of operator,
  3. and assembles the right-hand side C_app Tᵐ - Uᵐ + U⁺ of C_app Tᵐ⁺¹ - Δt ∂z(κ ∂z Tᵐ⁺¹) + Δt ∂z(C_w j_w T_upᵐ⁺¹) = C_app Tᵐ - Uᵐ + U⁺ (compute_implicit_energy_rhs).

implicit_update! then updates the internal energy conservatively: for any number of iterations, the change of column-integrated energy is exactly Δt times the net boundary flux in G_exp.

source
linearize!(
    operator::DiagonalOperator,
    rhs,
    state,
    grid,
    snow::SingleLayerSnow,
    constants::PhysicalConstants,
    explicit_water,
    Δt,
    ::Terrarium.Iteration{:snow_water_equivalent}
)

Fill the diagonal and right-hand side of the Newton step for the liquid fraction of the meltwater drainage (compute_implicit_meltwater_coefficients!), given the explicit stage W⁺ = explicit_water.

source
Terrarium.implicit_update! — Function
implicit_update!(state, model::AbstractModel, operator::AbstractLinearOperator, explicit_stage, Δt, ::Iteration{name})

Process hook of the implicit solver. Update the prognostic variable name in state conservatively from the solved target and explicit_stage, using the fluxes implied by operator.

source
Terrarium.implicit_closure! — Function
implicit_closure!(state, model::AbstractModel, ::Iteration{name, target, final})

Process hook of the implicit solver. Re-evaluate the closure of variable name (and the auxiliary variables the next linearization needs) after implicit_update!, for the process that owns the variable only. A model typically forwards it to the regular process closure! of that process; only a closure with once-per-step side effects (e.g. routing excess water to runoff) needs its own method, which must restrict them to final = true (see the soil hydrology for an example).

source