diagnostix: design details

The reference behind diagnostix: core types, the arrays callers pass in, the algorithms with their pitfalls, JAX notes, tests, dependencies, and a survey of the code diagnostix takes over, read on 5 Nov 2026.

Core types¶

Most diagnostics take arrays and a gridax geometry and return an array of the same grid. The types below are for the results that carry more than one array, and for the one input that needs structure, a control volume.

# Shapes: T times; Ny × Nx cells; K wavenumber bins; M modes; E members

import equinox as eqx
from jaxtyping import Array, Bool, Float, Int8


# proposed (diagnostix)
class ControlVolume(eqx.Module):
    """A region of a gridax grid and the signed faces that bound it."""

    inside: Bool[Array, "y x"]
    # +1 where the east face of cell (i, j) leaves the region, −1 enters, 0 not on it
    east_face_sign: Int8[Array, "y x"]
    # the same for the north face
    north_face_sign: Int8[Array, "y x"]


# proposed (diagnostix)
class Budget(eqx.Module):
    """d/dt ∫ q dV + ∮ q u·n dS = ∫ S dV, term by term, and what is left."""

    tendency: Float[Array, "time"]
    boundary_flux: Float[Array, "time"]
    source: Float[Array, "time"]
    residual: Float[Array, "time"]


# proposed (diagnostix)
class Spectrum(eqx.Module):
    """Power against wavenumber, isotropic or along one axis."""

    wavenumber: Float[Array, "k"]
    power: Float[Array, "*batch k"]
    kind: str = eqx.field(static=True)


# proposed (diagnostix)
class EOFs(eqx.Module):
    """Area-weighted EOFs of an anomaly field."""

    modes: Float[Array, "mode y x"]
    principal_components: Float[Array, "time mode"]
    explained_variance: Float[Array, "mode"]


# proposed (diagnostix)
class PSDScore(eqx.Module):
    """1 − PSD(error) / PSD(reference), and where it crosses a level."""

    wavenumber: Float[Array, "k"]
    score: Float[Array, "k"]
    resolved_wavelength: Float[Array, ""]


# proposed (diagnostix)
class SpreadSkill(eqx.Module):
    """Ensemble spread, RMSE of the mean, and their ratio."""

    spread: Float[Array, "*batch"]
    rmse: Float[Array, "*batch"]
    ratio: Float[Array, "*batch"]

Single-valued scores (RMSE, bias, correlation, CRPS) return plain arrays.

Types to adhere to¶

# Shapes: T times; Ny × Nx cells

import diagnostix as dgx
import gridax as gdx

# proposed (gridax)
grid = gdx.cells.LonLatGeometry.from_axes(lon, lat)
# proposed (gridax)
weights = gdx.metrics.cell_area(grid)  # → (Ny, Nx)
# proposed (diagnostix)
# a boolean region, made into signed boundary faces once, on the host
box = dgx.budgets.ControlVolume.from_region(grid, region)  # (Ny, Nx) → volume

NaN is not a mask: a diagnostic that sees NaN in a valid cell returns NaN there, so missing data must be marked in valid.

Algorithms¶

Geostrophy, vorticity, strain and Okubo–Weiss¶

From sea-surface height η (or geopotential) to geostrophic velocity, then the velocity-gradient invariants. On a lon–lat grid, λ is longitude, φ latitude, R the Earth’s radius, and u, v the eastward and northward components.

1. Take the derivatives with gridax: ζ is gridax’s spherical curl, and s_n, s_s are the covariant derivatives of the velocity, so they carry the tan φ connection terms, which scalar metric factors alone miss. On a Cartesian (projected) grid those terms vanish and the formulas reduce to ζ = ∂v/∂x − ∂u/∂y, s_n = ∂u/∂x − ∂v/∂y, s_s = ∂v/∂x + ∂u/∂y; on a curvilinear grid, use the grid’s own metric terms.

2. Centred differences on the A-grid inside the domain, one-sided at the edges, and NaN wherever the stencil touches an invalid cell.

3. Mask geostrophy where |φ| is below a latitude cut (default 5°), where f → 0.

4. Eddy cores are W < −W₀ with W₀ = 0.2 σ_W, the standard deviation of W over valid cells.

Cost: O(N) per field, one pass of stencils. Pitfalls: a stencil that reaches across a coast mixes land fill into the gradient (step 2 is the fix); the Cartesian formulas with only the scalar metric factors give solid-body rotation u = ωR cos φ a vorticity of ω sin φ, half the true 2ω sin φ, and a spurious shear strain, which then corrupts W; W₀ depends on the domain, so it is a parameter, not a constant.

Boundary-layer height by the bulk Richardson number¶

1. Virtual potential temperature on each level from T, p and specific humidity (geophysix).

2. Ri_b on every level above the surface reference.

3. The PBL height is the lowest z where Ri_b crosses the critical value Ri_c (default 0.25), linearly interpolated between the two levels that bracket the crossing.

4. No crossing in the column returns NaN, never the model top.

Cost: O(N_z) per column, vectorised over columns. JAX: find the first crossing with jnp.argmax on the boolean along z, then check it is a real crossing; no Python loop over levels. Pitfalls: a zero-shear column divides by zero without the u∗ term (b ≈ 100) or a floor; levels must be ordered upward, so the function checks that z is monotone.

Budgets with a true boundary flux¶

1. ControlVolume.from_region marks every cell face with one side inside the region and one outside, and gives it σ = +1 (flux leaves) or −1 (flux enters).

2. Interpolate q and the normal velocity to faces (the mean of the two neighbour cells on the A-grid; gridax’s face values on a C-grid).

3. Q(t) is the volume integral of q with gridax cell volumes; the tendency is a centred difference in time.

4. The residual is reported per time, with each term, so a budget is checked term by term.

Cost: O(N) per time. Pitfalls: xrops’ boundary_flux sums φ u A over every face of the field (its docstring says so), which is the net flux through a boundary only when the field covers exactly the region; diagnostix sums over the region’s own boundary. Snapshot fields for the tendency and time-mean fluxes for F is correct; the other way round is not. The face interpolation must follow the grid convention of the model that made the fields, or the residual is mostly interpolation error.

Climatology and anomalies¶

Here d is the calendar day of a sample’s (month, day) on a leap-year calendar, not its raw day of year: in a non-leap year 1 March is day 60 of the year but d = 61, the same slot as 1 March of a leap year, and d = 60 holds only 29 February.

1. Map each sample’s (month, day) to its calendar day d, on the host or with a 12-entry table of month offsets, then group with jax.ops.segment_sum and a static number of segments, counting only valid samples. The segment ID is zero-based, d − 1 ∈ [0, 366) (or month − 1 ∈ [0, 12)): segment_sum silently drops IDs outside [0, num_segments), so grouping by d itself would lose 31 December and leave segment 0 empty. Grouping by the raw day of year would put 1 March of non-leap years with 29 February, and shift every date after it by one between leap and non-leap years.

2. The smoothed climatology fits the first K harmonics (default 3) to the daily means by weighted least squares, the weights being the counts.

3. The anomaly subtracts the climatology at each sample’s day, looked up at the same zero-based ID d − 1; adding it back is the inverse.

Cost: O(N) to group, O(K²) per pixel to fit. Pitfalls: 29 February is its own segment with a quarter of the samples, so smoothing matters most there; a climatology fitted on the evaluation period leaks into the anomaly scores, so the base period is an argument.

Isotropic spectra and the PSD score¶

Here Δk_x = 1 / (N_x Δx) and Δk_y = 1 / (N_y Δy) are the spacings of the discrete wavenumbers, so P₂ is a 2-D density (Σ P₂ Δk_x Δk_y is the variance) and P is a 1-D density in |k| (Σ P(k) Δk is the same variance).

1. Remove the mean (or a plane), multiply by a 2-D Hann taper w, and divide by its energy U = mean(w²) so the taper does not remove variance.

2. jnp.fft.fft2, then |φ̂|² with the Δx Δy / (N_x N_y U) normalisation, giving the 2-D density P₂.

3. Bin |k| into K annuli of width Δk = min(Δk_x, Δk_y) with segment_sum (K static), and multiply each annulus sum by Δk_x Δk_y / Δk, the area of one wavenumber cell over the bin width. Without that factor P(k) is bin-integrated power, not a density, and its amplitude changes with the grid size and spacing.

4. The PSD score S(k) uses the error’s spectrum over the reference’s; the resolved wavelength is 1 / k∗ at the first k where S falls below the level (default 0.5), interpolated in log k.

Cost: O(N log N) per field. Pitfalls: the FFT needs a uniform Cartesian grid, so a lon–lat field is first put on a local tangent plane (gridax) or the domain kept small enough that cos φ barely changes; gaps cannot go through an FFT, so a field with invalid cells is filled (and the fraction reported) or refused; xrft and the numpy convention differ by the normalisation in step 2, so the oracle test fixes one; Parseval (Σ P(k) Δk = variance) is the check: on an 8 × 8 grid with unit spacing and Δk = 1/8, the factor is 1/8, and leaving it out integrates to eight times the taper-corrected mean square.

Area-weighted EOFs¶

1. X is the anomaly matrix, time × valid cells, with the climatology removed (previous section).

2. The weights w are gridax cell areas (∝ cos φ on a regular lon–lat grid), normalised to sum to the number of cells.

3. Thin SVD of X_w; keep M modes; modes are back-weighted so e_mᵀ diag(w) e_n = δ_mn.

4. Fix the sign of each mode so its largest-magnitude loading is positive.

Cost: O(T N min(T, N)) for the thin SVD; a randomized SVD (gaussx) for large fields. Pitfalls: xrops’ eof is sklearn PCA without weights, so polar cells count as much as tropical ones; the valid-cell selection must be static (made once on the host) to run under jit; mode signs are arbitrary until step 4, so comparisons across runs fail without it.

Weighted errors, CRPS and spread–skill¶

1. m is the validity mask (1 where valid) and w the cell areas; every score is a weighted mean over valid cells, then over time if asked.

2. CRPS uses the fair (E − 1) estimator by default; the pairwise term comes from the sorted members x₍ᵢ₎, O(E log E) instead of O(E²).

3. Spread–skill averages the ensemble variance before the square root, and includes the (E + 1)/E factor, so a reliable ensemble has SSR ≈ 1.

Cost: O(N) for pointwise scores, O(N E log E) for CRPS. Pitfalls: averaging RMSE per time step and then over time is not the RMSE of the whole period, so the reduction order is an argument; the biased CRPS penalises small ensembles; spread that is averaged after the square root is biased low.

JAX compatibility¶

Testing¶

• xrops as the oracle. Every diagnostic that xrops already has runs on the same arrays through both, inside the tolerance of float64: geostrophic velocities, vorticity, strain, Okubo–Weiss, KE, enstrophy, PBL height, climatology and anomaly, spectra (with the normalisation fixed), RMSE, CRPS, spread–skill.

• Closed forms. A solid-body rotation has ζ = 2ω and zero strain on a plane, and on the sphere ζ = 2ω sin φ with zero strain for rotation about the polar axis (u = ωR cos φ), the case that catches missing tan φ terms; rotation about any other axis gives ζ = 2ω · r̂ and zero strain too; a pure strain field has ζ = 0 and W > 0; a Gaussian eddy has W < 0 at its core.

• Conservation. A tracer advected by a divergence-free field in a closed box has a budget residual at round-off; a box cut through the middle has nonzero boundary flux equal to the tendency.

• Parseval. The binned spectrum integrates to the field’s variance (taper-corrected).

• EOFs. Weighted orthonormality of the modes; for a field made of two known patterns, the modes recover them up to sign; unweighted EOFs equal xrops’ eof.

• Scores. CRPS of a one-member ensemble equals MAE; the sorted and pairwise CRPS agree; SSR ≈ 1 within its sampling error for draws from the same distribution as the truth (bound from the number of samples, as the stack’s random-draw rule asks).

• Transforms. jit, vmap and grad on each row of the JAX table marked yes, and finite gradients on fields with masked cells.

External dependencies¶

pywt, xrft and sklearn are not dependencies: the routines that need them stay in xrops.

Where it stands¶

Read on 5 Nov 2026 at xrtoolz b3855d7, somax aea3a03 and finitevolX 82090a3 (0.1.0).