Learned WENO: an advection scheme that learns where to dissipate

The advection scheme is where a coarse ocean model loses information: it smooths fronts, dissipates eddies, and decides whether a run stays stable. This page keeps everything that makes WENO safe (flux form, fixed candidate stencils, convex weights, land masks) and learns one thing: how to weigh the candidates at each face. Same Gulf Stream QG example as the other pages.

The problem

The QG run of the other pages has \(\Delta x = 8\) km and a deformation radius \(L_d \approx 30\) km, so an eddy is about four grid cells across. The Gulf Stream itself is a PV front a few cells wide. Two classical choices both fail here:

The question this page asks: keep WENO5's structure, and learn how much to trust each candidate stencil from a high-resolution run of the same ocean.

Notation

Written in one dimension, along \(x\); the \(y\) direction is identical with \(v\).

Flux form: conservation for free

Integrate \(\partial_t q + \partial_x(uq) = 0\) over cell \(i\). The cell average changes only through its two faces:

Summing over all cells, every interior face appears once with each sign, so \(\frac{\mathrm{d}}{\mathrm{d}t}\sum_i q_i \Delta x\) depends only on the domain walls, where \(u=0\). Total PV is conserved whatever rule produces \(\hat q\). That is why the learning goes into \(\hat q\), never into a gradient: an advective form \(u \partial_x q\) with a learned gradient has no such identity.

Every scheme is a centred part plus a dissipation

Take the classical fifth-order upwind reconstruction for \(u\ge0\), on the points \(q_{i-2},\ldots,q_{i+3}\), and the sixth-order centred one on the same points:

Their difference is a fifth difference, so the upwind flux is the centred flux plus a term proportional to \(\vert u \vert\):

For a wave \(q = e^{\mathrm{i} kx}\), the centred part only moves it (zero damping); the second part damps it at rate \(\vert u \vert \Delta x^5 k^6/60\) for \(k\Delta x \ll 1\). In continuous form, fifth-order upwinding is the exact equation plus a sixth-order hyperdiffusion:

So "upwinding" is a choice of dissipation, and the dissipation is what decides stability and what the model forgets. The cleanest way to control it is to write any face flux in this split, with a non-negative face diffusivity \(\kappa\) (dimensionless, in units of \(\Delta x\)):

With periodic or walled boundaries, the second term changes the variance \(\frac{1}{2}\sum_i q_i^2\) by \(-\sum_i \vert u \vert \kappa (q_{i+1}-q_i)^2 /\Delta x \le 0\). It can only remove variance. The learned \(\kappa_{\gvect{\theta}}\) on this page enters exactly here.

Why a single learned kernel cannot win

A learned gradient or face kernel with fixed weights is a linear scheme: \(\hat q = \sum_j w_j q_{i+j}\). Godunov's theorem says a linear scheme that never creates new maxima or minima is at most first-order accurate. Training a fixed kernel therefore picks a point on one trade-off curve:

A fixed kernel can still tune the shape of the error across wavenumbers (how much each scale is delayed or damped), and that is a useful baseline. To be sharp in smooth water and safe at the front at the same time, the weights have to depend on the local field. That is nonlinearity, and WENO is the classical way to do it.

Classical WENO5, from local basis functions

Read WENO the way a neural operator is read: move the data into a basis, act on the coefficients, evaluate the result where it is needed. Here the basis is local: one small polynomial per candidate stencil.

Cell averages and a local coordinate

A finite-volume model stores cell averages, \(q_j = \frac{1}{\Delta x}\int_{\text{cell } j} q \mathrm{d}x\). Measure position from the centre of cell \(i\) in cell widths, \(\xi = (x - x_i)/\Delta x\). Cell \(i+j\) is then the interval \([j-\frac{1}{2}, j+\frac{1}{2}]\), and the face \(i+\frac12\) sits at \(\xi = \frac{1}{2}\). For \(u\ge0\) the face sees the window

Step 1 · A local basis

Each candidate \(r\) is one quadratic, written in the monomial basis \(\gvect{\varphi}(\xi) = (1, \xi, \xi^2)^\top\) with its own coefficients \(\vect{a}_r\):

Candidate \(r\) is fitted on three cells: \(\mathcal{S}_1 = \{i-2, i-1, i\}\), \(\mathcal{S}_2 = \{i-1, i, i+1\}\), \(\mathcal{S}_3 = \{i, i+1, i+2\}\).

Step 2 · Lift: cell averages to coefficients

Ask the quadratic to have the right average on each of its three cells. The average of \(\gvect{\varphi}\) over cell \(i+j\) is \((1, j, j^2 + \frac{1}{12})\), so this is a 3 × 3 linear system, solved once:

\(\mat{E}_r\) is \(\mat{A}_r^{-1}\) placed into the columns of its three cells, a fixed 3 × 5 matrix. For the centred candidate the coefficients are a mean, a slope and a curvature:

Step 3 · Evaluate at the face

The candidate's face value is its polynomial at \(\xi = \frac{1}{2}\):

The same thing in a cardinal basis: \(\ell^r_j(\xi) = [\gvect{\varphi}(\xi)^\top \mat{E}_r]_j\) is the quadratic whose average is 1 on cell \(j\) and 0 on the other two. Then \(p_r(\xi) = \sum_{j\in\mathcal{S}_r} q_j \ell^{r}_j(\xi)\), and every stencil coefficient is a basis function evaluated at the face: \(c^{r}_j = \ell^{r}_j(\frac{1}{2})\).

Step 4 · Smoothness is energy in the non-constant modes

The smoothness indicator measures the slope and curvature of each candidate over the face's own cell \(i\), in units of the cell width:

(\(a_{r,1}\) and \(a_{r,2}\) are the slope and half-curvature at the centre of cell \(i\); for candidates 1 and 3 that is the end of their stencil.) The constant mode has weight 0, so adding a constant to \(q\) never changes \(\beta\). This is the idea behind modal smoothness indicators in spectral methods: roughness is how much energy sits in the higher modes. In the window, \(\beta_r = \vect{q}^\top\mat{M}_r\vect{q}\) with \(\mat{M}_r = \mat{E}_r^\top\mat{S} \mat{E}_r\); expanded, these are the textbook formulas:

Step 5 · Mix

The nonlinear weights favour smooth candidates, with \(\tilde\beta_r = \beta_r/\sum_s\beta_s\):

With the linear weights, the mixture is exactly fifth-order upwinding: \(\sum_r d_r \vect{c}^r = \frac{1}{60}(2, -13, 47, 27, -3)\). In smooth water all \(\beta_r\) are similar and \(\gvect{\omega} \to \vect{d}\). Across a front, the candidates that straddle it have large \(\beta\) and get almost no weight. finitevolX normalises by \(\sum\beta\) (held out of the gradient), so rescaling \(q\) rescales \(\hat q\) exactly and nothing overflows in float32.

Two windows, worked through

Dashed curves are the three candidate quadratics \(p_r(\xi)\) over their own cells, solid bars the cell averages, small dots the candidate face values, and the large dot the blended \(\hat q\).

The whole scheme as an operator

window q            (q_{i−2} … q_{i+2})                (P,)
   │  lift          aᵣ = Eᵣ q                            (R, B)
   ▼
coefficients        mean, slope, curvature per candidate
   │  score         βᵣ = aᵣᵀ S aᵣ                        (R,)
   │  route         ω = softmax(log d − p log β̃)         (R,)
   │  evaluate      q̂ʳ = φ(½)ᵀ aᵣ                        (R,)
   ▼
face value          q̂ = ∑ᵣ ωᵣ q̂ʳ                         ()

In code, the lift, score and evaluate steps are three contractions over fixed tensors:

# Shapes: R = 3 candidates (candidate); P = 5 window points (point);
#         B = 3 basis functions 1, ξ, ξ² (basis); H, W grid (height, width)

# user
# Eᵣ = Aᵣ⁻¹ placed in the window: lift, cell averages → coefficients
lift = WENO5_LIFT                           # (R, B, P)
# user
# φ(½) = (1, ½, ¼): evaluate the basis at the face
phi_face = jnp.array([1.0, 0.5, 0.25])      # (B,)
# user
# S = diag(0, 1, 13/3): roughness of each basis function
roughness = jnp.array([0.0, 1.0, 13.0 / 3.0])  # (B,)

# --- lift ---
# aᵣ = Eᵣ q
coeffs = einx.dot(                          # (R, B, P), (P, H, W) → (R, B, H, W)
    "candidate basis point, point height width"
    " -> candidate basis height width",
    lift,
    window,
)

# --- score ---
# aᵣ,ₘ²
coeffs_squared = coeffs**2                  # (R, B, H, W)
# βᵣ = aᵣᵀ S aᵣ = aᵣ,₁² + (13/3) aᵣ,₂²
beta = einx.dot(                            # (R, B, H, W), (B,) → (R, H, W)
    "candidate basis height width, basis -> candidate height width",
    coeffs_squared,
    roughness,
)

# --- evaluate ---
# q̂ʳ = φ(½)ᵀ aᵣ
candidates = einx.dot(                      # (R, B, H, W), (B,) → (R, H, W)
    "candidate basis height width, basis -> candidate height width",
    coeffs,
    phi_face,
)

Written as a function of position inside cell \(i\), the reconstruction is a sum of coefficients times basis functions, which is the form neural operators use:

In DeepONet terms, the trunk is the fixed polynomial basis and the branch is \(\omega_r(\vect{q}) a_{r,m}(\vect{q})\). Every learned part of learned WENO sits in the routing \(\omega\); the basis, the lift and the evaluation stay fixed, and that is what keeps consistency and order.

The hand-chosen parts of classical WENO are the exponent 2, the floor \(\varepsilon\), and the linear weights. They decide where WENO dissipates, and they were tuned on shock-tube problems, not on ocean eddies.

Learned WENO

Write the WENO weights as a softmax over logits, then add a learned correction to the logits:

Since \(\mathrm{softmax}(\log\alpha) = \alpha/\sum\alpha\), the first term alone is exactly classical WENO5. The feature vector \(\vect{e}\) at each face holds only scale-free numbers:

Three design choices carry the safety:

• Zero initialisation. The last layer of \(g_\text{corr}\) starts at zero, so the untrained scheme is WENO5. Training only moves away from it as far as the data asks.
• Gating by \(\tilde\tau\). In smooth water \(\tilde\tau\to0\), the correction vanishes, and the weights return to WENO5's, which keeps fifth order where the field is smooth. The learning acts only near fronts and under-resolved eddies.
• Pointwise network. \(g_\text{corr}\) is a 1×1 convolution over the feature axis: the same small function at every face. The scheme is translation-equivariant and transfers to other domains and grid sizes.

A second head gives the extra dissipation \(\kappa_{\gvect{\theta}} = \mathrm{softplus}(g_\text{diss}(\vect{e}; \gvect{\theta})) \ge 0\), added through the variance-safe term of the flux. For \(u \lt 0\) the window is mirrored (\(q_{i+3},\ldots,q_{i-1}\)) and the same network is applied, so the scheme has no preferred direction.

Land: the same mechanism

Candidate \(r\) reads three cells. If any of them is dry, its logit is set to \(-\infty\) before the softmax, and the remaining weights renormalise:

In open water all three survive. One cell from a coast, only the candidate pointing away from land survives, and the face value is third order. At the coast itself the normal velocity is zero, so the flux vanishes whatever the reconstruction. finitevolX already does a hard version of this: Advection2D with a Mask2D picks, per upwind cell, the widest stencil tier whose cells are all wet (sizes 2 to 10), using the counts of contiguous wet cells in mask.stencil_capability. Learned WENO uses the same counts twice: to mask candidates, and as features \(n^\pm\), so the network can learn boundary-current behaviour (for example, more dissipation along the western boundary, where the grid is most under-resolved).

If the window reaches past the narrowest candidate (a single wet cell between two coasts), the face falls back to first-order upwind, as the finitevolX tiers already do.

What holds for every value of θ

Parallels

Strip WENO to its skeleton and it is a pattern machine learning uses everywhere: a few fixed local linear estimators, mixed by input-dependent normalised weights.

The four strongest parallels follow. All numbers use the two windows of the figures above, so they line up across sections. First, the exponent \(p\) swept from 0 to ∞ on both windows:

On the front, candidates 1 and 2 are both perfectly flat (\(\beta_1 = \beta_2 = 0\)), so for any \(p\ge1\) they keep their prior ratio 0.1 : 0.6 = 1 : 6. The table is the whole trade-off in one place: \(p = 0\) is most accurate in smooth water and undershoots at the front; raising \(p\) buys safety at the front with accuracy in smooth water.

p → ∞   hard attention (argmax)   →  ENO: take the smoothest stencil
p = 2   soft attention            →  WENO5
p = 0   attention to the bias     →  linear 5th-order upwind (ω = d)

# Shapes: R = 3 candidates (candidate); P = 5 points (point); H, W grid

# --- transformer attention, one query over R keys ---
# sᵣ = ⟨h_qry, h_key,r⟩ / √D_key + bᵣ
scores_attn = dot_scores / sqrt_dim + position_bias  # (R, H, W)

# --- WENO, same slot ---
# −p log(β̃ᵣ + ε): content score, "how smooth does candidate r see the field"
content = -temperature_inv * log_beta       # (R, H, W)
# sᵣ = log dᵣ − p log(β̃ᵣ + ε): bias log d plays the position bias
scores = einx.add(                          # (R,), (R, H, W) → (R, H, W)
    "candidate, candidate height width -> candidate height width",
    log_linear_weights,
    content,
)
# ωᵣ = softmax_r(sᵣ)
weights = einx.softmax("[candidate] height width", scores) #| (R, H, W)
# values h_val,r = q̂ʳ = ⟨cʳ, q⟩, fixed (never learned)
values = einx.dot(                          # (R, P), (P, H, W) → (R, H, W)
    "candidate point, point height width -> candidate height width",
    candidate_coeffs,
    window,
)
# q̂ = ∑ᵣ ωᵣ q̂ʳ
face = einx.dot(                            # (R, H, W) ×2 → (H, W)
    "candidate height width, candidate height width -> height width",
    weights,
    values,
)

# Shapes: R = 3 candidates (candidate); P = 5 points (point);
#         K = 64 test wavenumbers (wavenumber); H, W grid

# --- mixture of experts: mix the outputs ---
# q̂ = ∑ᵣ ωᵣ ⟨cʳ, q⟩
face_moe = einx.dot(                        # → (H, W)
    "candidate height width, candidate height width -> height width",
    weights,
    values,
)

# --- conditional convolution: mix the kernels, then apply once ---
# c(q) = ∑ᵣ ωᵣ cʳ, one 5-tap kernel per face
kernel = einx.dot(                          # (R, H, W), (R, P) → (P, H, W)
    "candidate height width, candidate point -> point height width",
    weights,
    candidate_coeffs,
)
# q̂ = ⟨c(q), q⟩, identical to face_moe because the experts are linear
face_conv = einx.dot(                       # (P, H, W) ×2 → (H, W)
    "point height width, point height width -> height width",
    kernel,
    window,
)

# --- diagnostic: local Fourier symbol of c, a map of numerical damping ---
# k j for test wavenumbers k Δx and offsets j = −2…2
angles = einx.multiply(                     # (K,), (P,) → (K, P)
    "wavenumber, point -> wavenumber point", k_dx, offsets
)
# e^{i k j}
phase = jnp.exp(1j * angles)                # (K, P)
# ĉ(k) = ∑ⱼ cⱼ e^{i k j}
symbol = einx.dot(                          # (K, P), (P, H, W) → (K, H, W)
    "wavenumber point, point height width -> wavenumber height width",
    phase,
    kernel,
)
# user
# λ(k) = −ĉ(k)(1 − e^{−ik}); Re λ < 0 is damping
growth = semi_discrete_rate(symbol, k_dx)   # (K, H, W)

# Shapes: R = 3 candidates (candidate); H, W grid

# --- WENO: power-law evidence ---
# log ωᵣ = log dᵣ − p log(β̃ᵣ + ε) + const
log_evidence_power = -2.0 * jnp.log(beta_norm + eps)  # (R, H, W)

# --- Bayesian model averaging: exponential (Gibbs) evidence ---
# log p(v | r) = −β̃ᵣ / T + const, T a temperature
log_evidence_gibbs = -beta_norm / temperature  # (R, H, W)
# log ℙ(r | q) = log dᵣ + log ℙ(q | r) + const
log_posterior = einx.add(                   # (R,), (R, H, W) → (R, H, W)
    "candidate, candidate height width -> candidate height width",
    log_linear_weights,
    log_evidence_gibbs,
)
# ℙ(r | q) = softmax_r(log ℙ)
weights = einx.softmax("[candidate] height width", log_posterior)  # (R, H, W)

# --- model uncertainty per face ---
# log ωᵣ
log_weights = jnp.log(weights + 1e-30)      # (R, H, W)
# ωᵣ log ωᵣ
plogp = weights * log_weights               # (R, H, W)
# ℍ[ω] = −∑ᵣ ωᵣ log ωᵣ, high where no stencil is clearly right
entropy = -einx.sum("candidate height width -> height width", plogp)  # (H, W)

# Shapes: H, W grid (height, width); E = 6 WENO features (feature)

# --- gradient at cell centres ---
# today (finitevolx)
# ∂ₓq at U-points
grad_x_faces = difference.diff_x_T_to_U(q)  # (H, W) → (H, W)
# today (finitevolx)
# ∂ᵧq at V-points
grad_y_faces = difference.diff_y_T_to_V(q)  # (H, W) → (H, W)
# user
# average the two faces of each cell back to its centre
grad_x = faces_to_center_x(grad_x_faces)    # (H, W) → (H, W)
# user
grad_y = faces_to_center_y(grad_y_faces)    # (H, W) → (H, W)

# --- local structure tensor M_struct = ∑_{3×3} ∇q ∇qᵀ ---
# ∂ₓq², ∂ₓq ∂ᵧq, ∂ᵧq²
gxx = grad_x * grad_x                       # (H, W)
gxy = grad_x * grad_y                       # (H, W)
gyy = grad_y * grad_y                       # (H, W)
# user
# M_xx, M_xy, M_yy summed over the 3×3 neighbourhood
jxx, jxy, jyy = box_sum_3x3(gxx, gxy, gyy)  # (H, W) ×3

# --- RAISR-style features ---
# M_xx + M_yy = σ²₊ + σ²₋
trace = jxx + jyy                           # (H, W)
# M_xx − M_yy
diagonal_gap = jxx - jyy                    # (H, W)
# 2 M_xy
off_diagonal = 2.0 * jxy                    # (H, W)
# σ²₊ − σ²₋ = √((M_xx − M_yy)² + 4 M_xy²)
eigen_gap = jnp.hypot(diagonal_gap, off_diagonal)  # (H, W)
# coherence = (σ²₊ − σ²₋)/(σ²₊ + σ²₋): 1 for a straight front, 0 for no direction
coherence = eigen_gap / (trace + 1e-30)     # (H, W)
# 2θ_dir = atan2(2 M_xy, M_xx − M_yy): front orientation, sign-free
two_theta = jnp.arctan2(off_diagonal, diagonal_gap)  # (H, W)
# (cos 2θ_dir, sin 2θ_dir): continuous, no wrap at ±π
orient_cos = jnp.cos(two_theta)             # (H, W)
orient_sin = jnp.sin(two_theta)             # (H, W)

# --- add to the router's features ---
# e′ = (e, cos 2θ_dir, sin 2θ_dir, coherence), all scale-free
features_2d = einx.rearrange(               # → (E + 3, H, W)
    "feature height width, height width, height width, height width"
    " -> (feature + 1 + 1 + 1) height width",
    features,
    orient_cos,
    orient_sin,
    coherence,
)

In one sentence: WENO is soft attention over a fixed bank of local polynomial experts, scored by their roughness. Learned WENO learns the score and keeps the experts.

As operators

Code sketch

The x-direction face flux of learned WENO5. WENO5_CANDIDATES holds the three candidate rows of the table above, padded to the five-point window; WENO5_BETA_MATRICES holds the three \(\mat{M}_r\). stencil_windows_x, candidate_wet_x and shift_east are short page-level helpers for the shifted views and the per-candidate wet test.

import einx
import equinox as eqx
import finitevolx as fvx
import jax
import jax.numpy as jnp
from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray

# Shapes: H, W = 130 × 130 grid with ghost ring (height, width);
#         P = 5 stencil points q_{i−2} … q_{i+2} (point);
#         R = 3 candidate stencils (candidate); E = 6 features (feature)


class LearnedWENO5(eqx.Module):
    """WENO5 whose candidate weights get a learned, gated logit correction.

    ωᵣ = softmax_r( log dᵣ − 2 log(β̃ᵣ + ε) + s(τ̃) g_corr(features; θ)ᵣ ).

    Algorithm:
        1. Build the three 3-point candidate face values q̂ʳ.
        2. Score each candidate's smoothness βᵣ and normalise: β̃ᵣ = βᵣ / ∑β.
        3. Classical WENO5 logits, plus a learned correction gated by the
           WENO-Z smoothness detector τ̃, so smooth flow keeps WENO5 exactly.
        4. Zero the weight of any candidate that reads a dry cell.
        5. Blend: q̂ = ∑ᵣ ωᵣ q̂ʳ, and add a learned non-negative diffusivity.

    Args:
        weight_net: Pointwise network g_corr, features → R logit corrections.
        diffusivity_net: Pointwise network, features → one raw diffusivity.
        candidate_coeffs: Fixed cʳⱼ, padded to the 5-point window.
        beta_matrices: Fixed quadratic forms with βᵣ = qᵀ Mᵣ q.
        log_linear_weights: log dᵣ, with d = (0.1, 0.6, 0.3).
    """

    weight_net: eqx.nn.Conv2d
    diffusivity_net: eqx.nn.Conv2d
    candidate_coeffs: Float[Array, "candidate point"]
    beta_matrices: Float[Array, "candidate point point"]
    log_linear_weights: Float[Array, "candidate"]
    eps_rel: float = 1e-6

    def __init__(self, *, key: PRNGKeyArray):
        weight_key, diffusivity_key = jax.random.split(key)
        # user
        weight_net = eqx.nn.Conv2d(6, 3, kernel_size=1, key=weight_key)
        # zero last layer: g_corr ≡ 0 at init, so training starts from WENO5
        self.weight_net = eqx.tree_at(
            lambda net: net.weight, weight_net, jnp.zeros_like(weight_net.weight)
        )
        # user
        self.diffusivity_net = eqx.nn.Conv2d(6, 1, kernel_size=1, key=diffusivity_key)
        self.candidate_coeffs = WENO5_CANDIDATES
        self.beta_matrices = WENO5_BETA_MATRICES
        self.log_linear_weights = jnp.log(jnp.array([0.1, 0.6, 0.3]))

    def face_value(
        self,
        window: Float[Array, "point height width"],
        courant: Float[Array, "height width"],
        wet_run: Int[Array, "2 height width"],
        candidate_wet: Bool[Array, "candidate height width"],
    ) -> tuple[Float[Array, "height width"], Float[Array, "height width"]]:
        # --- candidates ---
        # q̂ʳ = ∑ⱼ cʳⱼ q_{i+j}
        candidates = einx.dot(                        # (R, P), (P, H, W) → (R, H, W)
            "candidate point, point height width -> candidate height width",
            self.candidate_coeffs,
            window,
        )

        # --- smoothness ---
        # βᵣ = qᵀ Mᵣ q
        betas = einx.dot(                             # → (R, H, W)
            "candidate p1 p2, p1 height width, p2 height width"
            " -> candidate height width",
            self.beta_matrices,
            window,
            window,
        )
        # ∑β, held out of the gradient: the weights depend on ratios only
        beta_sum = einx.sum(                          # (R, H, W) → (H, W)
            "candidate height width -> height width", betas
        )
        beta_sum = jax.lax.stop_gradient(beta_sum)    # (H, W)
        # ∑β + ε_tiny, so a flat window gives the linear weights
        beta_sum = beta_sum + 1e-30                   # (H, W)
        # β̃ᵣ = βᵣ / ∑β
        betas_normalised = einx.divide(               # (R, H, W), (H, W) → (R, H, W)
            "candidate height width, height width -> candidate height width",
            betas,
            beta_sum,
        )
        # τ̃ = |β̃₁ − β̃₃|, the WENO-Z detector: ≈ 0 where the field is smooth
        beta_gap = betas_normalised[0] - betas_normalised[2]  # → (H, W)
        detector = jnp.abs(beta_gap)                  # (H, W)

        # --- features ---
        # e = (β̃₁, β̃₂, β̃₃, C, n⁻, n⁺), all scale-free
        wet_run_float = wet_run.astype(window.dtype)  # (2, H, W)
        features = einx.rearrange(                    # → (E, H, W)
            "candidate height width, height width, two height width"
            " -> (candidate + 1 + two) height width",
            betas_normalised,
            courant,
            wet_run_float,
        )

        # --- weights ---
        # β̃ᵣ + ε
        shifted_beta = betas_normalised + self.eps_rel  # (R, H, W)
        # log(β̃ᵣ + ε)
        log_beta = jnp.log(shifted_beta)              # (R, H, W)
        # 2 log(β̃ᵣ + ε)
        log_beta_twice = 2.0 * log_beta               # (R, H, W)
        # log αᵣ = log dᵣ − 2 log(β̃ᵣ + ε)
        classic_logits = einx.subtract(               # (R,), (R, H, W) → (R, H, W)
            "candidate, candidate height width -> candidate height width",
            self.log_linear_weights,
            log_beta_twice,
        )
        # user
        # g_corr(e; θ)
        correction = self.weight_net(features)        # (E, H, W) → (R, H, W)
        # s(τ̃) g_corr(e; θ), gated off where the field is smooth
        gated = einx.multiply(                        # → (R, H, W)
            "height width, candidate height width -> candidate height width",
            detector,
            correction,
        )
        logits = classic_logits + gated               # (R, H, W)
        # candidates that read a dry cell get logit −∞
        logits = jnp.where(candidate_wet, logits, -1e30)  # (R, H, W)
        # ωᵣ = softmax_r(logits)
        weights = einx.softmax("[candidate] height width", logits)  # (R, H, W)
        # q̂ = ∑ᵣ ωᵣ q̂ʳ
        face = einx.dot(                              # (R, H, W), (R, H, W) → (H, W)
            "candidate height width, candidate height width -> height width",
            weights,
            candidates,
        )

        # --- extra dissipation ---
        # user
        raw = self.diffusivity_net(features)          # (E, H, W) → (1, H, W)
        # κ = softplus(raw) ≥ 0
        kappa = jax.nn.softplus(raw[0])               # (H, W)
        return face, kappa

    def flux_x(
        self,
        q: Float[Array, "height width"],
        u: Float[Array, "height width"],
        courant: Float[Array, "height width"],
        mask: fvx.Mask2D,
    ) -> Float[Array, "height width"]:
        capability = mask.stencil_capability
        # user
        # windows q_{i−2…i+2} (u ≥ 0) and q_{i+3…i−1} (u < 0, mirrored)
        window_pos, window_neg = stencil_windows_x(q)  # (H, W) → (P, H, W) ×2
        # user
        wet_pos, wet_neg = candidate_wet_x(capability)  # → (R, H, W) ×2
        # today (finitevolx)
        # n⁻, n⁺: contiguous wet cells upstream and downstream
        runs_pos = einx.rearrange(                    # (H, W) ×2 → (2, H, W)
            "height width, height width -> (1 + 1) height width",
            capability.x_neg,
            capability.x_pos,
        )
        runs_neg = einx.rearrange(                    # (H, W) ×2 → (2, H, W)
            "height width, height width -> (1 + 1) height width",
            capability.x_pos,
            capability.x_neg,
        )

        # same network both ways: no directional bias
        face_pos, kappa_pos = self.face_value(window_pos, courant, runs_pos, wet_pos)
        face_neg, kappa_neg = self.face_value(window_neg, courant, runs_neg, wet_neg)

        # q̂_{i+½} from the upwind side
        upwind = u >= 0.0                             # (H, W)
        face = jnp.where(upwind, face_pos, face_neg)  # (H, W)
        kappa = jnp.where(upwind, kappa_pos, kappa_neg)  # (H, W)

        # user
        # q_{i+1}
        q_east = shift_east(q)                        # (H, W) → (H, W)
        # Δq_{i+½} = q_{i+1} − q_i
        jump = q_east - q                             # (H, W)
        # F = u q̂ − |u| κ (q_{i+1} − q_i)
        advective = u * face                          # (H, W)
        diffusive = jnp.abs(u) * kappa * jump         # (H, W)
        flux = advective - diffusive                  # (H, W)
        # dry faces carry no flux
        flux = jnp.where(mask.u, flux, 0.0)           # (H, W)
        return flux


# today (finitevolx)
# the g_corr ≡ 0 oracle: classical masked WENO5 on the same grid and mask
baseline = fvx.Advection2D(grid, mask=mask)(pv, u, v, method="weno5")  # (H, W)

Markers: today ships on finitevolX main (checked 24 Oct 2026, 831c69f): Advection2D with method="weno5" and a Mask2D, and Mask2D.stencil_capability. user is defined on this page. finitevolX's WENO candidates and indicators are private functions today, so the sketch rebuilds them; a learned-weights hook in finitevolX would be proposed.

Where it goes in the QG model

The model on the other pages advects PV with the Arakawa Jacobian. Two routes:

Because the discrete velocity comes from a streamfunction, it has zero discrete divergence, and the flux form and the advective form agree. Route (b) changes the reconstruction, not the physics.

Training

• Target. A QG run at \(\Delta x = 2\) km, coarse-grained to 8 km by cell averaging. This is a correction target in the taxonomy's terms: the scheme learns what the coarse grid misses near fronts.
• Loss. A posteriori, through the solver: rollout error of \(q\) over several days, plus the error of the kinetic energy spectrum at high wavenumbers, so the scheme is judged on the scales it controls.
• Start. From classical WENO5 (zero last layer), so every gradient step starts from a working, stable scheme.
• Check of what it learned. Map \(\omega\) and \(\kappa_{\gvect{\theta}}\) over a snapshot. They are physical diagnostics: where the scheme trusts the upwind side, and where it dissipates.

For data assimilation

• A true adjoint. Classical limiters use min, max and switches, whose adjoints are discontinuous; 4D-Var systems often swap in a smoother linearised scheme for the tangent-linear model. Here every weight is a smooth function of \(q\), so jax.grad of the forecast is the exact adjoint of the forecast that is run.
• Information over the window. Numerical diffusion also damps perturbations in the tangent-linear model. Over a 10-day window, a too-diffusive scheme makes day-10 maps nearly uninformative about day 0. Lower, targeted dissipation keeps that information, so the advection scheme sets how long a window can usefully be.
• A dissipation map as a control. \(\kappa_{\gvect{\theta}}\) (or a spatial scale on it) can be estimated from observations as a hyperparameter, the A1 route of the assimilation page.

What to watch

Test ladder

Each test has a known answer. Run each with \(g_\text{corr} = 0\) first (it must match finitevolX's method="weno5" to round-off), then with the trained weights.

Where it lives