Skip to content

Assimilation

3D/4D-Var cost, control, and solve scaffolding (optimistix), imported on demand.

The package resolves its public names lazily (PEP 562), so they are documented here from the submodules that define them.

assimilation

Variational data assimilation (3D-Var) for methane VMR retrieval.

Modules

  • :mod:.obs_operator — JAX forward operator H = SRF ∘ GSD ∘ PSF ∘ exp(−a·ΔVMR) built from a HAPI cross-section LUT.
  • :mod:.background — covariance B constructors: diagonal, Kronecker (separable AR(1) spatial), low-rank update from a sample stack.
  • :mod:.control — control-variable transforms: identity vs whitening (Cholesky / CVT) — the latter dramatically improves Hessian conditioning.
  • :mod:.cost — cost+grad+HVP bundles in either model space (δx) or whitened space (ξ). Reverse-mode AD handles the obs-term adjoint.
  • :mod:.solve — three solver paths: primal LBFGS, Gauss-Newton with a custom :mod:lineax linear solver, and the dual / PSAS reformulation.
  • :mod:.diagnostics — reduced χ², degrees-of-freedom-for-signal, posterior covariance proxy via CG on the Hessian.

The math derivation lives in notebooks/assimilation/00_3dvar_derivation.md and the end-to-end demo in notebooks/assimilation/06_3dvar_methane_retrieval.ipynb.

Notes

assimilation follows the same lazy-import pattern as :mod:plumax.radtran (PEP 562): submodules are imported only on first attribute access, so import plumax does not pull in optimistix / gaussx until you actually touch the assimilation surface.

cost

3D-Var cost functions and gradients.

The standard incremental 3D-Var minimises

J(δx) = ½ δxᵀ B⁻¹ δx + ½ ‖y − H(x_b + δx)‖²_R⁻¹

over the increment δx = x − x_b. Two versions live here:

  • :func:build_cost_x — model-space form. The prior term is evaluated via :func:gaussx.solve(B, δx) so any structured B from :mod:.background works without materialisation.
  • :func:build_cost_xi — whitened (CVT) form. The prior term collapses to ½ ξᵀ ξ and the obs term sees δx = U ξ via the supplied :class:~plumax.assimilation.control.WhiteningTransform. This is the version we use in practice — gradients are well-scaled and the Hessian is identity-plus-low-rank.

In both cases the obs term is differentiated through the JAX forward map via reverse-mode AD (jax.grad), so we don't need to maintain an explicit adjoint of the observation operator. The :class:Cost bundle bundles J, ∇J and a Hessian-vector product Hv so the optimiser / posterior-covariance machinery downstream can pick whichever it needs.

Cost dataclass 

Cost(value: Callable[[Array], Array], grad: Callable[[Array], Array], value_and_grad: Callable[[Array], tuple[Array, Array]], hvp: Callable[[Array, Array], Array])

Bundle of (value, gradient, Hessian-vector product) for an optimiser.

The three callables share the same input signature — a 1-D state vector — so the optimiser doesn't have to know which control-space we're working in.

build_cost_x 

build_cost_x(*, forward_fn: Callable[[Array], Array], background_op: AbstractLinearOperator, obs_inv_variance: float | Array, background_state: Array, observation: Array, state_shape: tuple[int, ...]) -> Cost

Build J(δx) = ½ δxᵀ B⁻¹ δx + ½ (y − H(x_b + δx))ᵀ R⁻¹ (y − H(x_b + δx)).

Parameters

forward_fn : Callable H : 2-D VMR field → obs cube. Must be JAX-traceable. background_op : lineax operator B (any structured form from :mod:.background). The cost calls gaussx.solve(B, δx_flat) once per gradient evaluation. obs_inv_variance : float or Array R⁻¹. Pass a scalar for R = σ² I or an array broadcastable against the obs cube for diagonal heteroscedastic noise. background_state, observation, state_shape x_b, y_obs, and the shape of the VMR field (e.g. (ny, nx)).

Source code in src/plumax/assimilation/cost.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def build_cost_x(
    *,
    forward_fn: Callable[[jax.Array], jax.Array],
    background_op: lx.AbstractLinearOperator,
    obs_inv_variance: float | jax.Array,
    background_state: jax.Array,
    observation: jax.Array,
    state_shape: tuple[int, ...],
) -> Cost:
    """Build ``J(δx) = ½ δxᵀ B⁻¹ δx + ½ (y − H(x_b + δx))ᵀ R⁻¹ (y − H(x_b + δx))``.

    Parameters
    ----------
    forward_fn : Callable
        ``H : 2-D VMR field → obs cube``. Must be JAX-traceable.
    background_op : lineax operator
        ``B`` (any structured form from :mod:`.background`). The cost calls
        ``gaussx.solve(B, δx_flat)`` once per gradient evaluation.
    obs_inv_variance : float or Array
        ``R⁻¹``. Pass a scalar for ``R = σ² I`` or an array broadcastable
        against the obs cube for diagonal heteroscedastic noise.
    background_state, observation, state_shape
        ``x_b``, ``y_obs``, and the shape of the VMR field (e.g. ``(ny, nx)``).
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)

    def value(delta_x_flat: jax.Array) -> jax.Array:
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        x = x_b + delta_x
        residual = y - forward_fn(x)
        # Prior term: ½ δxᵀ B⁻¹ δx. gx.quadratic_form is xᵀ B⁻¹ x via a single
        # structurally-dispatched solve — the same op gx.solve does, one call.
        prior = 0.5 * gx.quadratic_form(background_op, delta_x_flat)
        obs = 0.5 * jnp.sum(R_inv * residual * residual)
        return prior + obs

    val_and_grad = jax.value_and_grad(value)
    grad = lambda x: val_and_grad(x)[1]

    def hvp(x: jax.Array, v: jax.Array) -> jax.Array:
        # ∂²J/∂x² · v via forward-over-reverse — the standard JAX recipe.
        return jax.jvp(grad, (x,), (v,))[1]

    return Cost(value=value, grad=grad, value_and_grad=val_and_grad, hvp=hvp)

build_cost_xi 

build_cost_xi(*, forward_fn: Callable[[Array], Array], whitening: WhiteningTransform | IdentityTransform, obs_inv_variance: float | Array, background_state: Array, observation: Array, state_shape: tuple[int, ...]) -> Cost

Build J(ξ) = ½ ξᵀ ξ + ½ (y − H(x_b + U ξ))ᵀ R⁻¹ (…).

The whitening transform turns the prior into a clean ½‖ξ‖², so the Hessian is I + (HU)ᵀ R⁻¹ (HU) — a small perturbation of the identity. Even vanilla LBFGS converges in a handful of iterations regardless of how badly conditioned B was.

Pass an :class:~plumax.assimilation.control.IdentityTransform here to recover an unpreconditioned model-space cost with implicit B = I — useful for unit tests of the obs term in isolation.

Source code in src/plumax/assimilation/cost.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def build_cost_xi(
    *,
    forward_fn: Callable[[jax.Array], jax.Array],
    whitening: WhiteningTransform | IdentityTransform,
    obs_inv_variance: float | jax.Array,
    background_state: jax.Array,
    observation: jax.Array,
    state_shape: tuple[int, ...],
) -> Cost:
    """Build ``J(ξ) = ½ ξᵀ ξ + ½ (y − H(x_b + U ξ))ᵀ R⁻¹ (…)``.

    The whitening transform turns the prior into a clean ``½‖ξ‖²``, so the
    Hessian is ``I + (HU)ᵀ R⁻¹ (HU)`` — a small perturbation of the identity.
    Even vanilla LBFGS converges in a handful of iterations regardless of how
    badly conditioned ``B`` was.

    Pass an :class:`~plumax.assimilation.control.IdentityTransform`
    here to recover an unpreconditioned model-space cost with implicit
    ``B = I`` — useful for unit tests of the obs term in isolation.
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)

    def value(xi_flat: jax.Array) -> jax.Array:
        delta_x_flat = whitening.apply(xi_flat)
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        x = x_b + delta_x
        residual = y - forward_fn(x)
        prior = 0.5 * jnp.dot(xi_flat, xi_flat)
        obs = 0.5 * jnp.sum(R_inv * residual * residual)
        return prior + obs

    val_and_grad = jax.value_and_grad(value)
    grad = lambda xi: val_and_grad(xi)[1]

    def hvp(xi: jax.Array, v: jax.Array) -> jax.Array:
        return jax.jvp(grad, (xi,), (v,))[1]

    return Cost(value=value, grad=grad, value_and_grad=val_and_grad, hvp=hvp)

finite_difference_grad 

finite_difference_grad(fn: Callable[[Array], Array], x: Array, *, eps: float = 1e-06) -> np.ndarray

Central finite differences — used only by tests to validate jax.grad.

O(N) function evaluations, so keep x.size small.

Source code in src/plumax/assimilation/cost.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def finite_difference_grad(
    fn: Callable[[jax.Array], jax.Array], x: jax.Array, *, eps: float = 1e-6
) -> np.ndarray:
    """Central finite differences — used only by tests to validate ``jax.grad``.

    O(N) function evaluations, so keep ``x.size`` small.
    """
    x_np = np.asarray(x, dtype=float).copy()
    g = np.zeros_like(x_np)
    for i in range(x_np.size):
        orig = x_np.flat[i]
        x_np.flat[i] = orig + eps
        f_plus = float(fn(jnp.asarray(x_np)))
        x_np.flat[i] = orig - eps
        f_minus = float(fn(jnp.asarray(x_np)))
        g.flat[i] = (f_plus - f_minus) / (2.0 * eps)
        x_np.flat[i] = orig
    return g

control

Control-variable transforms for variational data assimilation.

3D-Var minimises a cost J(δx) in some "control space". The two choices worth keeping straight are:

  1. Model space — control variable is the increment δx = x − x_b. The prior term is ½ δxᵀ B⁻¹ δx; one B⁻¹-solve per gradient evaluation. Hessian conditioning = cond(B⁻¹ + H'ᵀ R⁻¹ H') blows up when B is poorly scaled.

  2. Whitened (CVT) space — control variable is ξ with δx = U ξ and B = U Uᵀ. The prior collapses to ½ ξᵀ ξ and the Hessian becomes I + (HU)ᵀ R⁻¹ (HU) — a small, dense identity-plus-low-rank matrix that even simple LBFGS solves in a handful of iterations regardless of how ill-conditioned B was. This is the single biggest practical lever you have in 3D-Var.

Both transforms expose the same minimal interface:

  • apply(xi) → δx (control-space → model-space)
  • apply_inverse(δx) → xi (model-space → control-space; needed once at init if the user gives an x-space prior mean)
  • project_gradient(g_x) → g_xi (chain rule for the prior gradient when you're working with the cost-in-x form; not used in JAX autodiff paths)

We rely on :func:gaussx.cholesky to produce U as a structured operator, so the whitening cost stays O(B-solve) rather than O(N²) for any of the structured B flavours in :mod:.background.

IdentityTransform dataclass 

IdentityTransform()

Pass-through transform. Use when the prior is already well-conditioned or when you want the textbook ½ δxᵀ B⁻¹ δx form for didactic clarity.

WhiteningTransform dataclass 

WhiteningTransform(cholesky_op: AbstractLinearOperator)

δx = U ξ where B = U Uᵀ.

Constructed once via :meth:from_background; afterwards each call is a single structured matvec / triangular solve from gaussx.

Notes

Whitening also defines a natural sampler from the prior: δx = U ε with ε ∼ 𝒩(0, I) is exactly distributed as 𝒩(0, B). We don't use that here but it's the same operator.

from_background classmethod 

from_background(background: LinearOperator) -> WhiteningTransform

Factor B = U Uᵀ via :func:gaussx.cholesky.

gaussx returns a structured triangular operator (preserving Kronecker / low-rank-update structure), so subsequent matvecs and inverse-solves stay in the cheap regime.

Source code in src/plumax/assimilation/control.py 
86
87
88
89
90
91
92
93
94
95
@classmethod
def from_background(cls, background: LinearOperator) -> WhiteningTransform:
    """Factor ``B = U Uᵀ`` via :func:`gaussx.cholesky`.

    gaussx returns a structured triangular operator (preserving Kronecker /
    low-rank-update structure), so subsequent matvecs and inverse-solves
    stay in the cheap regime.
    """
    U = gx.cholesky(background)
    return cls(cholesky_op=U)

apply 

apply(xi: Array) -> jax.Array

ξ → U ξ.

Source code in src/plumax/assimilation/control.py 
97
98
99
def apply(self, xi: jax.Array) -> jax.Array:
    """``ξ → U ξ``."""
    return self.cholesky_op.mv(jnp.asarray(xi))

apply_inverse 

apply_inverse(delta_x: Array) -> jax.Array

δx → U⁻¹ δx via :func:gaussx.solve for structured dispatch.

gaussx.cholesky returns operators that preserve their underlying structure (a Cholesky of a Kronecker is itself a Kronecker of Cholesky factors); gaussx.solve knows how to dispatch on those, whereas lineax.AutoLinearSolver would try to detect generic triangularity and raise NotImplementedError.

Source code in src/plumax/assimilation/control.py 
101
102
103
104
105
106
107
108
109
110
def apply_inverse(self, delta_x: jax.Array) -> jax.Array:
    """``δx → U⁻¹ δx`` via :func:`gaussx.solve` for structured dispatch.

    ``gaussx.cholesky`` returns operators that preserve their underlying
    structure (a Cholesky of a ``Kronecker`` is itself a ``Kronecker`` of
    Cholesky factors); ``gaussx.solve`` knows how to dispatch on those,
    whereas ``lineax.AutoLinearSolver`` would try to detect generic
    triangularity and raise ``NotImplementedError``.
    """
    return gx.solve(self.cholesky_op, jnp.asarray(delta_x))

project_gradient 

project_gradient(g_x: Array) -> jax.Array

∂J/∂ξ = Uᵀ ∂J/∂δx — the chain rule for the explicit-cost form.

Not needed when you build the cost in ξ-space and let jax.grad differentiate end-to-end (which is what :mod:.cost does).

Source code in src/plumax/assimilation/control.py 
112
113
114
115
116
117
118
def project_gradient(self, g_x: jax.Array) -> jax.Array:
    """``∂J/∂ξ = Uᵀ ∂J/∂δx`` — the chain rule for the explicit-cost form.

    Not needed when you build the cost in ξ-space and let ``jax.grad``
    differentiate end-to-end (which is what :mod:`.cost` does).
    """
    return self.cholesky_op.transpose().mv(jnp.asarray(g_x))

background

Background-error covariance B constructors for 3D-Var.

The prior term ½ δxᵀ B⁻¹ δx in the 3D-Var cost function is dominated, at the operator level, by how you write down B. We support three flavours that cover most of what a methane-retrieval workflow ever needs:

============= ====================================================== ================= Flavour Operator structure Cost of B⁻¹ ============= ====================================================== ================= diagonal B = diag(σ²) O(N) Kronecker B = ρ_y ⊗ ρ_x (separable, e.g. AR(1) in each axis) O(N · √N) low-rank B = λI + U diag(d) Uᵀ (Woodbury via gaussx) O(N · k) ============= ====================================================== =================

All three return a :class:lineax.AbstractLinearOperator (via gaussx) so the preconditioning transform :class:~plumax.assimilation.control.WhiteningTransform and the cost gradient can both call :func:gaussx.solve and :func:gaussx.cholesky uniformly without caring about structure.

The Kronecker constructor uses an exponential (AR(1)-style) correlation [exp(-|i-j|/L)] along each axis — simple, positive-definite by construction, and matches the prior families used in operational variational assimilation systems for spatially-smooth tracer fields.

build_diagonal_background 

build_diagonal_background(variance: float | ndarray, n_pixels: int | None = None) -> LinearOperator

Build B = diag(σ²) as a :class:lineax.DiagonalLinearOperator.

Parameters

variance : float or array Per-pixel prior variance. Scalar broadcasts to n_pixels. n_pixels : int, optional Required when variance is a scalar.

Returns

LinearOperator B as a diagonal operator. gaussx.solve(B, r) is then a per-element division.

Source code in src/plumax/assimilation/background.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def build_diagonal_background(
    variance: float | np.ndarray, n_pixels: int | None = None
) -> LinearOperator:
    """Build ``B = diag(σ²)`` as a :class:`lineax.DiagonalLinearOperator`.

    Parameters
    ----------
    variance : float or array
        Per-pixel prior variance. Scalar broadcasts to ``n_pixels``.
    n_pixels : int, optional
        Required when ``variance`` is a scalar.

    Returns
    -------
    LinearOperator
        ``B`` as a diagonal operator. ``gaussx.solve(B, r)`` is then a
        per-element division.
    """
    var = np.asarray(variance, dtype=float).ravel()
    if var.size == 1:
        if n_pixels is None:
            raise ValueError(
                "build_diagonal_background: scalar variance requires `n_pixels`."
            )
        var = np.full(int(n_pixels), float(var.item()))
    if np.any(var <= 0.0):
        raise ValueError("build_diagonal_background: variances must be > 0.")
    return lx.DiagonalLinearOperator(jnp.asarray(var))

build_kronecker_background 

build_kronecker_background(*, ny: int, nx: int, variance: float, length_scale_y: float, length_scale_x: float) -> LinearOperator

Build B = σ² · (ρ_y ⊗ ρ_x) using :class:gaussx.Kronecker.

The Kronecker structure makes B⁻¹ an O(ny³ + nx³) operation rather than O((ny·nx)³), which is the difference between feasible and not for realistic scene sizes (e.g. 256×256 → 65 k state vector vs. 256³ + 256³ ≈ 30 M flops vs. 280 T flops).

The convention here is B = σ² · K_y ⊗ K_x with row-major flattening δx[i·nx + j] = δx(i, j) — matches vmr_field.reshape(-1) order.

Source code in src/plumax/assimilation/background.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def build_kronecker_background(
    *,
    ny: int,
    nx: int,
    variance: float,
    length_scale_y: float,
    length_scale_x: float,
) -> LinearOperator:
    """Build ``B = σ² · (ρ_y ⊗ ρ_x)`` using :class:`gaussx.Kronecker`.

    The Kronecker structure makes ``B⁻¹`` an ``O(ny³ + nx³)`` operation rather
    than ``O((ny·nx)³)``, which is the difference between feasible and not for
    realistic scene sizes (e.g. 256×256 → 65 k state vector vs.
    256³ + 256³ ≈ 30 M flops vs. 280 T flops).

    The convention here is ``B = σ² · K_y ⊗ K_x`` with row-major flattening
    ``δx[i·nx + j] = δx(i, j)`` — matches ``vmr_field.reshape(-1)`` order.
    """
    if variance <= 0.0:
        raise ValueError("build_kronecker_background: `variance` must be > 0.")
    Ky = _ar1_correlation_matrix(ny, length_scale_y)
    Kx = _ar1_correlation_matrix(nx, length_scale_x)
    # Spread the σ² factor evenly between the two Kronecker factors so each is
    # individually well-scaled. Mathematically (σ a) ⊗ (σ b) = σ² (a ⊗ b).
    sigma = float(np.sqrt(variance))
    op_y = lx.MatrixLinearOperator(
        jnp.asarray(sigma * Ky),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    op_x = lx.MatrixLinearOperator(
        jnp.asarray(sigma * Kx),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    return gx.Kronecker(op_y, op_x)

build_lowrank_background 

build_lowrank_background(*, samples: ndarray, rank: int, regularization: float = 1e-06, random_state: int = 0) -> LinearOperator

Estimate B from a stack of state samples and wrap as λI + U D Uᵀ.

Parameters

samples : np.ndarray Sample matrix, shape (n_samples, n_state) — e.g. flattened VMR snapshots from a NWP ensemble or a climatology. rank : int Number of leading SVD components to keep. Costs O(n_state · k) per solve. regularization : float Diagonal floor λ so the resulting operator is strictly PD even when samples is rank-deficient. Must be > 0. random_state : int Seed for :func:gaussx.randomized_svd, used when rank < 0.5 * min(n_samples, n_state) (otherwise a full SVD runs).

Returns

LinearOperator gaussx.LowRankUpdate tagged symmetric + PSD; gaussx.solve(B, r) uses Woodbury internally.

Source code in src/plumax/assimilation/background.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def build_lowrank_background(
    *,
    samples: np.ndarray,
    rank: int,
    regularization: float = 1e-6,
    random_state: int = 0,
) -> LinearOperator:
    """Estimate ``B`` from a stack of state samples and wrap as ``λI + U D Uᵀ``.

    Parameters
    ----------
    samples : np.ndarray
        Sample matrix, shape ``(n_samples, n_state)`` — e.g. flattened VMR
        snapshots from a NWP ensemble or a climatology.
    rank : int
        Number of leading SVD components to keep. Costs ``O(n_state · k)``
        per solve.
    regularization : float
        Diagonal floor ``λ`` so the resulting operator is strictly PD even
        when ``samples`` is rank-deficient. Must be > 0.
    random_state : int
        Seed for :func:`gaussx.randomized_svd`, used when
        ``rank < 0.5 * min(n_samples, n_state)`` (otherwise a full SVD runs).

    Returns
    -------
    LinearOperator
        ``gaussx.LowRankUpdate`` tagged symmetric + PSD; ``gaussx.solve(B, r)``
        uses Woodbury internally.
    """
    if samples.ndim != 2:
        raise ValueError(
            f"build_lowrank_background: samples must be 2-D (n_samples, n_state); "
            f"got {samples.shape}."
        )
    n_samples, n_state = samples.shape
    if n_samples < 2:
        raise ValueError("build_lowrank_background: need ≥ 2 samples.")
    rank = max(1, min(int(rank), n_state, n_samples))
    if regularization <= 0.0:
        raise ValueError("build_lowrank_background: `regularization` must be > 0.")
    centred = samples - samples.mean(axis=0, keepdims=True)
    if rank < 0.5 * min(n_samples, n_state):
        # Few components wanted: randomized range finder, O(rank) passes over
        # the samples instead of a full O(N·n·min(N,n)) SVD.
        _, S_k, V_k = gx.randomized_svd(
            lx.MatrixLinearOperator(jnp.asarray(centred)),
            rank,
            n_power_iter=5,
            key=jax.random.key(random_state),
        )
    else:
        _, S, Vt = np.linalg.svd(centred, full_matrices=False)
        S_k = S[:rank]
        V_k = Vt[:rank]
    U = jnp.asarray(V_k.T)  # (n_state, rank)
    d = jnp.asarray(S_k**2 / max(n_samples - 1, 1))
    base = lx.DiagonalLinearOperator(
        jnp.asarray(regularization * np.ones(n_state, dtype=float))
    )
    return gx.LowRankUpdate(
        base,
        U,
        d,
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )

obs_operator

JAX observation operator for the methane VMR retrieval.

The full forward map factors as

x  ──┐
     │ ΔVMR = x − x_ref
     ▼
    Δτ(ν, i, j) = a(ν) · ΔVMR(i, j)         a(ν) = σ(ν, T, P) · N · ΔX · AMF
     │
     ▼
    L_hr(i, j, ν) = exp(−Δτ)                 (or 1 − Δτ for linear)
     │
     ▼ PSF (per channel)
     ▼ GSD (per channel)
     ▼ SRF (band integration)
    y_obs(I, J, b)

Everything below the LUT lookup is JAX, so

  • forward runs end-to-end under :func:jax.jit,
  • jax.grad(loss∘forward) produces the discrete adjoint for free,
  • jax.jvp(forward, …) produces the tangent-linear,

without us writing or maintaining any explicit adjoint. The reason radtran/instrument.py does expose hand-rolled adjoint methods is to let unit tests verify the inner-product identity independently — JAX's autodiff inherits the correctness, not the other way around.

The factory :meth:RadianceObservationModel.from_lut precomputes a(ν) once from the HAPI cross-section LUT so the per-iteration cost in 3D-Var is dominated by the cube-shaped operations (PSF, GSD, SRF).

RadianceObservationModel dataclass 

RadianceObservationModel(absorption_coeff_hr: Array, srf_matrix: Array, vmr_reference: float = 0.0, psf: PointSpreadFunction | None = None, gsd: GroundSamplingDistance | None = None)

Bundle of precomputed pieces that defines the forward operator H.

Attributes

absorption_coeff_hr : jnp.ndarray Per-wavenumber absorption coefficient a(ν) = σ(ν,T,P) · N · ΔX · AMF, shape (n_lambda,). Captures everything except the VMR field, so the forward map is just L_hr(i,j,ν) = exp(-a(ν) · ΔVMR(i,j)). srf_matrix : jnp.ndarray L1-normalised SRF matrix, shape (n_bands, n_lambda). vmr_reference : float Background VMR x_ref around which radiances are normalised. Setting vmr_reference = 0 recovers absolute Beer-Lambert; setting it to the scene's true background produces the differential signal that matched filter retrievals consume. psf, gsd : optional JAX-side instrument operators from :mod:plumax.radtran.instrument. None skips that stage — convenient for hyperspectral tests where the spectral cube is already at sensor resolution.

from_lut classmethod 

from_lut(lut: Dataset, *, srf: SpectralResponseFunction, T_K: float, p_atm: float, path_length_cm: float, amf: float, vmr_reference: float = 0.0, psf: PointSpreadFunction | None = None, gsd: GroundSamplingDistance | None = None, var: str = 'absorption_cross_section') -> RadianceObservationModel

Build the operator from a HAPI σ(ν,T,P) LUT and an SRF.

Pre-collapses σ · N · ΔX · AMF into a single a(ν) vector so the forward call only multiplies by ΔVMR per pixel.

The HR spectral grid is taken from the SRF (srf.wavelengths_hr_nm) and converted to wavenumbers for the σ lookup. This is the grid the SRF matrix's columns are indexed by, so the subsequent einsum("bl, ijl -> ijb") band integration contracts aligned spectral bins by construction. Earlier revisions accepted a separate nu_obs argument; when callers passed a wavenumber-ordered grid while the SRF was wavelength-ordered (common whenever wl = 1e7/ν was sorted), the einsum silently mixed non-corresponding bins — self-consistent for twin retrievals but physically wrong for anything cross-validated.

Source code in src/plumax/assimilation/obs_operator.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
@classmethod
def from_lut(
    cls,
    lut: xr.Dataset,
    *,
    srf: SpectralResponseFunction,
    T_K: float,
    p_atm: float,
    path_length_cm: float,
    amf: float,
    vmr_reference: float = 0.0,
    psf: PointSpreadFunction | None = None,
    gsd: GroundSamplingDistance | None = None,
    var: str = "absorption_cross_section",
) -> RadianceObservationModel:
    """Build the operator from a HAPI σ(ν,T,P) LUT and an SRF.

    Pre-collapses ``σ · N · ΔX · AMF`` into a single ``a(ν)`` vector so the
    forward call only multiplies by ``ΔVMR`` per pixel.

    The HR spectral grid is **taken from the SRF** (``srf.wavelengths_hr_nm``)
    and converted to wavenumbers for the σ lookup. This is the grid the SRF
    matrix's columns are indexed by, so the subsequent
    ``einsum("bl, ijl -> ijb")`` band integration contracts aligned spectral
    bins by construction. Earlier revisions accepted a separate ``nu_obs``
    argument; when callers passed a wavenumber-ordered grid while the SRF
    was wavelength-ordered (common whenever ``wl = 1e7/ν`` was sorted),
    the einsum silently mixed non-corresponding bins — self-consistent for
    twin retrievals but physically wrong for anything cross-validated.
    """
    if var not in lut:
        raise KeyError(
            f"RadianceObservationModel.from_lut: variable {var!r} not in "
            f"dataset (have {list(lut.data_vars)})."
        )
    import xarray as xr_mod

    # Source-of-truth HR grid: the SRF's own wavelength axis (ascending
    # by wavelength = descending by wavenumber). Convert to cm⁻¹ for σ
    # interpolation; keep the wavelength ordering for the output array so
    # axis l aligns with srf.matrix[:, l].
    wl_hr_nm = np.asarray(srf.wavelengths_hr_nm, dtype=float)
    nu_hr = 1e7 / wl_hr_nm
    nu_da = xr_mod.DataArray(nu_hr, dims=["obs_nu"])
    sigma = lut[var].interp(
        wavenumber=nu_da, temperature=T_K, pressure=p_atm, method="linear"
    )
    sigma_np = np.asarray(sigma.values, dtype=float)
    N_total = number_density_cm3(p_atm, T_K)
    a = sigma_np * N_total * path_length_cm * amf
    return cls(
        absorption_coeff_hr=jnp.asarray(a),
        srf_matrix=jnp.asarray(srf.matrix),
        vmr_reference=float(vmr_reference),
        psf=psf,
        gsd=gsd,
    )

forward 

forward(vmr_field: Array, *, linear: bool = False) -> jax.Array

Map a 2-D VMR field to a band-integrated obs cube.

Parameters

vmr_field : Array Per-pixel VMR, shape (ny_hr, nx_hr). linear : bool If True, replace exp(-Δτ) by its Maclaurin-1 approximation 1 - Δτ. The tangent-linear of the nonlinear forward at vmr_field == vmr_reference coincides with this; the matched filter is exactly this case with δx propagated through.

Returns

Array Sensor-space radiance cube, shape (ny_lr, nx_lr, n_bands) (ny_lr = ny_hr // f, nx_lr = nx_hr // f if a GSD is set).

Source code in src/plumax/assimilation/obs_operator.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def forward(self, vmr_field: jax.Array, *, linear: bool = False) -> jax.Array:
    """Map a 2-D VMR field to a band-integrated obs cube.

    Parameters
    ----------
    vmr_field : Array
        Per-pixel VMR, shape ``(ny_hr, nx_hr)``.
    linear : bool
        If True, replace ``exp(-Δτ)`` by its Maclaurin-1 approximation
        ``1 - Δτ``. The tangent-linear of the nonlinear forward at
        ``vmr_field == vmr_reference`` coincides with this; the matched
        filter is exactly this case with ``δx`` propagated through.

    Returns
    -------
    Array
        Sensor-space radiance cube, shape ``(ny_lr, nx_lr, n_bands)``
        (``ny_lr = ny_hr // f``, ``nx_lr = nx_hr // f`` if a GSD is set).
    """
    if vmr_field.ndim != 2:
        raise ValueError(
            f"RadianceObservationModel.forward: vmr_field must be 2-D, got shape {vmr_field.shape}."
        )
    delta_vmr = vmr_field - self.vmr_reference
    # Outer product (ny, nx) ⊗ (n_lambda,) → (ny, nx, n_lambda).
    delta_tau = delta_vmr[..., None] * self.absorption_coeff_hr
    # `linear` is a Python bool (static); branch on it directly so JAX
    # doesn't trace the unused branch — `exp(huge δτ)` would otherwise
    # overflow on stress-tests like the adjoint inner-product check.
    L_hr = (1.0 - delta_tau) if linear else jnp.exp(-delta_tau)
    if self.psf is not None:
        L_hr = self.psf.apply(L_hr)
    if self.gsd is not None:
        L_hr = self.gsd.apply(L_hr)
    # Band integration: contract last (HR-λ) axis with SRF.
    return jnp.einsum("bl, ijl -> ijb", self.srf_matrix, L_hr)

make_forward 

make_forward(*, linear: bool = False) -> Callable[[jax.Array], jax.Array]

Return a closure that captures the model so it can be jit-compiled.

jax.jit(model.forward) works in JAX ≥ 0.4 because frozen dataclasses without register_pytree_node_class are hashable and become static arguments — but explicit closures are easier to compose with jax.value_and_grad and friends, and they work uniformly across JAX versions.

Source code in src/plumax/assimilation/obs_operator.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def make_forward(self, *, linear: bool = False) -> Callable[[jax.Array], jax.Array]:
    """Return a closure that captures the model so it can be ``jit``-compiled.

    ``jax.jit(model.forward)`` works in JAX ≥ 0.4 because frozen dataclasses
    without ``register_pytree_node_class`` are hashable and become static
    arguments — but explicit closures are easier to compose with
    ``jax.value_and_grad`` and friends, and they work uniformly across
    JAX versions.
    """
    linear_flag = bool(linear)

    def _fwd(vmr_field: jax.Array) -> jax.Array:
        return self.forward(vmr_field, linear=linear_flag)

    return _fwd

solve

Optimisers for 3D-Var: primal LBFGS, Gauss-Newton, and a dual PSAS path.

Three entry points, each tuned to a different conditioning story:

  • :func:run_lbfgs — black-box minimiser. Pass any cost+grad bundle from :mod:.cost. Robust, but iteration count tracks cond(Hessian), so use this with a :class:~plumax.assimilation.control.WhiteningTransform.
  • :func:run_gauss_newton — assumes the cost is a sum of squared residuals (which 3D-Var is, by construction). Each outer step solves a linear system (JᵀJ) δ = -Jᵀr via a user-chosen :mod:lineax solver (CG by default, so structured B keeps Woodbury / Kronecker on the hot path).
  • :func:run_dual_psas — the observation-space reformulation. When dim(y) ≪ dim(x) (typical for compact plumes), solving (H'BH'ᵀ + R) λ = d and recovering δx = B H'ᵀ λ is dramatically cheaper than the primal — and exactly equivalent for the linear forward.

The dual path is also the "matched filter, generalised": the matched-filter target spectrum t falls out as H' applied to a unit increment, and the PSAS solution δx̂ = B H'ᵀ (H'BH'ᵀ + R)⁻¹ d is the BLUE estimator that the pixel-wise matched filter approximates with B = ∞ I, R = Σ_pixel.

SolveResult dataclass 

SolveResult(state: ndarray, n_steps: int, cost_history: ndarray)

Common result type for all three solvers.

Attributes

state : np.ndarray Optimal control variable (ξ for whitened, δx for model-space, δx for the dual after the back-transform). n_steps : int Number of iterations consumed (-1 if the solver doesn't report). cost_history : np.ndarray Cost trajectory; for solvers that don't expose it, this is empty.

run_lbfgs 

run_lbfgs(cost: Cost, initial_state: Array, *, rtol: float = 1e-06, atol: float = 1e-09, max_steps: int = 200) -> SolveResult

Minimise cost.value via :class:optimistix.LBFGS.

The cost is wrapped to match optimistix's (y, args) → scalar signature. The implicit-function adjoint is the default — fine for this use case because we don't differentiate through the solve.

Source code in src/plumax/assimilation/solve.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def run_lbfgs(
    cost: Cost,
    initial_state: jax.Array,
    *,
    rtol: float = 1e-6,
    atol: float = 1e-9,
    max_steps: int = 200,
) -> SolveResult:
    """Minimise ``cost.value`` via :class:`optimistix.LBFGS`.

    The cost is wrapped to match optimistix's ``(y, args) → scalar`` signature.
    The implicit-function adjoint is the default — fine for this use case
    because we don't differentiate *through* the solve.
    """
    fn = lambda y, _args: cost.value(y)
    solver = optx.LBFGS(rtol=rtol, atol=atol)
    sol = optx.minimise(fn, solver, initial_state, max_steps=max_steps, throw=False)
    return SolveResult(
        state=np.asarray(sol.value),
        n_steps=int(sol.stats.get("num_steps", -1)),
        cost_history=np.asarray([]),
    )

run_gauss_newton 

run_gauss_newton(*, residual_fn: Callable[[Array], Array], initial_state: Array, linear_solver: AbstractLinearSolver | None = None, rtol: float = 1e-06, atol: float = 1e-09, max_steps: int = 50) -> SolveResult

Run :class:optimistix.GaussNewton for min ½ ‖r(y)‖².

residual_fn should return a flat residual vector [B^{-½}(δx); R^{-½}(y - H(x_b + δx))] (the "augmented" 3D-Var residual); optimistix internally constructs the Jacobian and solves the normal equations (JᵀJ) δ = -Jᵀr using linear_solver. Pass lineax.CG(...) for a matrix-free inner solve (useful when the Jacobian has Kronecker / low-rank structure that full Cholesky would materialise), or lineax.Cholesky() for small dense systems.

Source code in src/plumax/assimilation/solve.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def run_gauss_newton(
    *,
    residual_fn: Callable[[jax.Array], jax.Array],
    initial_state: jax.Array,
    linear_solver: lx.AbstractLinearSolver | None = None,
    rtol: float = 1e-6,
    atol: float = 1e-9,
    max_steps: int = 50,
) -> SolveResult:
    """Run :class:`optimistix.GaussNewton` for ``min ½ ‖r(y)‖²``.

    ``residual_fn`` should return a flat residual vector
    ``[B^{-½}(δx); R^{-½}(y - H(x_b + δx))]`` (the "augmented" 3D-Var residual);
    optimistix internally constructs the Jacobian and solves the normal
    equations ``(JᵀJ) δ = -Jᵀr`` using ``linear_solver``. Pass
    ``lineax.CG(...)`` for a matrix-free inner solve (useful when the
    Jacobian has Kronecker / low-rank structure that full Cholesky would
    materialise), or ``lineax.Cholesky()`` for small dense systems.
    """
    if linear_solver is None:
        linear_solver = lx.AutoLinearSolver(well_posed=None)
    fn = lambda y, _args: residual_fn(y)
    solver = optx.GaussNewton(rtol=rtol, atol=atol, linear_solver=linear_solver)
    sol = optx.least_squares(
        fn, solver, initial_state, max_steps=max_steps, throw=False
    )
    return SolveResult(
        state=np.asarray(sol.value),
        n_steps=int(sol.stats.get("num_steps", -1)),
        cost_history=np.asarray([]),
    )

run_dual_psas 

run_dual_psas(*, forward_fn: Callable[[Array], Array], background_op: AbstractLinearOperator, obs_inv_variance: float, background_state: Array, observation: Array, state_shape: tuple[int, ...], cg_rtol: float = 1e-06, cg_atol: float = 1e-09, cg_max_steps: int = 200) -> SolveResult

Solve the dual 3D-Var system in observation space.

For a forward operator that is linear in the state perturbation the optimum of J(δx) = ½ δxᵀ B⁻¹ δx + ½(y - H(x_b + δx))ᵀ R⁻¹(y - H(x_b + δx)) satisfies

(H' B H'ᵀ + R) λ = d,           with d = y − H(x_b)
δx̂ = B H'ᵀ λ.

When dim(y) ≪ dim(x) this is much cheaper than primal — and it's exactly the estimator that the pixel-wise matched filter approximates.

Parameters

forward_fn : Callable Full-state forward x → y (i.e. the same object returned by :meth:RadianceObservationModel.make_forward). For the dual to be correct, H must be linear in δx = x − x_b. We extract the tangent-linear action by differencing at x_b once (see base_obs below) and the adjoint via :func:jax.vjp. Callers typically pass model.make_forward(linear=True). background_op : lineax operator B. obs_inv_variance : float Scalar R⁻¹ (heteroscedastic R can be folded into forward_fn at the cost of clarity — keep it scalar for the demo). background_state, observation, state_shape Self-explanatory. cg_rtol, cg_atol, cg_max_steps Tolerances for the CG solve of the dual system.

Source code in src/plumax/assimilation/solve.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def run_dual_psas(
    *,
    forward_fn: Callable[[jax.Array], jax.Array],
    background_op: lx.AbstractLinearOperator,
    obs_inv_variance: float,
    background_state: jax.Array,
    observation: jax.Array,
    state_shape: tuple[int, ...],
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-9,
    cg_max_steps: int = 200,
) -> SolveResult:
    """Solve the **dual** 3D-Var system in observation space.

    For a forward operator that is linear in the state perturbation the
    optimum of ``J(δx) = ½ δxᵀ B⁻¹ δx + ½(y - H(x_b + δx))ᵀ R⁻¹(y - H(x_b + δx))``
    satisfies

        (H' B H'ᵀ + R) λ = d,           with d = y − H(x_b)
        δx̂ = B H'ᵀ λ.

    When ``dim(y) ≪ dim(x)`` this is much cheaper than primal — and it's
    exactly the estimator that the pixel-wise matched filter approximates.

    Parameters
    ----------
    forward_fn : Callable
        **Full-state forward** ``x → y`` (i.e. the same object returned by
        :meth:`RadianceObservationModel.make_forward`). For the dual to be
        correct, ``H`` must be linear in ``δx = x − x_b``. We extract the
        tangent-linear action by differencing at ``x_b`` once (see
        ``base_obs`` below) and the adjoint via :func:`jax.vjp`. Callers
        typically pass ``model.make_forward(linear=True)``.
    background_op : lineax operator
        ``B``.
    obs_inv_variance : float
        Scalar ``R⁻¹`` (heteroscedastic R can be folded into ``forward_fn``
        at the cost of clarity — keep it scalar for the demo).
    background_state, observation, state_shape
        Self-explanatory.
    cg_rtol, cg_atol, cg_max_steps
        Tolerances for the CG solve of the dual system.
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = float(obs_inv_variance)

    # Evaluate H(x_b) once — used both for the innovation and as the affine
    # anchor in the tangent-linear probe. Hoisting this out of the matvec
    # avoids a redundant forward pass per CG iteration.
    base_obs = forward_fn(x_b)
    d = (y - base_obs).reshape(-1)
    n_obs = d.size
    n_state = int(np.prod(state_shape))

    def H_op(delta_x_flat: jax.Array) -> jax.Array:
        """``H' · δx``: linearity ⇒ ``H(x_b + δx) − H(x_b)`` is exact."""
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        return (forward_fn(x_b + delta_x) - base_obs).reshape(-1)

    # Build the VJP closure *once*, outside the CG matvec, so every iteration
    # reuses the traced pullback instead of re-tracing jax.vjp(H_op, 0) each
    # time.  This is materially faster — was the P2 perf finding in review.
    _, _Ht_pullback = jax.vjp(H_op, jnp.zeros(n_state))

    def Ht_op(lam: jax.Array) -> jax.Array:
        (out,) = _Ht_pullback(lam)
        return out

    def dual_matvec(lam: jax.Array) -> jax.Array:
        """``(H' B H'ᵀ + R) λ``, with ``R λ = λ / R_inv`` for scalar R."""
        BHt_lam = background_op.mv(Ht_op(lam))
        return H_op(BHt_lam) + lam / R_inv

    dual_op = lx.FunctionLinearOperator(
        dual_matvec,
        input_structure=jax.eval_shape(lambda: jnp.zeros(n_obs)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    sol = lx.linear_solve(
        dual_op,
        d,
        solver=lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps),
        throw=False,  # surface non-convergence in `sol.result` instead of raising
    )
    lam_star = sol.value
    delta_x_flat = background_op.mv(Ht_op(lam_star))
    return SolveResult(
        state=np.asarray(delta_x_flat),
        n_steps=int(sol.stats.get("num_steps", -1)) if hasattr(sol, "stats") else -1,
        cost_history=np.asarray([]),
    )

diagnostics

Posterior diagnostics for a converged 3D-Var solution.

After the optimiser has produced x̂, three quick scalars tell you whether the result is statistically consistent with the prior + likelihood you specified:

  • :func:reduced_chi_squared — χ² = ‖y − H(x̂)‖²_R⁻¹ / dim(y). ≈ 1 means the posterior fits the observations to the noise level. Much smaller → over-fitting (R too generous); much larger → systematic errors not in R.
  • :func:degrees_of_freedom_for_signal — DFS = trace(I − A) with the averaging kernel A = Bₐ B⁻¹ (Rodgers, 2000). For low-rank B this is cheap; for the dense form we use Hutchinson's stochastic trace estimator.
  • :func:posterior_covariance_proxy — the inverse Hessian at the optimum. Returned as a callable that takes a probe vector v and returns Bₐ v (i.e. you can evaluate variances per-pixel without forming the full Bₐ matrix).

These are diagnostics, not exact statistics — for the linear-Gaussian case they coincide with the textbook posterior; for the nonlinear case they're Laplace approximations around the optimum.

reduced_chi_squared 

reduced_chi_squared(*, forward_fn: Callable[[Array], Array], estimated_state: Array, observation: Array, obs_inv_variance: float | Array) -> float

χ²_red = (y - H(x̂))ᵀ R⁻¹ (y - H(x̂)) / dim(y).

Values near 1 indicate the posterior is consistent with the observation error model; ≪ 1 → noise overestimated; ≫ 1 → noise underestimated or systematic forward-model error.

Source code in src/plumax/assimilation/diagnostics.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def reduced_chi_squared(
    *,
    forward_fn: Callable[[jax.Array], jax.Array],
    estimated_state: jax.Array,
    observation: jax.Array,
    obs_inv_variance: float | jax.Array,
) -> float:
    """``χ²_red = (y - H(x̂))ᵀ R⁻¹ (y - H(x̂)) / dim(y)``.

    Values near 1 indicate the posterior is consistent with the observation
    error model; ≪ 1 → noise overestimated; ≫ 1 → noise underestimated or
    systematic forward-model error.
    """
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)
    residual = y - forward_fn(estimated_state)
    chi2 = float(jnp.sum(R_inv * residual * residual))
    return chi2 / int(y.size)

degrees_of_freedom_for_signal 

degrees_of_freedom_for_signal(*, hessian_vector_product: Callable[[Array], Array], state_size: int, background_op: AbstractLinearOperator | None = None, whitened: bool = False, n_probes: int = 32, seed: int = 0, cg_rtol: float = 1e-06, cg_atol: float = 1e-10, cg_max_steps: int = 2000) -> float

Estimate DFS = trace(I − Bₐ B⁻¹) (Rodgers 2000, §2.5).

For the linear-Gaussian case the averaging kernel is A = I − Bₐ B⁻¹ and DFS = trace(A) measures how many independent pieces of information the observations contributed (DFS = 0 → no information; DFS = N → observations fully pin the state). The posterior covariance is the inverse Hessian at the optimum, Bₐ ≈ Hess⁻¹. Because DFS is invariant under the control-variable change δx = U ξ, the estimate must be computed in the space the supplied hessian_vector_product lives in — mixing an ξ-space Hessian with a model-space B gives a wrong answer for any B ≠ I.

Two modes, one per cost builder:

  • Model space (whitened=False, the default) — pass the Hessian from :func:plumax.assimilation.cost.build_cost_x and background_op = B. Each Hutchinson probe needs one B⁻¹ z via :func:gaussx.solve (structured dispatch) and one Hess⁻¹ w via CG: DFS ≈ E[zᵀz − zᵀ Hess⁻¹ B⁻¹ z].
  • Whitened space (whitened=True) — pass the Hessian from :func:plumax.assimilation.cost.build_cost_xi (Hess_ξ = I + (HU)ᵀR⁻¹HU). The whitened prior is the identity, so B⁻¹ drops out: DFS ≈ E[zᵀz − zᵀ Hess_ξ⁻¹ z] and background_op must be omitted.

Cost per probe: O(n_CG · forward) — dominated by the CG solve.

Earlier versions computed trace(B · Hess), which is not DFS: in the zero-information limit (Hess = B⁻¹) it returns N instead of the correct 0. See PR-review thread on diagnostics.py.

Source code in src/plumax/assimilation/diagnostics.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def degrees_of_freedom_for_signal(
    *,
    hessian_vector_product: Callable[[jax.Array], jax.Array],
    state_size: int,
    background_op: lx.AbstractLinearOperator | None = None,
    whitened: bool = False,
    n_probes: int = 32,
    seed: int = 0,
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-10,
    cg_max_steps: int = 2000,
) -> float:
    """Estimate ``DFS = trace(I − Bₐ B⁻¹)`` (Rodgers 2000, §2.5).

    For the linear-Gaussian case the averaging kernel is ``A = I − Bₐ B⁻¹``
    and ``DFS = trace(A)`` measures how many independent pieces of information
    the observations contributed (``DFS = 0`` → no information; ``DFS = N`` →
    observations fully pin the state). The posterior covariance is the
    inverse Hessian at the optimum, ``Bₐ ≈ Hess⁻¹``. Because DFS is invariant
    under the control-variable change ``δx = U ξ``, the estimate must be
    computed in the space the supplied ``hessian_vector_product`` lives in —
    mixing an ξ-space Hessian with a model-space ``B`` gives a wrong answer for
    any ``B ≠ I``.

    Two modes, one per cost builder:

    - **Model space** (``whitened=False``, the default) — pass the Hessian from
      :func:`plumax.assimilation.cost.build_cost_x` and ``background_op = B``.
      Each Hutchinson probe needs one ``B⁻¹ z`` via :func:`gaussx.solve`
      (structured dispatch) and one ``Hess⁻¹ w`` via CG:
      ``DFS ≈ E[zᵀz − zᵀ Hess⁻¹ B⁻¹ z]``.
    - **Whitened space** (``whitened=True``) — pass the Hessian from
      :func:`plumax.assimilation.cost.build_cost_xi`
      (``Hess_ξ = I + (HU)ᵀR⁻¹HU``). The whitened prior is the identity, so
      ``B⁻¹`` drops out: ``DFS ≈ E[zᵀz − zᵀ Hess_ξ⁻¹ z]`` and ``background_op``
      must be omitted.

    Cost per probe: ``O(n_CG · forward)`` — dominated by the CG solve.

    Earlier versions computed ``trace(B · Hess)``, which is not DFS: in the
    zero-information limit (``Hess = B⁻¹``) it returns ``N`` instead of the
    correct ``0``. See PR-review thread on diagnostics.py.
    """
    if whitened:
        if background_op is not None:
            raise ValueError(
                "degrees_of_freedom_for_signal: `whitened=True` uses the "
                "whitened prior B_ξ = I; do not pass `background_op` (the ξ-space "
                "Hessian already absorbs B via δx = U ξ)."
            )

        def b_inv(z: jax.Array) -> jax.Array:
            return z
    else:
        if background_op is None:
            raise ValueError(
                "degrees_of_freedom_for_signal: model-space mode "
                "(`whitened=False`) requires `background_op = B`; pass the "
                "Hessian from build_cost_x, or set `whitened=True` for a "
                "build_cost_xi (ξ-space) Hessian."
            )
        import gaussx as gx

        def b_inv(z: jax.Array) -> jax.Array:
            return gx.solve(background_op, z)

    # Build a CG-solvable operator for ``Hess⁻¹``.
    hess_op = lx.FunctionLinearOperator(
        lambda v: hessian_vector_product(jnp.asarray(v)),
        input_structure=jax.eval_shape(lambda: jnp.zeros(state_size)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    cg_solver = lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps)

    rng = np.random.default_rng(seed)
    z_batch = rng.choice([-1.0, 1.0], size=(n_probes, state_size))
    total = 0.0
    for z in z_batch:
        z_j = jnp.asarray(z, dtype=jnp.float64)
        # w = B⁻¹ z (model space) or z (whitened, since B_ξ = I).
        w = b_inv(z_j)
        # u = Hess⁻¹ w  (matrix-free CG)
        u = lx.linear_solve(hess_op, w, solver=cg_solver, throw=False).value
        # zᵀ (I − Bₐ B⁻¹) z ≈ zᵀ z − zᵀ (Hess⁻¹ B⁻¹) z
        total += float(jnp.dot(z_j, z_j) - jnp.dot(z_j, u))
    return total / n_probes

posterior_covariance_proxy 

posterior_covariance_proxy(*, hessian_vector_product: Callable[[Array], Array], state_size: int, whitening: WhiteningTransform | None = None, cg_rtol: float = 1e-06, cg_atol: float = 1e-09, cg_max_steps: int = 200) -> Callable[[jax.Array], jax.Array]

Return a callable v → Bₐ v (model-space posterior covariance) via CG.

Bₐ ≈ Hess⁻¹ is the Laplace-approximation posterior covariance. We expose it as a matvec (rather than materialising) so per-pixel variance estimates remain O(state · cg_iters) — fine for moderate scenes.

Pass the Hessian that matches whitening:

  • whitening=None (default) — a model-space Hessian (from :func:plumax.assimilation.cost.build_cost_x); the CG solve returns Hess⁻¹ v = Bₐ v directly.
  • a :class:~plumax.assimilation.control.WhiteningTransform — a whitened Hessian (from :func:plumax.assimilation.cost.build_cost_xi), whose inverse is the ξ-space covariance, not Bₐ. The model-space covariance is recovered as Bₐ = U Hess_ξ⁻¹ Uᵀ, so each apply un-whitens on both sides: v → U (Hess_ξ⁻¹ (Uᵀ v)).
Source code in src/plumax/assimilation/diagnostics.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def posterior_covariance_proxy(
    *,
    hessian_vector_product: Callable[[jax.Array], jax.Array],
    state_size: int,
    whitening: WhiteningTransform | None = None,
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-9,
    cg_max_steps: int = 200,
) -> Callable[[jax.Array], jax.Array]:
    """Return a callable ``v → Bₐ v`` (model-space posterior covariance) via CG.

    ``Bₐ ≈ Hess⁻¹`` is the Laplace-approximation posterior covariance. We
    expose it as a matvec (rather than materialising) so per-pixel variance
    estimates remain ``O(state · cg_iters)`` — fine for moderate scenes.

    Pass the Hessian that matches ``whitening``:

    - ``whitening=None`` (default) — a **model-space** Hessian (from
      :func:`plumax.assimilation.cost.build_cost_x`); the CG solve returns
      ``Hess⁻¹ v = Bₐ v`` directly.
    - a :class:`~plumax.assimilation.control.WhiteningTransform` — a **whitened**
      Hessian (from :func:`plumax.assimilation.cost.build_cost_xi`), whose
      inverse is the ξ-space covariance, *not* ``Bₐ``. The model-space
      covariance is recovered as ``Bₐ = U Hess_ξ⁻¹ Uᵀ``, so each apply
      un-whitens on both sides: ``v → U (Hess_ξ⁻¹ (Uᵀ v))``.
    """

    def matvec(v: jax.Array) -> jax.Array:
        return hessian_vector_product(jnp.asarray(v))

    op = lx.FunctionLinearOperator(
        matvec,
        input_structure=jax.eval_shape(lambda: jnp.zeros(state_size)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )

    def apply(v: jax.Array) -> jax.Array:
        v = jnp.asarray(v)
        # Whitened Hessian: Bₐ = U Hess_ξ⁻¹ Uᵀ, so map into ξ-space first.
        rhs = whitening.project_gradient(v) if whitening is not None else v
        # ``throw=False`` returns the best partial solution at max_steps rather
        # than raising — for diagnostics we'd rather get a noisy estimate than
        # crash the whole notebook on a stiff probe direction.
        sol = lx.linear_solve(
            op,
            rhs,
            solver=lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps),
            throw=False,
        ).value
        return whitening.apply(sol) if whitening is not None else sol

    return apply