Hybrid model taxonomy: what can be learned in an ocean model

The physical model is not one object. It is a stack of pieces: parameters, discretisation, terms, a tendency, an integrator, a step map, a rollout. Every hybrid or emulator learns one of those pieces and glues it back in one of a few ways. This page lists the ten that matter, on the same Gulf Stream QG example as the state-estimation page.

Notation

As on the state-estimation page, plus the pieces of the model.

Note the clash: \(\nu\) in the equations below is the viscosity, as is standard in QG. The variational parameters \(\gvect{\nu}\) of the last page never appear on this page.

The physical model, taken apart

The 1.5-layer QG model of the last page, written layer by layer. Let \(\nabla := \nabla_{\vect{s}}\) act on the horizontal coordinates \(\vect{s} = (s_1, s_2)\), \(\mathcal{A} = \nabla^2 - L_d^{-2}\) be the Helmholtz operator and \(J(a, b) = \partial_{s_1} a \partial_{s_2} b - \partial_{s_2} a \partial_{s_1} b\) the Jacobian.

Below the terms sits the discretisation: \(J\) is the Arakawa stencil and \(\nabla^2\) the five-point stencil on the C-grid, and \(\mathcal{A}^{-1}\) is a spectral or sine-transform solve. Each layer is a function with a signature, and each can be swapped.

The sketch below is the model in code. Every type in this page is an edit to one or two of its lines.

import equinox as eqx
import finitevolx as fvx
import jax
from jaxtyping import Array, Float

# Shapes: T = 10 days (time); H, W = 130 × 130 grid with ghost ring
#         (height, width); V covariates, β s_2 first (variable);
#         48 RK4 substeps per day


class QGParams(eqx.Module):
    """Physical parameters θ = (L_d, ν) of the 1.5-layer QG model."""

    deformation_radius: Float[Array, ""]
    viscosity: Float[Array, ""]


def qg_tendency(
    pv: Float[Array, "height width"],
    covariates: Float[Array, "variable height width"],
    forcing: Float[Array, "height width"],
    params: QGParams,
    grid: fvx.CartesianGrid2D,
) -> Float[Array, "height width"]:
    # today (finitevolx)
    difference = fvx.Difference2D(grid)
    # today (finitevolx)
    jacobian = fvx.ArakawaJacobian2D(grid)
    beta_y = covariates[0]                            # (H, W)

    # --- diagnostic solve ---
    # q − β s_2
    relative_pv = pv - beta_y                         # (H, W) → (H, W)
    # L_d²
    radius_squared = params.deformation_radius**2     # → ()
    # L_d⁻²
    helmholtz_shift = 1.0 / radius_squared            # → ()
    # today* (finitevolx)
    # ψ = 𝒜⁻¹[q − β s_2],  𝒜 = ∇² − L_d⁻²
    streamfunction = fvx.pv_inversion(                # (H, W) → (H, W)
        relative_pv, grid.dx, grid.dy, lambda_=helmholtz_shift, bc="dst"
    )

    # --- terms ---
    # today (finitevolx)
    # J(ψ, q)
    jacobian_term = jacobian(streamfunction, pv)      # (H, W)
    # advection = −J(ψ, q)
    advection = -jacobian_term                        # (H, W)
    # today (finitevolx)
    # ∇²ψ
    laplacian_1 = difference.laplacian(streamfunction)  # (H, W)
    # today (finitevolx)
    # ∇⁴ψ = ∇²(∇²ψ)
    laplacian_2 = difference.laplacian(laplacian_1)   # (H, W)
    # today (finitevolx)
    # ∇⁶ψ = ∇²(∇⁴ψ)
    laplacian_3 = difference.laplacian(laplacian_2)   # (H, W)
    # dissipation = −ν ∇⁶ψ
    dissipation = -params.viscosity * laplacian_3     # (H, W)

    # ℱ[q] = −J(ψ, q) − ν∇⁶ψ + u^wind
    tendency = advection + dissipation + forcing      # (H, W)
    return tendency


def qg_day(
    pv: Float[Array, "height width"],
    covariates: Float[Array, "variable height width"],
    forcing: Float[Array, "height width"],
    params: QGParams,
    grid: fvx.CartesianGrid2D,
    *,
    dt: float = 1800.0,
    n_substeps: int = 48,
) -> Float[Array, "height width"]:
    def rhs(state):
        # user
        # ℱ[q; u_k, θ]
        return qg_tendency(state, covariates, forcing, params, grid)

    def substep(state, _):
        # today (finitevolx)
        # q ← RK4_Δt[ℱ](q), classic RK4
        state = fvx.rk4_step(state, rhs, dt)          # (H, W) → (H, W)
        return state, None

    # f(x_k, u_k; θ) = RK4_Δt[ℱ]⁴⁸(x_k)
    next_pv, _ = jax.lax.scan(                        # (H, W) → (H, W)
        substep, pv, None, length=n_substeps
    )
    return next_pv


def qg_rollout(
    initial_pv: Float[Array, "height width"],
    covariates: Float[Array, "variable height width"],
    forcing: Float[Array, "time height width"],
    params: QGParams,
    grid: fvx.CartesianGrid2D,
) -> Float[Array, "time height width"]:
    def day(state, daily_forcing):
        # user
        # x_k = f(x_{k−1}, u_k; θ)
        state = qg_day(                               # (H, W) → (H, W)
            state, covariates, daily_forcing, params, grid
        )
        return state, state

    # f_{1:T←0}(x₀, u_{1:T}; θ) = (x₁, …, x_T)
    _, trajectory = jax.lax.scan(                     # (H, W) → (T, H, W)
        day, initial_pv, forcing
    )
    return trajectory

Markers: today ships on main (checked 24 Oct 2026: finitevolX 831c69f, gaussx eeda235, somax 9cb19b3); today* ships, arguments sketched; user is defined on these pages; not pinned is a library whose API was not checked. Shape column letters are defined in the header.

somax ships barotropic and baroclinic QG models but no 1.5-layer one, so the tendency here is page-defined. BaroclinicQG.create(n_layers=1) may be the shipped route, but it is untested.

Three axes

Which layer is learned

The diagram above. Lower layers (T1, T2) keep the most physics and fix the least; upper layers (T8, T9) keep the least and are fastest.

How the learned piece is glued in

What it is trained against

Two jobs, three sources of reference, four loss levels. Let \(\vect{x}^\text{ref}\) be the reference trajectory and \(\mathcal{F}_\text{ref}\) its tendency.

The loss can sit at four levels, from cheapest to most faithful to how the model will be used:

For correction against a finer run, the tendency target is the subgrid forcing: what the coarse physics misses when applied to the coarse-grained state,

The tendency and one-step losses are a priori: cheap, but the model never sees its own errors, so long rollouts often go unstable. The rollout and observation losses are a posteriori: they differentiate through the solver, which is what a JAX model makes affordable. The observation loss also needs each window's \(\vect{x}_0\), so learning from data is 4D-Var with \(\gvect{\theta}_\text{emu}\) added to the control vector.

The ten types

Each card gives the equation, the operator signature with what is learned, and the training target as a column. Shapes use the named dimensions of the last page.

Selection matrix

● yes   ◐ partly or with care   ○ no

"Faster" is for this QG model, where the physics is already cheap. T2 is faster because a learned stencil lets a coarser grid match a finer one, not because a stencil is quicker. "Can fix model error ◐" for T8–T10 means only when trained on reanalysis or data instead of step 1.

What the structure implies

Where each type lives