Models

The seven analysis methods are peers — sibling equinox.Module classes with no inheritance between them, each owning exactly the assumptions its math requires. Pick by problem structure: a single analysis time with a linear observation operator wants OI or 3DVar; an assimilation window with perfect-model dynamics wants strong-constraint 4DVar; admitting model error turns that into weak-constraint 4DVar; the operational linearise-and-iterate formulation with control-variable transform is incremental 4DVar; and replacing the inner-loop optimiser with a trained ConvLSTM gives 4DVarNet.

Every model exposes .as_analysis_step(), returning a lightweight wrapper that satisfies the pipekit-cycle AnalysisStep Protocol — that is the seam through which all seven plug into VarDACycle / VarSmootherCycle interchangeably.

Classical methods

Closed-form and optimisation-based analyses. OptimalInterpolation is the linear-Gaussian BLUE solution and refuses non-linear observation operators at construction; ThreeDVar minimises the same cost iteratively and accepts non-linear operators; the three 4DVar variants extend the cost over a time window. IncrementalConfig collects the outer/inner-loop knobs of IncrementalFourDVar.

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

OptimalInterpolation

Bases: Module

BLUE / OI — closed-form linear-Gaussian analysis.

\[x^* = x_b + B H^\top (H B H^\top + R)^{-1} (y - H x_b)\]

Algorithm

1. For each sample (jax.vmap), form the masked operator \(m \odot H(\cdot)\) and the innovation \(d = m \odot y - m \odot H(x_b)\).
2. Linearise the masked operator at \(x_b\) to get \(H\) and build the innovation operator \(H B H^\top + R\) without materialising it.
3. Solve \((H B H^\top + R) v = d\) with CG, or with solver when given.
4. Return \(x_b + B H^\top v\).

Attributes:

Name	Type	Description
obs_op	Any	Observation operator. Must be linear (its linearize(x) must return an operator independent of x). Use ThreeDVar for non-linear H.
prior_mean	Float[Array, '*state']	Background \(x_b\).
prior_cov_op	AbstractLinearOperator	Background-error covariance \(B\) as a lineax.AbstractLinearOperator.
obs_cov_op	AbstractLinearOperator	Observation-error covariance \(R\) as a lineax.AbstractLinearOperator.
cg_atol	float	CG absolute tolerance for the inner solve.
cg_rtol	float	CG relative tolerance for the inner solve.
cg_max_steps	int	CG iteration cap.
solver	AbstractLinearSolver | None	Optional linear solver for the innovation system, used instead of CG (the cg_* fields are then ignored), e.g. lineax.Cholesky() to factorise \(HBH^\top + R\) exactly. It is also the default solver of the posterior adapter.

Examples:

With \(B = R = I\), identity \(H\) and everything observed, the analysis splits the innovation in half: \(x^* = y / 2\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
>>> oi = vardax.OptimalInterpolation(
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros((1, 3)),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
>>> xa = oi(batch)
>>> xa.shape
(1, 1, 3)
>>> bool(jnp.allclose(xa, 0.5, atol=1e-4))
True

Unobserved entries (zero mask) keep the background, and an exact Cholesky solver replaces CG:

>>> exact = vardax.OptimalInterpolation(
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros((1, 3)),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     solver=lx.Cholesky(),
... )
>>> gap = vardax.Batch1D(
...     input=jnp.ones((1, 1, 3)), mask=jnp.array([[[1.0, 0.0, 1.0]]])
... )
>>> bool(jnp.allclose(exact(gap), jnp.array([0.5, 0.0, 0.5]), atol=1e-5))
True

Composed with the Laplace adapter, the posterior covariance of this problem is \((I + I)^{-1} = I / 2\):

>>> post = oi.posterior(batch)
>>> bool(jnp.allclose(post.cov.mv(jnp.ones((1, 3))), 0.5, atol=1e-2))
True

Source code in src/vardax/_src/models/optimal_interpolation.py

class OptimalInterpolation(eqx.Module):
    r"""BLUE / OI — closed-form linear-Gaussian analysis.

    $$x^* = x_b + B H^\top (H B H^\top + R)^{-1} (y - H x_b)$$

    Algorithm:
        1. For each sample (``jax.vmap``), form the masked operator
           $m \odot H(\cdot)$ and the innovation
           $d = m \odot y - m \odot H(x_b)$.
        2. Linearise the masked operator at $x_b$ to get $H$ and build the
           innovation operator $H B H^\top + R$ without materialising it.
        3. Solve $(H B H^\top + R) v = d$ with CG, or with ``solver`` when
           given.
        4. Return $x_b + B H^\top v$.

    Attributes:
        obs_op: Observation operator. Must be linear (its
            ``linearize(x)`` must return an operator independent of
            ``x``). Use [`ThreeDVar`][vardax.ThreeDVar] for non-linear
            ``H``.
        prior_mean: Background $x_b$.
        prior_cov_op: Background-error covariance $B$ as a
            ``lineax.AbstractLinearOperator``.
        obs_cov_op: Observation-error covariance $R$ as a
            ``lineax.AbstractLinearOperator``.
        cg_atol: CG absolute tolerance for the inner solve.
        cg_rtol: CG relative tolerance for the inner solve.
        cg_max_steps: CG iteration cap.
        solver: Optional linear solver for the innovation system, used
            instead of CG (the ``cg_*`` fields are then ignored), e.g.
            ``lineax.Cholesky()`` to factorise $HBH^\top + R$ exactly. It
            is also the default solver of the ``posterior`` adapter.

    Examples:
        With $B = R = I$, identity $H$ and everything observed, the
        analysis splits the innovation in half: $x^* = y / 2$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
        >>> oi = vardax.OptimalInterpolation(
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros((1, 3)),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
        >>> xa = oi(batch)
        >>> xa.shape
        (1, 1, 3)
        >>> bool(jnp.allclose(xa, 0.5, atol=1e-4))
        True

        Unobserved entries (zero mask) keep the background, and an exact
        Cholesky solver replaces CG:

        >>> exact = vardax.OptimalInterpolation(
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros((1, 3)),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     solver=lx.Cholesky(),
        ... )
        >>> gap = vardax.Batch1D(
        ...     input=jnp.ones((1, 1, 3)), mask=jnp.array([[[1.0, 0.0, 1.0]]])
        ... )
        >>> bool(jnp.allclose(exact(gap), jnp.array([0.5, 0.0, 0.5]), atol=1e-5))
        True

        Composed with the Laplace adapter, the posterior covariance of this
        problem is $(I + I)^{-1} = I / 2$:

        >>> post = oi.posterior(batch)
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones((1, 3))), 0.5, atol=1e-2))
        True
    """

    obs_op: Any  # ObservationOperator (linear)
    prior_mean: Float[Array, "*state"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator
    cg_atol: float = eqx.field(static=True, default=1e-6)
    cg_rtol: float = eqx.field(static=True, default=1e-6)
    cg_max_steps: int = eqx.field(static=True, default=200)
    solver: lx.AbstractLinearSolver | None = eqx.field(static=True, default=None)

    def __call__(self, batch: Batch1D) -> Float[Array, " batch *state"]:
        """BLUE analysis for each sample in the batch.

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D].

        Returns:
            The analysis of every sample, with a leading batch axis.
        """

        def _one(
            sample_obs: Float[Array, "*obs"], sample_mask: Float[Array, "*obs"]
        ) -> Float[Array, "*state"]:
            return self._analyse(
                self.prior_mean, sample_obs, sample_mask, self.obs_op, self.obs_cov_op
            )

        return jax.vmap(_one)(batch.input, batch.mask)

    def _analyse(
        self,
        background: Float[Array, "*state"],
        sample_obs: Float[Array, "*obs"],
        sample_mask: Float[Array, "*obs"],
        obs_op: Any,
        obs_cov_op: lx.AbstractLinearOperator,
    ) -> Float[Array, "*state"]:
        """BLUE analysis of one sample around the ``background``."""

        # The effective observation operator is mask ⊙ H(·) — apply the
        # mask both to the predicted obs (innovation) AND to the
        # tangent-linear used to assemble (H B H^T + R). Without the
        # mask in `linearize`, missing entries are treated as
        # zero-innovation observations and bias the analysis whenever
        # B couples observed and missing components.
        def masked_obs(state: Float[Array, "*state"]) -> Float[Array, "*obs"]:
            raw = (
                obs_op(state, mask=sample_mask)
                if _accepts_mask(obs_op)
                else obs_op(state)
            )
            return sample_mask * raw

        predicted_obs = masked_obs(background)
        # d = m ⊙ y − m ⊙ H(x_b)
        innovation = (sample_obs * sample_mask) - predicted_obs

        # Build (H B H^T + R) as a composed lineax operator. We
        # tag it positive-semidefinite so lineax.CG accepts it
        # (the composition itself doesn't carry the tag).
        # H = ∂(m ⊙ H(x))/∂x at x_b
        obs_jacobian = lx.JacobianLinearOperator(
            lambda state, _args=None: masked_obs(state), background
        )
        # H B Hᵀ + R
        inner_raw = (
            obs_jacobian @ self.prior_cov_op @ obs_jacobian.transpose() + obs_cov_op
        )
        inner_op = lx.TaggedLinearOperator(inner_raw, lx.positive_semidefinite_tag)

        solver = self.solver or lx.CG(
            atol=self.cg_atol,
            rtol=self.cg_rtol,
            max_steps=self.cg_max_steps,
        )
        # v = (H B Hᵀ + R)⁻¹ d
        dual_weights = lx.linear_solve(inner_op, innovation, solver=solver).value

        # xₐ = x_b + B Hᵀ v
        return background + self.prior_cov_op.mv(
            obs_jacobian.transpose().mv(dual_weights)
        )

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$, and with this model's ``solver``
        when one is set.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch.
        """
        from vardax._src.posterior._precision import model_posterior

        if adapter is None and self.solver is not None:
            from vardax._src.posterior.laplace import LaplaceCovariance

            adapter = LaplaceCovariance(solver=self.solver)
        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _OIAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            An analysis step that returns the BLUE analysis around a
            supplied forecast, treating ``NaN`` observations as missing.
        """
        return _OIAnalysisStep(self)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\), and with this model's solver when one is set.

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch.

Source code in src/vardax/_src/models/optimal_interpolation.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$, and with this model's ``solver``
    when one is set.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch.
    """
    from vardax._src.posterior._precision import model_posterior

    if adapter is None and self.solver is not None:
        from vardax._src.posterior.laplace import LaplaceCovariance

        adapter = LaplaceCovariance(solver=self.solver)
    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _OIAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_OIAnalysisStep	An analysis step that returns the BLUE analysis around a
_OIAnalysisStep	supplied forecast, treating NaN observations as missing.

Source code in src/vardax/_src/models/optimal_interpolation.py

def as_analysis_step(self) -> _OIAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        An analysis step that returns the BLUE analysis around a
        supplied forecast, treating ``NaN`` observations as missing.
    """
    return _OIAnalysisStep(self)

ThreeDVar

Bases: Module

3D variational analysis.

\[J(x) = \tfrac12 \|x - x_b\|^2_{B^{-1}} + \tfrac12 \|y - H(x)\|^2_{R^{-1}}\]

with the residual masked by \(m\) where observations are missing.

\[x^* = \arg\min_x J(x)\]

Algorithm

1. For each sample in the batch (jax.vmap), start from the background \(x_b\) (or the parameterization's initial control).
2. Evaluate the background term \(\tfrac12 \delta x^\top B^{-1} \delta x\) with a cov_solver solve; under a parameterization the penalty \(\mathrm{pen}(\theta)\) replaces it.
3. Evaluate the observation term from the masked residual \(m \odot (y - H(x))\) and an \(R^{-1}\) solve; a point batch uses the \(\sigma_i\)-weighted point cost instead.
4. Minimise the sum with minimiser for at most max_steps iterations and return the minimising state.

Attributes:

Name	Type	Description
obs_op	Any	Observation operator (linear or nonlinear).
prior_mean	Float[Array, '*state']	Background \(x_b\).
prior_cov_op	AbstractLinearOperator	\(B\). lineax.AbstractLinearOperator.
obs_cov_op	AbstractLinearOperator | None	\(R\). lineax.AbstractLinearOperator.
minimiser	AbstractMinimiser | Minimiser	optimistix.AbstractMinimiser for the inner iteration (default optimistix.BFGS(rtol=1e-6, atol=1e-6)), or a Minimiser such as OptaxMinimiser.
minimiser_adjoint	AbstractAdjoint	optimistix.AbstractAdjoint for differentiating through the minimum (used only if the ThreeDVar analysis is itself nested inside a larger training loop). Default ImplicitAdjoint — exact at the optimum.
max_steps	int	Iteration cap on the inner solver when minimiser is an optimistix solver; a vardax Minimiser keeps its own.
cov_solver	AbstractLinearSolver	lineax solver for the covariance solves in the cost (\(B^{-1}\), \(R^{-1}\)). Default lx.CG(atol=1e-6, rtol=1e-6), which suits large matrix-free operators but is inexact on an ill-conditioned covariance (for example a tiny \(R\)); pass lx.AutoLinearSolver(well_posed=True) or lx.Cholesky() for an exact solve when the operator can be materialised.
parameterization	StateParameterization | None	Optional StateParameterization. When given, the control is \(\theta\), the state is \(\psi(\theta)\) (reshaped to prior_mean.shape) and \(\mathrm{pen}(\theta)\) replaces the background term, so prior_mean and prior_cov_op are not used — nor is the forecast of the analysis-step adapter. Default None: the control is the state itself.

Examples:

With \(B = R = I\), identity \(H\) and everything observed, the analysis agrees with the closed-form BLUE, \(x^* = y / 2\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
>>> three = vardax.ThreeDVar(
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros((1, 3)),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
>>> xa = three(batch)
>>> xa.shape
(1, 1, 3)
>>> bool(jnp.allclose(xa, 0.5, atol=1e-3))
True

A nonlinear \(H(x) = x^2\) through CallableObs. Where the mask is zero the analysis falls back to the background, and analyse also returns the minimiser diagnostics:

>>> sq = vardax.ThreeDVar(
...     obs_op=vardax.CallableObs(fn=lambda x: x**2),
...     prior_mean=jnp.ones((1, 3)),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> half = vardax.Batch1D(
...     input=jnp.full((1, 1, 3), 4.0),
...     mask=jnp.array([[[1.0, 0.0, 0.0]]]),
... )
>>> xa, result = sq.analyse(half)
>>> xa.shape, result.cost.shape
((1, 1, 3), (1,))
>>> bool(jnp.allclose(xa[0, 0, 1:], 1.0, atol=1e-3))
True

Composition with a posterior adapter: the default LaplaceCovariance returns \(P^* = (I + I)^{-1}\) for the first example's model.

>>> post = three.posterior(batch)
>>> bool(jnp.allclose(post.cov.mv(jnp.ones((1, 3))), 0.5, atol=1e-2))
True

Source code in src/vardax/_src/models/threedvar.py

class ThreeDVar(eqx.Module):
    r"""3D variational analysis.

    $$J(x) = \tfrac12 \|x - x_b\|^2_{B^{-1}} + \tfrac12 \|y - H(x)\|^2_{R^{-1}}$$

    with the residual masked by $m$ where observations are missing.

    $$x^* = \arg\min_x J(x)$$

    Algorithm:
        1. For each sample in the batch (``jax.vmap``), start from the
           background $x_b$ (or the parameterization's initial control).
        2. Evaluate the background term
           $\tfrac12 \delta x^\top B^{-1} \delta x$ with a ``cov_solver`` solve; under a
           ``parameterization`` the penalty $\mathrm{pen}(\theta)$
           replaces it.
        3. Evaluate the observation term from the masked residual
           $m \odot (y - H(x))$ and an $R^{-1}$ solve; a point batch uses
           the $\sigma_i$-weighted point cost instead.
        4. Minimise the sum with ``minimiser`` for at most ``max_steps``
           iterations and return the minimising state.

    Attributes:
        obs_op: Observation operator (linear or nonlinear).
        prior_mean: Background $x_b$.
        prior_cov_op: $B$. ``lineax.AbstractLinearOperator``.
        obs_cov_op: $R$. ``lineax.AbstractLinearOperator``.
        minimiser: ``optimistix.AbstractMinimiser`` for the inner
            iteration (default ``optimistix.BFGS(rtol=1e-6, atol=1e-6)``),
            or a [`Minimiser`][vardax.Minimiser] such as
            [`OptaxMinimiser`][vardax.OptaxMinimiser].
        minimiser_adjoint: ``optimistix.AbstractAdjoint`` for
            differentiating *through* the minimum (used only if the
            ``ThreeDVar`` analysis is itself nested inside a larger
            training loop). Default ``ImplicitAdjoint`` — exact at the
            optimum.
        max_steps: Iteration cap on the inner solver when ``minimiser``
            is an ``optimistix`` solver; a vardax `Minimiser` keeps its own.
        cov_solver: ``lineax`` solver for the covariance solves in the
            cost ($B^{-1}$, $R^{-1}$). Default
            ``lx.CG(atol=1e-6, rtol=1e-6)``, which suits large
            matrix-free operators but is inexact on an ill-conditioned
            covariance (for example a tiny $R$); pass
            ``lx.AutoLinearSolver(well_posed=True)`` or ``lx.Cholesky()``
            for an exact solve when the operator can be materialised.
        parameterization: Optional
            [`StateParameterization`][vardax.StateParameterization]. When
            given, the control is $\theta$, the state is
            $\psi(\theta)$ (reshaped to ``prior_mean.shape``) and
            $\mathrm{pen}(\theta)$ replaces the background term, so
            ``prior_mean`` and ``prior_cov_op`` are not used — nor is the
            ``forecast`` of the analysis-step adapter. Default ``None``:
            the control is the state itself.

    Examples:
        With $B = R = I$, identity $H$ and everything observed, the
        analysis agrees with the closed-form BLUE, $x^* = y / 2$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((1, 3), jnp.float32))
        >>> three = vardax.ThreeDVar(
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros((1, 3)),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
        >>> xa = three(batch)
        >>> xa.shape
        (1, 1, 3)
        >>> bool(jnp.allclose(xa, 0.5, atol=1e-3))
        True

        A nonlinear $H(x) = x^2$ through
        [`CallableObs`][vardax.CallableObs]. Where the mask is zero the
        analysis falls back to the background, and ``analyse`` also
        returns the minimiser diagnostics:

        >>> sq = vardax.ThreeDVar(
        ...     obs_op=vardax.CallableObs(fn=lambda x: x**2),
        ...     prior_mean=jnp.ones((1, 3)),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> half = vardax.Batch1D(
        ...     input=jnp.full((1, 1, 3), 4.0),
        ...     mask=jnp.array([[[1.0, 0.0, 0.0]]]),
        ... )
        >>> xa, result = sq.analyse(half)
        >>> xa.shape, result.cost.shape
        ((1, 1, 3), (1,))
        >>> bool(jnp.allclose(xa[0, 0, 1:], 1.0, atol=1e-3))
        True

        Composition with a posterior adapter: the default
        [`LaplaceCovariance`][vardax.LaplaceCovariance] returns
        $P^* = (I + I)^{-1}$ for the first example's model.

        >>> post = three.posterior(batch)
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones((1, 3))), 0.5, atol=1e-2))
        True
    """

    obs_op: Any
    prior_mean: Float[Array, "*state"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator | None
    minimiser: optx.AbstractMinimiser | Minimiser = eqx.field(static=True)
    minimiser_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    cov_solver: lx.AbstractLinearSolver = eqx.field(static=True)
    max_steps: int = eqx.field(static=True, default=200)
    parameterization: StateParameterization | None = None

    def __init__(
        self,
        obs_op: Any,
        prior_mean: Float[Array, "*state"],
        prior_cov_op: lx.AbstractLinearOperator,
        obs_cov_op: lx.AbstractLinearOperator | None,
        minimiser: optx.AbstractMinimiser | Minimiser | None = None,
        minimiser_adjoint: optx.AbstractAdjoint | None = None,
        max_steps: int = 200,
        parameterization: StateParameterization | None = None,
        cov_solver: lx.AbstractLinearSolver | None = None,
    ) -> None:
        self.obs_op = obs_op
        self.prior_mean = prior_mean
        self.prior_cov_op = prior_cov_op
        self.obs_cov_op = obs_cov_op
        self.minimiser = minimiser or optx.BFGS(rtol=1e-6, atol=1e-6)
        self.minimiser_adjoint = minimiser_adjoint or optx.ImplicitAdjoint()
        self.max_steps = max_steps
        self.cov_solver = cov_solver or lx.CG(atol=1e-6, rtol=1e-6)
        self.parameterization = parameterization

    def __call__(self, batch: Batch1D | PointObsBatch) -> Float[Array, " batch *state"]:
        r"""3DVar analysis for each sample in the batch.

        A [`PointObsBatch`][vardax.PointObsBatch] needs a point operator as
        ``obs_op``: a [`GridPointObs`][vardax.GridPointObs] on the state
        ``(T, N)`` (``time_index`` selects the row), or a
        [`FieldPointObs`][vardax.FieldPointObs] with a
        [`NeuralFieldParameterization`][vardax.NeuralFieldParameterization],
        which evaluates the field at the coordinates directly. The error
        variances are the batch's $\sigma_i^2$; ``obs_cov_op`` must be
        ``None`` or an identity.

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D] or a
                [`PointObsBatch`][vardax.PointObsBatch].

        Returns:
            The analysis of every sample, with a leading batch axis.

        Raises:
            ObservationError: For a point batch with a non-trivial
                ``obs_cov_op``.
            MissingInputError: For a gridded batch when ``obs_cov_op`` is
                ``None``.
            InputTypeError: For a point batch whose ``obs_op`` is not a
                point operator, or a ``FieldPointObs`` without a
                ``NeuralFieldParameterization``.
        """
        if isinstance(batch, PointObsBatch):
            check_point_obs_cov(self.obs_cov_op)
            return map_point_batch(
                lambda window: self._analyse_points(self.prior_mean, window), batch
            )

        obs_cov_op = require_obs_cov(self.obs_cov_op)

        def _one(
            sample_obs: Float[Array, "*obs"], sample_mask: Float[Array, "*obs"]
        ) -> Float[Array, "*state"]:
            return self._analyse(
                self.prior_mean, sample_obs, sample_mask, self.obs_op, obs_cov_op
            )

        return jax.vmap(_one)(batch.input, batch.mask)

    def analyse(
        self, batch: Batch1D | PointObsBatch
    ) -> tuple[Float[Array, " batch *state"], MinimiseResult]:
        r"""The analysis, as ``__call__``, with the minimiser's diagnostics.

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D].

        Returns:
            ``(analysis, result)``: the analysis ``(B, T, N)`` and a
            [`MinimiseResult`][vardax.MinimiseResult] whose leaves carry a
            leading sample axis: ``cost`` $J$ and ``grad_norm``
            $\lVert \nabla J \rVert_2$ at the minimum, ``num_steps`` (``-1``
            for a custom minimiser) and ``value``, the minimising control
            (the state, or $\theta$ under a ``parameterization``).

        Raises:
            InputTypeError: For a ``PointObsBatch``.
            MissingInputError: If ``obs_cov_op`` is ``None``.
        """
        if not isinstance(batch, Batch1D):
            raise InputTypeError(
                f"ThreeDVar.analyse takes a Batch1D; got {type(batch).__name__}."
            )
        obs_cov_op = require_obs_cov(self.obs_cov_op)

        def _one(
            sample_obs: Float[Array, "*obs"], sample_mask: Float[Array, "*obs"]
        ) -> tuple[Float[Array, "*state"], MinimiseResult]:
            return self._analyse(
                self.prior_mean,
                sample_obs,
                sample_mask,
                self.obs_op,
                obs_cov_op,
                with_result=True,
            )

        return jax.vmap(_one)(batch.input, batch.mask)

    def _analyse(
        self,
        background: Float[Array, "*state"],
        sample_obs: Float[Array, "*obs"],
        sample_mask: Float[Array, "*obs"],
        obs_op: Any,
        obs_cov_op: lx.AbstractLinearOperator,
        *,
        with_result: bool = False,
    ) -> Any:
        """Minimise the 3D-Var cost of one sample around the ``background``."""

        def obs_cost(state: Float[Array, "*state"]) -> Float[Array, ""]:
            # Observation term — apply R^{-1} via lineax solve.
            # Mask predictions AND observations symmetrically: for
            # obs_ops that don't take a `mask` kwarg (e.g.
            # LinearObs), unmasked predictions at missing entries
            # would otherwise penalise the analysis there.
            predicted_obs = (
                obs_op(state, mask=sample_mask)
                if _accepts_mask(obs_op)
                else obs_op(state)
            )
            # r = m ⊙ (y − H(x))
            residual = sample_mask * (sample_obs - predicted_obs)
            # R⁻¹r
            weighted_residual = lx.linear_solve(
                obs_cov_op,
                residual,
                solver=self.cov_solver,
            ).value
            # J_o = ½ rᵀ R⁻¹ r
            return 0.5 * jnp.sum(residual * weighted_residual)

        return self._solve(background, obs_cost, with_result=with_result)

    def _analyse_points(
        self, background: Float[Array, "*state"], batch: PointObsBatch
    ) -> Float[Array, "*state"]:
        """Minimise the 3D-Var cost of one point-observation sample."""
        if isinstance(self.obs_op, FieldPointObs):
            param = self.parameterization
            if not isinstance(param, NeuralFieldParameterization):
                raise InputTypeError(
                    "FieldPointObs evaluates a field: it needs "
                    "parameterization=NeuralFieldParameterization(...)."
                )
            field_op = self.obs_op

            def control_obs_cost(theta: Any) -> Float[Array, ""]:
                return point_obs_cost(field_op(param.as_field(theta), batch), batch)

            return self._solve(background, None, control_obs_cost=control_obs_cost)

        def obs_cost(state: Float[Array, "*state"]) -> Float[Array, ""]:
            return grid_point_obs_cost(self.obs_op, state, batch)

        return self._solve(background, obs_cost)

    def _solve(
        self,
        background: Float[Array, "*state"],
        obs_cost: Callable[[Array], Array] | None,
        *,
        control_obs_cost: Callable[[Any], Array] | None = None,
        with_result: bool = False,
    ) -> Any:
        """Minimise background (or penalty) plus ``obs_cost`` from ``background``.

        Returns the analysis, or ``(analysis, MinimiseResult)`` when
        ``with_result``.
        """
        run = minimise_with_result if with_result else minimise
        param = self.parameterization
        if param is not None:
            # J(θ) = pen(θ) + J_o(ψ(θ)): the penalty replaces the B term.
            def decode(theta: Any) -> Float[Array, "*state"]:
                return as_shape(param.decode(theta), background.shape)

            def param_cost(theta: Any, _args: Any) -> Float[Array, ""]:
                if control_obs_cost is not None:
                    return param.penalty(theta) + control_obs_cost(theta)
                assert obs_cost is not None
                return param.penalty(theta) + obs_cost(decode(theta))

            result = run(
                self.minimiser,
                param_cost,
                param.init_control(),
                adjoint=self.minimiser_adjoint,
                max_steps=self.max_steps,
            )
            if with_result:
                return decode(result.value), result
            return decode(result)

        def cost(state: Float[Array, "*state"], _args: Any) -> Float[Array, ""]:
            # Background term — apply B^{-1} via lineax solve.
            # δx = x − x_b
            increment = state - background
            # B⁻¹δx
            weighted_increment = lx.linear_solve(
                self.prior_cov_op,
                increment,
                solver=self.cov_solver,
            ).value
            # J_b = ½ δxᵀ B⁻¹ δx
            background_cost = 0.5 * jnp.sum(increment * weighted_increment)
            assert obs_cost is not None
            return background_cost + obs_cost(state)

        result = run(
            self.minimiser,
            cost,
            background,
            adjoint=self.minimiser_adjoint,
            max_steps=self.max_steps,
        )
        if with_result:
            return result.value, result
        return result

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch.
        """
        from vardax._src.posterior._precision import model_posterior

        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _ThreeDVarAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            An analysis step that minimises the cost around a supplied
            forecast, treating ``NaN`` observations as missing.
        """
        return _ThreeDVarAnalysisStep(self)

analyse

analyse(
    batch: Batch1D | PointObsBatch,
) -> tuple[Float[Array, " batch *state"], MinimiseResult]

The analysis, as __call__, with the minimiser's diagnostics.

Parameters:

Name	Type	Description	Default
batch	Batch1D | PointObsBatch	A gridded Batch1D.	required

Returns:

Type	Description
Float[Array, ' batch *state']	(analysis, result): the analysis (B, T, N) and a
MinimiseResult	MinimiseResult whose leaves carry a
tuple[Float[Array, ' batch *state'], MinimiseResult]	leading sample axis: cost \(J\) and grad_norm
tuple[Float[Array, ' batch *state'], MinimiseResult]	\(\lVert \nabla J \rVert_2\) at the minimum, num_steps (-1
tuple[Float[Array, ' batch *state'], MinimiseResult]	for a custom minimiser) and value, the minimising control
tuple[Float[Array, ' batch *state'], MinimiseResult]	(the state, or \(\theta\) under a parameterization).

Raises:

Type	Description
InputTypeError	For a PointObsBatch.
MissingInputError	If obs_cov_op is None.

Source code in src/vardax/_src/models/threedvar.py

def analyse(
    self, batch: Batch1D | PointObsBatch
) -> tuple[Float[Array, " batch *state"], MinimiseResult]:
    r"""The analysis, as ``__call__``, with the minimiser's diagnostics.

    Args:
        batch: A gridded [`Batch1D`][vardax.Batch1D].

    Returns:
        ``(analysis, result)``: the analysis ``(B, T, N)`` and a
        [`MinimiseResult`][vardax.MinimiseResult] whose leaves carry a
        leading sample axis: ``cost`` $J$ and ``grad_norm``
        $\lVert \nabla J \rVert_2$ at the minimum, ``num_steps`` (``-1``
        for a custom minimiser) and ``value``, the minimising control
        (the state, or $\theta$ under a ``parameterization``).

    Raises:
        InputTypeError: For a ``PointObsBatch``.
        MissingInputError: If ``obs_cov_op`` is ``None``.
    """
    if not isinstance(batch, Batch1D):
        raise InputTypeError(
            f"ThreeDVar.analyse takes a Batch1D; got {type(batch).__name__}."
        )
    obs_cov_op = require_obs_cov(self.obs_cov_op)

    def _one(
        sample_obs: Float[Array, "*obs"], sample_mask: Float[Array, "*obs"]
    ) -> tuple[Float[Array, "*state"], MinimiseResult]:
        return self._analyse(
            self.prior_mean,
            sample_obs,
            sample_mask,
            self.obs_op,
            obs_cov_op,
            with_result=True,
        )

    return jax.vmap(_one)(batch.input, batch.mask)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\).

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch.

Source code in src/vardax/_src/models/threedvar.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch.
    """
    from vardax._src.posterior._precision import model_posterior

    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _ThreeDVarAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_ThreeDVarAnalysisStep	An analysis step that minimises the cost around a supplied
_ThreeDVarAnalysisStep	forecast, treating NaN observations as missing.

Source code in src/vardax/_src/models/threedvar.py

def as_analysis_step(self) -> _ThreeDVarAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        An analysis step that minimises the cost around a supplied
        forecast, treating ``NaN`` observations as missing.
    """
    return _ThreeDVarAnalysisStep(self)

StrongFourDVar

Bases: Module

Strong-constraint 4DVar.

\[J(x_0) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}} + \tfrac12 \sum_t \|m_t \odot (y_t - H(M_t(x_0)))\|^2_{R^{-1}}\]

Algorithm

1. For each window (jax.vmap), start from the background \(x_b\) (or the parameterization's initial control).
2. Roll the forward model out from \(x_0\) for one step fewer than the window has times, with optional jax.checkpoint rematerialisation.
3. Sum the masked, \(R^{-1}\)-weighted observation misfit over every time of the trajectory; a point batch uses the \(\sigma_i\)-weighted cost at each time_index instead.
4. Add the background term \(\tfrac12 \delta x_0^\top B^{-1} \delta x_0\) (or the parameterization penalty).
5. Minimise over \(x_0\) with minimiser, differentiating through the rollout by reverse-mode autodiff.

Attributes:

Name	Type	Description
forward	Any	pipekit_cycle.ForwardModel supplying step(state, dt) -> state.
obs_op	Any	Observation operator.
prior_mean	Float[Array, ' N']	Background \(x_b\) — the initial state vector at \(t=0\).
prior_cov_op	AbstractLinearOperator	\(B\).
obs_cov_op	AbstractLinearOperator | None	\(R\).
minimiser	AbstractMinimiser | Minimiser	optimistix.AbstractMinimiser for the outer optimisation over \(x_0\) (default optx.BFGS(rtol=1e-6, atol=1e-6)), or a Minimiser such as OptaxMinimiser for a large control.
minimiser_adjoint	AbstractAdjoint	optimistix.AbstractAdjoint for differentiating through the minimum.
forward_adjoint	Any	diffrax.AbstractAdjoint passed to forward.with_adjoint(...) before the rollout. Requires a forward implementing SupportsForwardAdjoint; any other forward raises UnsupportedInputError. Default None (use the forward's own).
max_steps	int	Iteration cap on the outer solver when minimiser is an optimistix solver; a vardax Minimiser keeps its own.
cov_solver	AbstractLinearSolver	lineax solver for the covariance solves in the cost (\(B^{-1}\), \(R^{-1}\)). Default lx.CG(atol=1e-6, rtol=1e-6), which suits large matrix-free operators but is inexact on an ill-conditioned covariance (for example a tiny \(R\)); pass lx.AutoLinearSolver(well_posed=True) or lx.Cholesky() for an exact solve when the operator can be materialised.
parameterization	StateParameterization | None	Optional StateParameterization. When given, the control is \(\theta\), \(x_0 = \psi(\theta)\) (reshaped to prior_mean.shape) and \(\mathrm{pen}(\theta)\) replaces the background term, so prior_mean and prior_cov_op are not used. Default None: the control is \(x_0\).
checkpoint	bool | int	Rematerialise the rollout in reverse mode. False (default) stores every step's intermediates; True wraps each step in jax.checkpoint; an int k wraps blocks of k steps. Same gradient, less memory, about one extra forward run.

Examples:

Trivial dynamics (\(M_t(x) = x\)) and a single timestep reduce strong 4DVar to 3DVar, so with \(B = R = I\) the analysis is \(x_0^* = y / 2\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> class Identity:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return x
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
>>> strong = vardax.StrongFourDVar(
...     forward=Identity(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
>>> strong(batch).shape
(1, 3)

A decaying forward model \(M(x) = x / 2\) and observations at \(t = 0\) and \(t = 1\) both inform \(x_0\): with \(x_b = 0\) the analysis is \(x_0^* = (1 + 1/2) / (1 + 1 + 1/4) = 2/3\).

>>> class Decay:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return 0.5 * x
>>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
>>> decay = vardax.StrongFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> bool(jnp.allclose(decay(window), 2 / 3, atol=1e-3))
True

checkpoint=True rematerialises the rollout in reverse mode and gives the same analysis; the Laplace posterior of the window is \(P^* = 1 / 2.25\):

>>> remat = vardax.StrongFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     checkpoint=True,
... )
>>> bool(jnp.allclose(remat(window), decay(window), atol=1e-4))
True
>>> post = remat.posterior(window)
>>> bool(jnp.allclose(post.cov.mv(jnp.ones(3)), 1 / 2.25, atol=1e-3))
True

A forward without with_adjoint cannot take a forward_adjoint:

>>> vardax.StrongFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     forward_adjoint=object(),
... )
Traceback (most recent call last):
...
UnsupportedInputError: forward_adjoint was set but Decay has no ...

Source code in src/vardax/_src/models/strong_fourdvar.py

class StrongFourDVar(eqx.Module):
    r"""Strong-constraint 4DVar.

    $$J(x_0) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}}
    + \tfrac12 \sum_t \|m_t \odot (y_t - H(M_t(x_0)))\|^2_{R^{-1}}$$

    Algorithm:
        1. For each window (``jax.vmap``), start from the background
           $x_b$ (or the parameterization's initial control).
        2. Roll the forward model out from $x_0$ for one step fewer than
           the window has times, with optional ``jax.checkpoint``
           rematerialisation.
        3. Sum the masked, $R^{-1}$-weighted observation misfit over every
           time of the trajectory; a point batch uses the $\sigma_i$-weighted
           cost at each ``time_index`` instead.
        4. Add the background term
           $\tfrac12 \delta x_0^\top B^{-1} \delta x_0$
           (or the parameterization penalty).
        5. Minimise over $x_0$ with ``minimiser``, differentiating through
           the rollout by reverse-mode autodiff.

    Attributes:
        forward: ``pipekit_cycle.ForwardModel`` supplying
            ``step(state, dt) -> state``.
        obs_op: Observation operator.
        prior_mean: Background $x_b$ — the initial state vector at
            $t=0$.
        prior_cov_op: $B$.
        obs_cov_op: $R$.
        minimiser: ``optimistix.AbstractMinimiser`` for the outer
            optimisation over $x_0$ (default ``optx.BFGS(rtol=1e-6,
            atol=1e-6)``), or a [`Minimiser`][vardax.Minimiser] such as
            [`OptaxMinimiser`][vardax.OptaxMinimiser] for a large control.
        minimiser_adjoint: ``optimistix.AbstractAdjoint`` for
            differentiating through the minimum.
        forward_adjoint: ``diffrax.AbstractAdjoint`` passed to
            ``forward.with_adjoint(...)`` before the rollout. Requires a
            forward implementing
            [`SupportsForwardAdjoint`][vardax.SupportsForwardAdjoint];
            any other forward raises ``UnsupportedInputError``. Default
            ``None``
            (use the forward's own).
        max_steps: Iteration cap on the outer solver when ``minimiser``
            is an ``optimistix`` solver; a vardax `Minimiser` keeps its own.
        cov_solver: ``lineax`` solver for the covariance solves in the
            cost ($B^{-1}$, $R^{-1}$). Default
            ``lx.CG(atol=1e-6, rtol=1e-6)``, which suits large
            matrix-free operators but is inexact on an ill-conditioned
            covariance (for example a tiny $R$); pass
            ``lx.AutoLinearSolver(well_posed=True)`` or ``lx.Cholesky()``
            for an exact solve when the operator can be materialised.
        parameterization: Optional
            [`StateParameterization`][vardax.StateParameterization]. When
            given, the control is $\theta$, $x_0 = \psi(\theta)$ (reshaped
            to ``prior_mean.shape``) and $\mathrm{pen}(\theta)$ replaces
            the background term, so ``prior_mean`` and ``prior_cov_op``
            are not used. Default ``None``: the control is $x_0$.
        checkpoint: Rematerialise the rollout in reverse mode. ``False``
            (default) stores every step's intermediates; ``True`` wraps
            each step in ``jax.checkpoint``; an int ``k`` wraps blocks of
            ``k`` steps. Same gradient, less memory, about one extra
            forward run.

    Examples:
        Trivial dynamics ($M_t(x) = x$) and a single timestep reduce
        strong 4DVar to 3DVar, so with $B = R = I$ the analysis is
        $x_0^* = y / 2$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> class Identity:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return x
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> strong = vardax.StrongFourDVar(
        ...     forward=Identity(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
        >>> strong(batch).shape
        (1, 3)

        A decaying forward model $M(x) = x / 2$ and observations at $t = 0$
        and $t = 1$ both inform $x_0$: with $x_b = 0$ the analysis is
        $x_0^* = (1 + 1/2) / (1 + 1 + 1/4) = 2/3$.

        >>> class Decay:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return 0.5 * x
        >>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
        >>> decay = vardax.StrongFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> bool(jnp.allclose(decay(window), 2 / 3, atol=1e-3))
        True

        ``checkpoint=True`` rematerialises the rollout in reverse mode and
        gives the same analysis; the Laplace posterior of the window is
        $P^* = 1 / 2.25$:

        >>> remat = vardax.StrongFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     checkpoint=True,
        ... )
        >>> bool(jnp.allclose(remat(window), decay(window), atol=1e-4))
        True
        >>> post = remat.posterior(window)
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones(3)), 1 / 2.25, atol=1e-3))
        True

        A forward without ``with_adjoint`` cannot take a ``forward_adjoint``:

        >>> vardax.StrongFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     forward_adjoint=object(),
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        UnsupportedInputError: forward_adjoint was set but Decay has no ...
    """

    forward: Any
    obs_op: Any
    prior_mean: Float[Array, " N"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator | None
    minimiser: optx.AbstractMinimiser | Minimiser = eqx.field(static=True)
    minimiser_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    forward_adjoint: Any = eqx.field(static=True)
    cov_solver: lx.AbstractLinearSolver = eqx.field(static=True)
    max_steps: int = eqx.field(static=True, default=200)
    parameterization: StateParameterization | None = None
    checkpoint: bool | int = eqx.field(static=True, default=False)

    def __init__(
        self,
        forward: Any,
        obs_op: Any,
        prior_mean: Float[Array, " N"],
        prior_cov_op: lx.AbstractLinearOperator,
        obs_cov_op: lx.AbstractLinearOperator | None,
        minimiser: optx.AbstractMinimiser | Minimiser | None = None,
        minimiser_adjoint: optx.AbstractAdjoint | None = None,
        forward_adjoint: Any = None,
        max_steps: int = 200,
        parameterization: StateParameterization | None = None,
        checkpoint: bool | int = False,
        cov_solver: lx.AbstractLinearSolver | None = None,
    ) -> None:
        # fail at construction, not on the first traced call
        checkpoint_every(checkpoint)
        resolve_forward(forward, forward_adjoint)
        self.forward = forward
        self.obs_op = obs_op
        self.prior_mean = prior_mean
        self.prior_cov_op = prior_cov_op
        self.obs_cov_op = obs_cov_op
        self.minimiser = minimiser or optx.BFGS(rtol=1e-6, atol=1e-6)
        self.minimiser_adjoint = minimiser_adjoint or optx.ImplicitAdjoint()
        self.forward_adjoint = forward_adjoint
        self.max_steps = max_steps
        self.cov_solver = cov_solver or lx.CG(atol=1e-6, rtol=1e-6)
        self.parameterization = parameterization
        self.checkpoint = checkpoint

    def _rollout(
        self, initial_state: Float[Array, " N"], n_steps: int
    ) -> Float[Array, "T N"]:
        """Step the forward model ``n_steps`` times, returning ``T + 1`` states."""
        forward = resolve_forward(self.forward, self.forward_adjoint)
        return rollout(forward, initial_state, n_steps, checkpoint=self.checkpoint)

    def __call__(self, batch: Batch1D | PointObsBatch) -> Float[Array, "*batch N"]:
        r"""Strong-4DVar analysis: minimise over $x_0$.

        A [`PointObsBatch`][vardax.PointObsBatch] needs
        ``obs_op=GridPointObs(...)``: observation $i$ is compared with the
        trajectory state $x_{\tau_i}$, $\tau_i$ = ``time_index[i]``, over a
        window of ``n_times`` states, with error variance $\sigma_i^2$
        (``obs_cov_op`` must be ``None`` or an identity).

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D] whose time axis
                holds the window, or a
                [`PointObsBatch`][vardax.PointObsBatch].

        Returns:
            The analysed initial state $x_0^*$ of every window.

        Raises:
            ObservationError: For a point batch with a non-trivial
                ``obs_cov_op`` or an out-of-range ``time_index``.
            MissingInputError: For a gridded batch when ``obs_cov_op`` is
                ``None``, or a point batch with ``time_index`` but no
                ``n_times``.
            InputTypeError: For a point batch whose ``obs_op`` is not a
                ``GridPointObs``.
        """
        if isinstance(batch, PointObsBatch):
            check_point_obs_cov(self.obs_cov_op)
            n_steps = window_steps(batch)

            def _one_points(window: PointObsBatch) -> Float[Array, " N"]:
                def obs_cost(initial_state: Float[Array, " N"]) -> Float[Array, ""]:
                    trajectory = self._rollout(initial_state, n_steps=n_steps)
                    return grid_point_obs_cost(self.obs_op, trajectory, window)

                return self._solve(obs_cost)

            return map_point_batch(_one_points, batch)

        obs_cov_op = require_obs_cov(self.obs_cov_op)
        # rollout steps (input has T+1 timesteps)
        n_steps = batch.input.shape[1] - 1

        def _one(
            window_obs: Float[Array, "T_plus_1 N"],
            window_mask: Float[Array, "T_plus_1 N"],
        ) -> Float[Array, " N"]:
            def obs_cost(initial_state: Float[Array, " N"]) -> Float[Array, ""]:
                # Roll out the trajectory.
                trajectory = self._rollout(initial_state, n_steps=n_steps)

                # Observation term, summed across time.
                def _per_step(state, observation, step_mask):
                    predicted_obs = (
                        self.obs_op(state, mask=step_mask)
                        if _accepts_mask(self.obs_op)
                        else self.obs_op(state)
                    )
                    # r = m ⊙ (y − H(x))
                    residual = step_mask * (observation - predicted_obs)
                    # R⁻¹r
                    weighted_residual = lx.linear_solve(
                        obs_cov_op,
                        residual,
                        solver=self.cov_solver,
                    ).value
                    # J_o,t = ½ rᵀ R⁻¹ r
                    return 0.5 * jnp.sum(residual * weighted_residual)

                per_step_costs = jax.vmap(_per_step)(
                    trajectory, window_obs, window_mask
                )
                return jnp.sum(per_step_costs)

            return self._solve(obs_cost)

        return jax.vmap(_one)(batch.input, batch.mask)

    def _solve(
        self,
        obs_cost: Callable[[Array], Array],
    ) -> Float[Array, " N"]:
        """Minimise background (or penalty) plus ``obs_cost(x_0)``."""
        param = self.parameterization
        if param is not None:
            # x_0 = ψ(θ); J(θ) = pen(θ) + J_o(M(ψ(θ))).
            def decode(theta: Any) -> Float[Array, " N"]:
                return as_shape(param.decode(theta), self.prior_mean.shape)

            def param_cost(theta: Any, _args: Any) -> Float[Array, ""]:
                return param.penalty(theta) + obs_cost(decode(theta))

            theta_star = minimise(
                self.minimiser,
                param_cost,
                param.init_control(),
                adjoint=self.minimiser_adjoint,
                max_steps=self.max_steps,
            )
            return decode(theta_star)

        def cost(initial_state: Float[Array, " N"], _args: Any) -> Float[Array, ""]:
            # Background term
            # δx = x₀ − x_b
            increment = initial_state - self.prior_mean
            # B⁻¹δx
            weighted_increment = lx.linear_solve(
                self.prior_cov_op,
                increment,
                solver=self.cov_solver,
            ).value
            # J_b = ½ δxᵀ B⁻¹ δx
            background_cost = 0.5 * jnp.sum(increment * weighted_increment)
            return background_cost + obs_cost(initial_state)

        return minimise(
            self.minimiser,
            cost,
            self.prior_mean,
            adjoint=self.minimiser_adjoint,
            max_steps=self.max_steps,
        )

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch.
        """
        from vardax._src.posterior._precision import model_posterior

        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _StrongFourDVarAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep``.

        Returns:
            An analysis step that treats non-finite observations as
            missing and returns the analysed initial state.
        """
        return _StrongFourDVarAnalysisStep(self)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\).

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch.

Source code in src/vardax/_src/models/strong_fourdvar.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch.
    """
    from vardax._src.posterior._precision import model_posterior

    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _StrongFourDVarAnalysisStep

Adapt to pipekit_cycle.AnalysisStep.

Returns:

Type	Description
_StrongFourDVarAnalysisStep	An analysis step that treats non-finite observations as
_StrongFourDVarAnalysisStep	missing and returns the analysed initial state.

Source code in src/vardax/_src/models/strong_fourdvar.py

def as_analysis_step(self) -> _StrongFourDVarAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep``.

    Returns:
        An analysis step that treats non-finite observations as
        missing and returns the analysed initial state.
    """
    return _StrongFourDVarAnalysisStep(self)

WeakFourDVar

Bases: Module

Weak-constraint 4DVar with augmented control vector.

\[J(x_0, \eta) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}} + J_o + J_q\]

\[J_o = \tfrac12 \sum_t \|y_t - H(x_t)\|^2_{R^{-1}}\]

\[J_q = \tfrac12 \sum_t \|\eta_t\|^2_{Q^{-1}}\]

\[x_t = M^\text{free}(x_{t-1}) + \eta_t\]

Algorithm

1. Start from the background \(x_0 = x_b\) and zero model errors.
2. Roll the free model out with the forcing \(\eta_t\) added after every step.
3. Sum the masked, \(R^{-1}\)-weighted observation misfit over the trajectory, the \(Q^{-1}\)-weighted model-error cost over the \(T\) steps, and the background term (or the parameterization penalty).
4. Minimise jointly over \((x_0, \eta)\) with minimiser and return the pair.

Attributes:

Name	Type	Description
forward	Any	pipekit_cycle.ForwardModel (the free model \(M_t^\text{free}\)).
obs_op	Any	Observation operator.
prior_mean	Float[Array, ' N']	Background \(x_b\) — the initial state.
prior_cov_op	AbstractLinearOperator	\(B\).
obs_cov_op	AbstractLinearOperator | None	\(R\).
model_err_cov_op	AbstractLinearOperator	\(Q\) — covariance of the per-step model error \(\eta_t\). Defaults to identity scaled by a small variance if not supplied (effectively strong-constraint with a tiny relaxation).
minimiser	AbstractMinimiser | Minimiser	As in StrongFourDVar.
minimiser_adjoint	AbstractAdjoint	As in StrongFourDVar.
max_steps	int	As in StrongFourDVar.
cov_solver	AbstractLinearSolver	lineax solver for the covariance solves in the cost (\(B^{-1}\), \(R^{-1}\), \(Q^{-1}\)). Default lx.CG(atol=1e-6, rtol=1e-6), which suits large matrix-free operators but is inexact on an ill-conditioned covariance (for example a tiny \(R\)); pass lx.AutoLinearSolver(well_posed=True) or lx.Cholesky() for an exact solve when the operator can be materialised.
parameterization	StateParameterization | None	Optional StateParameterization of the initial state. When given, the control is \((\theta, \eta)\), \(x_0 = \psi(\theta)\) and \(\mathrm{pen}(\theta)\) replaces the background term (prior_mean and prior_cov_op are not used). Default None: the control is \((x_0, \eta)\).
forward_adjoint	Any	As in StrongFourDVar: a diffrax.AbstractAdjoint for a forward implementing SupportsForwardAdjoint.
checkpoint	bool | int	As in StrongFourDVar: False, True or a block length k for jax.checkpoint over the rollout.

Examples:

Trivial dynamics and a single timestep: no model-error steps remain (\(T = 0\)) and the analysis initial state reduces to the 3DVar / BLUE answer \(x_0^* = y / 2\) for \(B = R = I\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> class Identity:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return x
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
>>> weak = vardax.WeakFourDVar(
...     forward=Identity(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     model_err_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
>>> x0_star, eta_star = weak(batch)
>>> x0_star.shape, eta_star.shape
((1, 3), (1, 0, 3))

A two-time window has one model-error step. With a tiny model-error covariance the free model is almost exact and the analysis approaches the strong-constraint \(x_0^* = 2/3\) for \(M(x) = x / 2\):

>>> class Decay:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return 0.5 * x
>>> tiny = lx.TaggedLinearOperator(
...     lx.DiagonalLinearOperator(jnp.full(3, 1e-4)),
...     lx.positive_semidefinite_tag,
... )
>>> weak = vardax.WeakFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     model_err_cov_op=tiny,
... )
>>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
>>> x0_star, eta_star = weak(window)
>>> eta_star.shape
(1, 1, 3)
>>> bool(jnp.allclose(x0_star, 2 / 3, atol=5e-2))
True

Composed with the Laplace adapter, posterior returns the \(x_0\) block of the inverse augmented precision:

>>> weak.posterior(window).cov.mv(jnp.ones(3)).shape
(3,)

Source code in src/vardax/_src/models/weak_fourdvar.py

class WeakFourDVar(eqx.Module):
    r"""Weak-constraint 4DVar with augmented control vector.

    $$J(x_0, \eta) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}} + J_o + J_q$$

    $$J_o = \tfrac12 \sum_t \|y_t - H(x_t)\|^2_{R^{-1}}$$

    $$J_q = \tfrac12 \sum_t \|\eta_t\|^2_{Q^{-1}}$$

    $$x_t = M^\text{free}(x_{t-1}) + \eta_t$$

    Algorithm:
        1. Start from the background $x_0 = x_b$ and zero model errors.
        2. Roll the free model out with the forcing $\eta_t$ added after
           every step.
        3. Sum the masked, $R^{-1}$-weighted observation misfit over the
           trajectory, the $Q^{-1}$-weighted model-error cost over the
           $T$ steps, and the background term (or the parameterization
           penalty).
        4. Minimise jointly over $(x_0, \eta)$ with ``minimiser`` and
           return the pair.

    Attributes:
        forward: ``pipekit_cycle.ForwardModel`` (the free model
            $M_t^\text{free}$).
        obs_op: Observation operator.
        prior_mean: Background $x_b$ — the initial state.
        prior_cov_op: $B$.
        obs_cov_op: $R$.
        model_err_cov_op: $Q$ — covariance of the per-step model
            error $\eta_t$. Defaults to identity scaled by a
            small variance if not supplied (effectively
            strong-constraint with a tiny relaxation).
        minimiser: As in [`StrongFourDVar`][vardax.StrongFourDVar].
        minimiser_adjoint: As in [`StrongFourDVar`][vardax.StrongFourDVar].
        max_steps: As in [`StrongFourDVar`][vardax.StrongFourDVar].
        cov_solver: ``lineax`` solver for the covariance solves in the
            cost ($B^{-1}$, $R^{-1}$, $Q^{-1}$). Default
            ``lx.CG(atol=1e-6, rtol=1e-6)``, which suits large
            matrix-free operators but is inexact on an ill-conditioned
            covariance (for example a tiny $R$); pass
            ``lx.AutoLinearSolver(well_posed=True)`` or ``lx.Cholesky()``
            for an exact solve when the operator can be materialised.
        parameterization: Optional
            [`StateParameterization`][vardax.StateParameterization] of the
            initial state. When given, the control is $(\theta, \eta)$,
            $x_0 = \psi(\theta)$ and $\mathrm{pen}(\theta)$ replaces the
            background term (``prior_mean`` and ``prior_cov_op`` are not
            used). Default ``None``: the control is $(x_0, \eta)$.
        forward_adjoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            a ``diffrax.AbstractAdjoint`` for a forward implementing
            [`SupportsForwardAdjoint`][vardax.SupportsForwardAdjoint].
        checkpoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            ``False``, ``True`` or a block length ``k`` for
            ``jax.checkpoint`` over the rollout.

    Examples:
        Trivial dynamics and a single timestep: no model-error steps
        remain ($T = 0$) and the analysis initial state reduces to the
        3DVar / BLUE answer $x_0^* = y / 2$ for $B = R = I$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> class Identity:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return x
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> weak = vardax.WeakFourDVar(
        ...     forward=Identity(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     model_err_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
        >>> x0_star, eta_star = weak(batch)
        >>> x0_star.shape, eta_star.shape
        ((1, 3), (1, 0, 3))

        A two-time window has one model-error step. With a tiny model-error
        covariance the free model is almost exact and the analysis
        approaches the strong-constraint $x_0^* = 2/3$ for $M(x) = x / 2$:

        >>> class Decay:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return 0.5 * x
        >>> tiny = lx.TaggedLinearOperator(
        ...     lx.DiagonalLinearOperator(jnp.full(3, 1e-4)),
        ...     lx.positive_semidefinite_tag,
        ... )
        >>> weak = vardax.WeakFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     model_err_cov_op=tiny,
        ... )
        >>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
        >>> x0_star, eta_star = weak(window)
        >>> eta_star.shape
        (1, 1, 3)
        >>> bool(jnp.allclose(x0_star, 2 / 3, atol=5e-2))
        True

        Composed with the Laplace adapter, ``posterior`` returns the $x_0$
        block of the inverse augmented precision:

        >>> weak.posterior(window).cov.mv(jnp.ones(3)).shape
        (3,)
    """

    forward: Any
    obs_op: Any
    prior_mean: Float[Array, " N"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator | None
    model_err_cov_op: lx.AbstractLinearOperator
    minimiser: optx.AbstractMinimiser | Minimiser = eqx.field(static=True)
    minimiser_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    cov_solver: lx.AbstractLinearSolver = eqx.field(static=True)
    max_steps: int = eqx.field(static=True, default=200)
    parameterization: StateParameterization | None = None
    forward_adjoint: Any = eqx.field(static=True, default=None)
    checkpoint: bool | int = eqx.field(static=True, default=False)

    def __init__(
        self,
        forward: Any,
        obs_op: Any,
        prior_mean: Float[Array, " N"],
        prior_cov_op: lx.AbstractLinearOperator,
        obs_cov_op: lx.AbstractLinearOperator | None,
        model_err_cov_op: lx.AbstractLinearOperator,
        minimiser: optx.AbstractMinimiser | Minimiser | None = None,
        minimiser_adjoint: optx.AbstractAdjoint | None = None,
        max_steps: int = 200,
        parameterization: StateParameterization | None = None,
        forward_adjoint: Any = None,
        checkpoint: bool | int = False,
        cov_solver: lx.AbstractLinearSolver | None = None,
    ) -> None:
        # fail at construction, not on the first traced call
        checkpoint_every(checkpoint)
        resolve_forward(forward, forward_adjoint)
        self.forward = forward
        self.obs_op = obs_op
        self.prior_mean = prior_mean
        self.prior_cov_op = prior_cov_op
        self.obs_cov_op = obs_cov_op
        self.model_err_cov_op = model_err_cov_op
        self.minimiser = minimiser or optx.BFGS(rtol=1e-6, atol=1e-6)
        self.minimiser_adjoint = minimiser_adjoint or optx.ImplicitAdjoint()
        self.max_steps = max_steps
        self.cov_solver = cov_solver or lx.CG(atol=1e-6, rtol=1e-6)
        self.parameterization = parameterization
        self.forward_adjoint = forward_adjoint
        self.checkpoint = checkpoint

    def _rollout_with_error(
        self,
        initial_state: Float[Array, " N"],
        model_errors: Float[Array, "T N"],
    ) -> Float[Array, "T_plus_1 N"]:
        r"""Roll out with per-step model error.

        Each step applies $x_t = M_t^\text{free}(x_{t-1}) + \eta_t$.
        """
        forward = resolve_forward(self.forward, self.forward_adjoint)
        n_steps = model_errors.shape[0]
        return rollout(
            forward,
            initial_state,
            n_steps,
            forcing=model_errors,
            checkpoint=self.checkpoint,
        )

    def __call__(
        self,
        batch: Batch1D | PointObsBatch,
    ) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch T N"]]:
        r"""Weak-4DVar analysis.

        Returns ``(x_0_star, eta_star)`` — the analysis initial state
        and the per-step model-error trajectory (one row per
        timestep, total ``T``).

        A [`PointObsBatch`][vardax.PointObsBatch] needs
        ``obs_op=GridPointObs(...)``: observation $i$ is compared with the
        trajectory state at ``time_index[i]`` over a window of ``n_times``
        states (so ``T = n_times - 1``), with error variance
        $\sigma_i^2$ (``obs_cov_op`` must be ``None`` or an identity).

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D] whose time axis
                holds the window, or a
                [`PointObsBatch`][vardax.PointObsBatch].

        Returns:
            The pair $(x_0^*, \eta^*)$ of every window.

        Raises:
            ObservationError: For a point batch with a non-trivial
                ``obs_cov_op`` or an out-of-range ``time_index``.
            MissingInputError: For a gridded batch when ``obs_cov_op`` is
                ``None``, or a point batch with ``time_index`` but no
                ``n_times``.
            InputTypeError: For a point batch whose ``obs_op`` is not a
                ``GridPointObs``.
        """
        state_dim = self.prior_mean.shape[0]
        if isinstance(batch, PointObsBatch):
            check_point_obs_cov(self.obs_cov_op)
            n_steps = window_steps(batch)

            def _one_points(window: PointObsBatch):
                def traj_obs_cost(
                    trajectory: Float[Array, "T_plus_1 N"],
                ) -> Float[Array, ""]:
                    return grid_point_obs_cost(self.obs_op, trajectory, window)

                return self._solve(traj_obs_cost, n_steps, state_dim)

            return map_point_batch(_one_points, batch)

        obs_cov_op = require_obs_cov(self.obs_cov_op)
        # T_plus_1 timesteps in the batch ⇒ T model-error steps.
        n_times = batch.input.shape[1]
        n_steps = n_times - 1
        state_dim = batch.input.shape[2]

        def _one(
            window_obs: Float[Array, "T_plus_1 N"],
            window_mask: Float[Array, "T_plus_1 N"],
        ):
            def traj_obs_cost(
                trajectory: Float[Array, "T_plus_1 N"],
            ) -> Float[Array, ""]:
                def _per_step(state, observation, step_mask):
                    predicted_obs = (
                        self.obs_op(state, mask=step_mask)
                        if _accepts_mask(self.obs_op)
                        else self.obs_op(state)
                    )
                    # r = m ⊙ (y − H(x))
                    residual = step_mask * (observation - predicted_obs)
                    # R⁻¹r
                    weighted_residual = lx.linear_solve(
                        obs_cov_op,
                        residual,
                        solver=self.cov_solver,
                    ).value
                    # J_o,t = ½ rᵀ R⁻¹ r
                    return 0.5 * jnp.sum(residual * weighted_residual)

                per_step_costs = jax.vmap(_per_step)(
                    trajectory, window_obs, window_mask
                )
                return jnp.sum(per_step_costs)

            return self._solve(traj_obs_cost, n_steps, state_dim)

        analysis, model_error_analysis = jax.vmap(_one)(batch.input, batch.mask)
        return analysis, model_error_analysis

    def _solve(
        self,
        traj_obs_cost: Callable[[Array], Array],
        n_steps: int,
        state_dim: int,
    ) -> tuple[Float[Array, " N"], Float[Array, "T N"]]:
        """Minimise over $(x_0, \\eta)$ (or $(\\theta, \\eta)$) for one sample."""

        def model_and_obs_cost(
            initial_state: Float[Array, " N"],
            model_errors: Float[Array, "T N"],
        ) -> Float[Array, ""]:
            # Model-error term
            def _model_error_cost(step_error):
                # Q⁻¹η
                weighted_error = lx.linear_solve(
                    self.model_err_cov_op,
                    step_error,
                    solver=self.cov_solver,
                ).value
                # J_q,t = ½ ηᵀ Q⁻¹ η
                return 0.5 * jnp.sum(step_error * weighted_error)

            # A one-state window has no model-error steps (vmap over an
            # empty axis would trip lineax's CG on a zero-size operand).
            model_error_cost = (
                jnp.sum(jax.vmap(_model_error_cost)(model_errors))
                if model_errors.shape[0] > 0
                else jnp.zeros((), dtype=initial_state.dtype)
            )

            # Observation term
            obs_cost = traj_obs_cost(
                self._rollout_with_error(initial_state, model_errors)
            )
            return obs_cost + model_error_cost

        # Initial guess: x_0 at the background, etas at zero.
        init_model_errors = jnp.zeros((n_steps, state_dim))

        param = self.parameterization
        if param is not None:
            # Control (θ, η); x_0 = ψ(θ) and pen(θ) replaces the B term.
            def decode(theta: Any) -> Float[Array, " N"]:
                return as_shape(param.decode(theta), self.prior_mean.shape)

            def param_cost(control: Any, _args: Any) -> Float[Array, ""]:
                theta, model_errors = control
                return param.penalty(theta) + model_and_obs_cost(
                    decode(theta), model_errors
                )

            theta_star, model_errors_star = minimise(
                self.minimiser,
                param_cost,
                (param.init_control(), init_model_errors),
                adjoint=self.minimiser_adjoint,
                max_steps=self.max_steps,
            )
            return decode(theta_star), model_errors_star

        def cost(control: Float[Array, "T_plus_1 N"], _args: Any) -> Float[Array, ""]:
            initial_state = control[0]
            model_errors = control[1:]

            # Background term
            # δx = x₀ − x_b
            increment = initial_state - self.prior_mean
            # B⁻¹δx
            weighted_increment = lx.linear_solve(
                self.prior_cov_op,
                increment,
                solver=self.cov_solver,
            ).value
            # J_b = ½ δxᵀ B⁻¹ δx
            background_cost = 0.5 * jnp.sum(increment * weighted_increment)
            return background_cost + model_and_obs_cost(initial_state, model_errors)

        # control [x_b; η₁, …, η_T]; jnp, not einx: einx cannot size T = 0
        initial_state = jnp.expand_dims(self.prior_mean, 0)
        initial_control = jnp.concatenate([initial_state, init_model_errors])

        control_star = minimise(
            self.minimiser,
            cost,
            initial_control,
            adjoint=self.minimiser_adjoint,
            max_steps=self.max_steps,
        )
        return control_star[0], control_star[1:]

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$.

        The covariance is that of $x_0$: the $x_0$ block of the inverse
        Gauss–Newton precision over $(x_0, \\eta)$.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch, or for a model with
                a ``parameterization``.
        """
        from vardax._src.posterior._precision import model_posterior

        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _WeakFourDVarAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep``.

        Returns:
            An analysis step that treats non-finite observations as
            missing and returns the analysed initial state only.
        """
        return _WeakFourDVarAnalysisStep(self)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\).

The covariance is that of \(x_0\): the \(x_0\) block of the inverse Gauss–Newton precision over \((x_0, \eta)\).

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch, or for a model with a parameterization.

Source code in src/vardax/_src/models/weak_fourdvar.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$.

    The covariance is that of $x_0$: the $x_0$ block of the inverse
    Gauss–Newton precision over $(x_0, \\eta)$.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch, or for a model with
            a ``parameterization``.
    """
    from vardax._src.posterior._precision import model_posterior

    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _WeakFourDVarAnalysisStep

Adapt to pipekit_cycle.AnalysisStep.

Returns:

Type	Description
_WeakFourDVarAnalysisStep	An analysis step that treats non-finite observations as
_WeakFourDVarAnalysisStep	missing and returns the analysed initial state only.

Source code in src/vardax/_src/models/weak_fourdvar.py

def as_analysis_step(self) -> _WeakFourDVarAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep``.

    Returns:
        An analysis step that treats non-finite observations as
        missing and returns the analysed initial state only.
    """
    return _WeakFourDVarAnalysisStep(self)

IncrementalFourDVar

Bases: Module

Operational incremental 4DVar (Decision D11).

\[J(x_0) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}} + \tfrac12 \|r(x_0)\|^2_{R^{-1}}\]

\[r_t(x_0) = m_t \odot (y_t - H_t(M_t(x_0)))\]

Algorithm

1. Start from \(x^{(0)} = x_b\) (the control \(\chi = 0\) under the control-variable transform).
2. For each of n_outer iterations, linearise the residual \(r\) at the current iterate with jax.linearize and jax.vjp.
3. Solve the Gauss-Newton system \((P + r'^\top R^{-1} r') \, \delta = -\nabla J\) with CG for at most n_inner steps, where \(P = B^{-1}\) in state space and \(P = I\) under the transform.
4. Update the iterate with \(\delta\) and map back to the state.

Functionally equivalent to StrongFourDVar (same problem, same answer in the converged limit) but with a specialised inner solver: Gauss-Newton outer iterations and CG inner iterations on the linearised cost. Use this for production / long-window 4DVar.

Attributes:

Name	Type	Description
forward	Any	pipekit_cycle.ForwardModel.
obs_op	Any	Observation operator.
prior_mean	Float[Array, ' N']	Background \(x_b\) — the initial state.
prior_cov_op	AbstractLinearOperator	\(B\).
obs_cov_op	AbstractLinearOperator | None	\(R\).
config	IncrementalConfig	IncrementalConfig.
prior_sqrt_op	AbstractLinearOperator | None	Optional square root \(S\) of \(B\) (\(S S^\top = B\)), used by the control-variable transform: a Cholesky factor, a structured operator's own square root, or a Lanczos gaussx.sqrt(B, lanczos_order=k) for a large unstructured \(B\). Default None: gaussx.sqrt(prior_cov_op), closed form for diagonal / Kronecker / block-diagonal \(B\) and a dense eigendecomposition otherwise.
forward_adjoint	Any	As in StrongFourDVar: a diffrax.AbstractAdjoint for a forward implementing SupportsForwardAdjoint.
checkpoint	bool | int	As in StrongFourDVar: False, True or a block length k for jax.checkpoint over the rollout (the outer nonlinear run and its tangent-linear / adjoint).

Examples:

Trivial dynamics (\(M_t(x) = x\)) and a single timestep reduce the problem to 3DVar, so with \(B = R = I\) the converged analysis is \(x_0^* = y / 2\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> class Identity:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return x
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
>>> inc = vardax.IncrementalFourDVar(
...     forward=Identity(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
>>> xa = inc(batch)
>>> xa.shape
(1, 3)
>>> bool(jnp.allclose(xa, 0.5, atol=1e-2))
True

On a linear decaying model \(M(x) = x / 2\) with observations at \(t = 0, 1\) the Gauss-Newton step is exact, so one outer iteration reproduces the strong-constraint answer \(x_0^* = 2/3\):

>>> class Decay:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return 0.5 * x
>>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
>>> one = vardax.IncrementalFourDVar(
...     forward=Decay(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(3),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
...     config=vardax.IncrementalConfig(n_outer=1, cvt=False),
... )
>>> bool(jnp.allclose(one(window), 2 / 3, atol=1e-3))
True

Composed with the Laplace adapter (via posterior), the window's covariance is \(P^* = 1 / 2.25\):

>>> post = one.posterior(window)
>>> bool(jnp.allclose(post.cov.mv(jnp.ones(3)), 1 / 2.25, atol=1e-3))
True

Source code in src/vardax/_src/models/incremental_fourdvar.py

class IncrementalFourDVar(eqx.Module):
    r"""Operational incremental 4DVar (Decision D11).

    $$J(x_0) = \tfrac12 \|x_0 - x_b\|^2_{B^{-1}} + \tfrac12 \|r(x_0)\|^2_{R^{-1}}$$

    $$r_t(x_0) = m_t \odot (y_t - H_t(M_t(x_0)))$$

    Algorithm:
        1. Start from $x^{(0)} = x_b$ (the control $\chi = 0$ under the
           control-variable transform).
        2. For each of ``n_outer`` iterations, linearise the residual
           $r$ at the current iterate with ``jax.linearize`` and
           ``jax.vjp``.
        3. Solve the Gauss-Newton system
           $(P + r'^\top R^{-1} r') \, \delta = -\nabla J$ with CG for at
           most ``n_inner`` steps, where $P = B^{-1}$ in state space and
           $P = I$ under the transform.
        4. Update the iterate with $\delta$ and map back to the state.

    Functionally equivalent to [`StrongFourDVar`][vardax.StrongFourDVar]
    (same problem, same answer in the converged limit) but with a
    specialised inner solver: Gauss-Newton outer iterations and CG
    inner iterations on the linearised cost. Use this for production /
    long-window 4DVar.

    Attributes:
        forward: ``pipekit_cycle.ForwardModel``.
        obs_op: Observation operator.
        prior_mean: Background $x_b$ — the initial state.
        prior_cov_op: $B$.
        obs_cov_op: $R$.
        config: [`IncrementalConfig`][vardax.IncrementalConfig].
        prior_sqrt_op: Optional square root $S$ of $B$ ($S S^\top = B$),
            used by the control-variable transform: a Cholesky factor, a
            structured operator's own square root, or a Lanczos
            ``gaussx.sqrt(B, lanczos_order=k)`` for a large unstructured
            $B$. Default ``None``: ``gaussx.sqrt(prior_cov_op)``, closed
            form for diagonal / Kronecker / block-diagonal $B$ and a dense
            eigendecomposition otherwise.
        forward_adjoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            a ``diffrax.AbstractAdjoint`` for a forward implementing
            [`SupportsForwardAdjoint`][vardax.SupportsForwardAdjoint].
        checkpoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            ``False``, ``True`` or a block length ``k`` for
            ``jax.checkpoint`` over the rollout (the
            outer nonlinear run and its tangent-linear / adjoint).

    Examples:
        Trivial dynamics ($M_t(x) = x$) and a single timestep reduce
        the problem to 3DVar, so with $B = R = I$ the converged
        analysis is $x_0^* = y / 2$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> class Identity:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return x
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> inc = vardax.IncrementalFourDVar(
        ...     forward=Identity(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 3)), mask=jnp.ones((1, 1, 3)))
        >>> xa = inc(batch)
        >>> xa.shape
        (1, 3)
        >>> bool(jnp.allclose(xa, 0.5, atol=1e-2))
        True

        On a linear decaying model $M(x) = x / 2$ with observations at
        $t = 0, 1$ the Gauss-Newton step is exact, so one outer iteration
        reproduces the strong-constraint answer $x_0^* = 2/3$:

        >>> class Decay:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return 0.5 * x
        >>> window = vardax.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
        >>> one = vardax.IncrementalFourDVar(
        ...     forward=Decay(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(3),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ...     config=vardax.IncrementalConfig(n_outer=1, cvt=False),
        ... )
        >>> bool(jnp.allclose(one(window), 2 / 3, atol=1e-3))
        True

        Composed with the Laplace adapter (via ``posterior``), the window's
        covariance is $P^* = 1 / 2.25$:

        >>> post = one.posterior(window)
        >>> bool(jnp.allclose(post.cov.mv(jnp.ones(3)), 1 / 2.25, atol=1e-3))
        True
    """

    forward: Any
    obs_op: Any
    prior_mean: Float[Array, " N"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator | None
    config: IncrementalConfig
    prior_sqrt_op: lx.AbstractLinearOperator | None = None
    forward_adjoint: Any = eqx.field(static=True, default=None)
    checkpoint: bool | int = eqx.field(static=True, default=False)

    def __init__(
        self,
        forward: Any,
        obs_op: Any,
        prior_mean: Float[Array, " N"],
        prior_cov_op: lx.AbstractLinearOperator,
        obs_cov_op: lx.AbstractLinearOperator | None,
        config: IncrementalConfig | None = None,
        forward_adjoint: Any = None,
        checkpoint: bool | int = False,
        prior_sqrt_op: lx.AbstractLinearOperator | None = None,
    ) -> None:
        # fail at construction, not on the first traced call
        checkpoint_every(checkpoint)
        resolve_forward(forward, forward_adjoint)
        self.forward = forward
        self.obs_op = obs_op
        self.prior_mean = prior_mean
        self.prior_cov_op = prior_cov_op
        self.obs_cov_op = obs_cov_op
        self.config = config or IncrementalConfig()
        self.prior_sqrt_op = prior_sqrt_op
        self.forward_adjoint = forward_adjoint
        self.checkpoint = checkpoint

    def _rollout(
        self,
        initial_state: Float[Array, " N"],
        n_steps: int,
    ) -> Float[Array, "T_plus_1 N"]:
        forward = resolve_forward(self.forward, self.forward_adjoint)
        return rollout(forward, initial_state, n_steps, checkpoint=self.checkpoint)

    def __call__(self, batch: Batch1D | PointObsBatch) -> Float[Array, "*batch N"]:
        r"""Incremental-4DVar analysis: GN outer + CG inner.

        A [`PointObsBatch`][vardax.PointObsBatch] needs
        ``obs_op=GridPointObs(...)``: observation $i$ is compared with the
        trajectory state at ``time_index[i]`` over a window of ``n_times``
        states, with error variance $\sigma_i^2$ (``obs_cov_op`` must
        be ``None`` or an identity).

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D] whose time axis
                holds the window, or a
                [`PointObsBatch`][vardax.PointObsBatch].

        Returns:
            The analysed initial state $x_0^*$ of every window.

        Raises:
            ObservationError: For a point batch with a non-trivial
                ``obs_cov_op`` or an out-of-range ``time_index``.
            MissingInputError: For a gridded batch when ``obs_cov_op`` is
                ``None``, or a point batch with ``time_index`` but no
                ``n_times``.
            InputTypeError: For a point batch whose ``obs_op`` is not a
                ``GridPointObs``.
        """
        if isinstance(batch, PointObsBatch):
            check_point_obs_cov(self.obs_cov_op)
            n_steps = window_steps(batch)
            return map_point_batch(
                lambda window: self._analyse_points(window, n_steps), batch
            )

        tol = self.config.covariance_tol
        obs_cov = _as_covariance(require_obs_cov(self.obs_cov_op))

        def apply_obs_precision(
            residuals: Float[Array, "T_plus_1 N"],
        ) -> Float[Array, "T_plus_1 N"]:
            # R⁻¹ rₜ at each time
            def solve_one(step_residual: Float[Array, " N"]) -> Float[Array, " N"]:
                solver = lx.CG(atol=tol, rtol=tol)
                return lx.linear_solve(obs_cov, step_residual, solver=solver).value

            return jax.vmap(solve_one)(residuals)

        def analyse_one(
            window_obs: Float[Array, "T_plus_1 N"],
            window_mask: Float[Array, "T_plus_1 N"],
        ) -> Float[Array, " N"]:
            # G(x₀) = [mₜ ⊙ Hₜ(Mₜ(x₀))]ₜ, shared with the posterior adapters
            observed = observed_trajectory(self, window_mask)
            masked_observations = window_mask * window_obs

            def residual(
                initial_state: Float[Array, " N"],
            ) -> Float[Array, "T_plus_1 N"]:
                # rₜ = mₜ ⊙ (yₜ - Hₜ(Mₜ(x₀)))
                return masked_observations - observed(initial_state)

            return self._gauss_newton(residual, apply_obs_precision)

        return jax.vmap(analyse_one)(batch.input, batch.mask)

    def _analyse_points(self, batch: PointObsBatch, n_steps: int) -> Float[Array, " N"]:
        r"""GN outer + CG inner on a point batch.

        With the whitened, padding-masked residual
        $r_i(x_0) = m_i (y_i - \mathcal H(M^{\tau_i}(x_0))(s_i)) / \sigma_i$,
        the cost is
        $\tfrac12 \lVert x_0 - x_b \rVert^2_{B^{-1}} + \tfrac12 \lVert r \rVert^2$
        and the Gauss-Newton Hessian is $B^{-1} + r'^\top r'$
        ($I + S^\top r'^\top r' S$ with the control-variable transform).
        """
        obs_op = require_grid_point_obs(self.obs_op)

        def residual(initial_state: Float[Array, " N"]) -> Float[Array, " n_obs"]:
            predictions = obs_op(self._rollout(initial_state, n_steps=n_steps), batch)
            return whitened_residual(predictions, batch)

        def identity(residuals: Float[Array, " n_obs"]) -> Float[Array, " n_obs"]:
            # the residual is already whitened: R = I
            return residuals

        return self._gauss_newton(residual, identity)

    def _gauss_newton(
        self,
        residual: Callable[[Array], Array],
        apply_obs_precision: Callable[[Array], Array],
    ) -> Float[Array, " N"]:
        r"""Gauss–Newton outer loop with a CG inner solve, for one window.

        J(x₀) = ½ ‖x₀ − x_b‖²_{B⁻¹} + ½ r(x₀)ᵀ W r(x₀),  W = R⁻¹.

        Algorithm:
            1. Pick the control c and its map to the state: c = x₀
               (state space), or c = χ with x₀ = x_b + S χ, S Sᵀ = B
               (control-variable transform), where the background term
               becomes ½ ‖χ‖².
            2. At the current control c⁽ᵏ⁾, linearise the residual,
               r(c⁽ᵏ⁾ + δc) ≈ r(c⁽ᵏ⁾) + r' δc.
            3. Solve (P + r'ᵀ W r') δc = −∇J(c⁽ᵏ⁾) by CG, with P = B⁻¹ in
               state space and P = I under the transform.
            4. Update c⁽ᵏ⁺¹⁾ = c⁽ᵏ⁾ + δc and repeat ``n_outer`` times.

        Args:
            residual: x₀ ↦ r(x₀), masked or whitened observation misfits.
            apply_obs_precision: r ↦ W r.

        Returns:
            The analysis x₀*.
        """
        config = self.config
        background = self.prior_mean
        tol = config.covariance_tol

        if config.cvt:
            prior_sqrt = self._prior_sqrt()

            def to_state(control: Float[Array, " N"]) -> Float[Array, " N"]:
                # x₀ = x_b + S χ
                return background + prior_sqrt.mv(control)

            def background_cost(control: Float[Array, " N"]) -> Float[Array, ""]:
                # ½ ‖χ‖²
                return 0.5 * jnp.sum(control**2)

            def apply_prior_precision(direction: Float[Array, " N"]):
                # the prior Hessian in χ is the identity
                return direction

            control = jnp.zeros_like(background)
        else:

            def to_state(control: Float[Array, " N"]) -> Float[Array, " N"]:
                return control

            prior_cov = _as_covariance(self.prior_cov_op)

            def apply_prior_precision(direction: Float[Array, " N"]):
                # B⁻¹ v
                solver = lx.CG(atol=tol, rtol=tol)
                return lx.linear_solve(prior_cov, direction, solver=solver).value

            def background_cost(control: Float[Array, " N"]) -> Float[Array, ""]:
                # ½ (x₀ - x_b)ᵀ B⁻¹ (x₀ - x_b)
                departure = control - background
                return 0.5 * jnp.sum(departure * apply_prior_precision(departure))

            control = background

        def control_residual(control: Float[Array, " N"]) -> Array:
            return residual(to_state(control))

        def cost(control: Float[Array, " N"]) -> Float[Array, ""]:
            # J = background + ½ rᵀ W r
            misfit = control_residual(control)
            weighted = apply_obs_precision(misfit)
            return background_cost(control) + 0.5 * jnp.sum(misfit * weighted)

        structure = jax.ShapeDtypeStruct(control.shape, control.dtype)
        inner_solver = lx.CG(
            atol=config.cg_atol, rtol=config.cg_rtol, max_steps=config.n_inner
        )
        for _ in range(config.n_outer):
            # --- linearise at c⁽ᵏ⁾ ---
            _, residual_tangent = jax.linearize(control_residual, control)
            _, residual_adjoint = jax.vjp(control_residual, control)

            def hessian_mv(
                direction: Float[Array, " N"],
                tangent: Callable = residual_tangent,
                adjoint: Callable = residual_adjoint,
            ) -> Float[Array, " N"]:
                # (P + r'ᵀ W r') v
                weighted = apply_obs_precision(tangent(direction))
                (data_term,) = adjoint(weighted)
                return apply_prior_precision(direction) + data_term

            hessian = lx.FunctionLinearOperator(
                hessian_mv, structure, lx.positive_semidefinite_tag
            )

            # --- inner solve: (P + r'ᵀ W r') δc = -∇J(c⁽ᵏ⁾) ---
            negative_gradient = -jax.grad(cost)(control)
            # truncated CG: stop at n_inner whether or not it has converged
            solution = lx.linear_solve(
                hessian, negative_gradient, solver=inner_solver, throw=False
            )
            # c⁽ᵏ⁺¹⁾ = c⁽ᵏ⁾ + δc
            control = control + solution.value

        return to_state(control)

    def _prior_sqrt(self) -> lx.AbstractLinearOperator:
        """S with S Sᵀ = B: the given ``prior_sqrt_op``, else ``gaussx.sqrt(B)``."""
        if self.prior_sqrt_op is not None:
            return self.prior_sqrt_op
        return gaussx.sqrt(_as_covariance(self.prior_cov_op))

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch.
        """
        from vardax._src.posterior._precision import model_posterior

        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _IncrementalFourDVarAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep``.

        Returns:
            An analysis step that treats non-finite observations as
            missing and returns the analysed initial state.
        """
        return _IncrementalFourDVarAnalysisStep(self)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\).

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch.

Source code in src/vardax/_src/models/incremental_fourdvar.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch.
    """
    from vardax._src.posterior._precision import model_posterior

    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _IncrementalFourDVarAnalysisStep

Adapt to pipekit_cycle.AnalysisStep.

Returns:

Type	Description
_IncrementalFourDVarAnalysisStep	An analysis step that treats non-finite observations as
_IncrementalFourDVarAnalysisStep	missing and returns the analysed initial state.

Source code in src/vardax/_src/models/incremental_fourdvar.py

def as_analysis_step(self) -> _IncrementalFourDVarAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep``.

    Returns:
        An analysis step that treats non-finite observations as
        missing and returns the analysed initial state.
    """
    return _IncrementalFourDVarAnalysisStep(self)

IncrementalConfig

Bases: Module

Configuration for IncrementalFourDVar.

A plain container of the outer/inner iteration budget and the control-variable-transform switch.

Attributes:

Name	Type	Description
n_outer	int	Number of Gauss-Newton outer iterations (typical: 3).
n_inner	int	Max CG iterations per outer (typical: 20-50). The inner solve stops there whether or not it has met the tolerance (truncated CG, as in operational incremental 4DVar).
cg_atol	float	Inner CG absolute tolerance.
cg_rtol	float	Inner CG relative tolerance.
cvt	bool	Solve the inner system over the control variable \(\chi\) with \(x_0 = x_b + B^{1/2} \chi\) (default). The background term is \(\tfrac12 \lVert \chi \rVert^2\), the inner Hessian \(I + B^{1/2\top} G'^\top R^{-1} G' B^{1/2}\) has eigenvalues ≥ 1, and \(B\) is never solved against. False solves in state space with \(B^{-1}\) applied by CG at every Hessian product.
covariance_tol	float	Absolute and relative tolerance of the CG solves that apply \(R^{-1}\) (and \(B^{-1}\) when cvt=False). Outer accuracy cannot beat it.

Examples:

The defaults are three outer and thirty inner iterations with the control-variable transform on:

>>> import vardax as vdx
>>> cfg = vdx.IncrementalConfig()
>>> cfg.n_outer, cfg.n_inner, cfg.cvt
(3, 30, True)

State-space solving (cvt=False) with a bigger budget:

>>> vdx.IncrementalConfig(n_outer=5, cvt=False).n_outer
5

Source code in src/vardax/_src/models/incremental_fourdvar.py

class IncrementalConfig(eqx.Module):
    r"""Configuration for ``IncrementalFourDVar``.

    A plain container of the outer/inner iteration budget and the
    control-variable-transform switch.

    Attributes:
        n_outer: Number of Gauss-Newton outer iterations (typical: 3).
        n_inner: Max CG iterations per outer (typical: 20-50). The inner
            solve stops there whether or not it has met the tolerance
            (truncated CG, as in operational incremental 4DVar).
        cg_atol: Inner CG absolute tolerance.
        cg_rtol: Inner CG relative tolerance.
        cvt: Solve the inner system over the control variable $\chi$ with
            $x_0 = x_b + B^{1/2} \chi$ (default). The background term is
            $\tfrac12 \lVert \chi \rVert^2$, the inner Hessian
            $I + B^{1/2\top} G'^\top R^{-1} G' B^{1/2}$ has eigenvalues
            ≥ 1, and $B$ is never solved against. ``False`` solves in
            state space with $B^{-1}$ applied by CG at every Hessian
            product.
        covariance_tol: Absolute and relative tolerance of the CG solves
            that apply $R^{-1}$ (and $B^{-1}$ when ``cvt=False``). Outer
            accuracy cannot beat it.

    Examples:
        The defaults are three outer and thirty inner iterations with the
        control-variable transform on:

        >>> import vardax as vdx
        >>> cfg = vdx.IncrementalConfig()
        >>> cfg.n_outer, cfg.n_inner, cfg.cvt
        (3, 30, True)

        State-space solving (``cvt=False``) with a bigger budget:

        >>> vdx.IncrementalConfig(n_outer=5, cvt=False).n_outer
        5
    """

    n_outer: int = eqx.field(static=True, default=3)
    n_inner: int = eqx.field(static=True, default=30)
    cg_atol: float = eqx.field(static=True, default=1e-6)
    cg_rtol: float = eqx.field(static=True, default=1e-6)
    cvt: bool = eqx.field(static=True, default=True)
    covariance_tol: float = eqx.field(static=True, default=1e-6)

    def get_config(self) -> dict[str, Any]:
        return {
            "n_outer": self.n_outer,
            "n_inner": self.n_inner,
            "cg_atol": self.cg_atol,
            "cg_rtol": self.cg_rtol,
            "cvt": self.cvt,
            "covariance_tol": self.covariance_tol,
        }

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> IncrementalConfig:
        return cls(**config)

ReducedBasisFourDVar

Bases: Module

Reduced-basis 4DVar.

\[J(X) = \tfrac12 X^\top Q^{-1} X + J_o(x_b + \Phi X)\]

\[J_o(x_0) = \tfrac12 \sum_t \|m_t \odot (y_t - H(M^t(x_0)))\|^2_{R^{-1}}\]

Without a basis the control is \(x_0\) itself and the first term is \(\tfrac12 \|x_0 - x_b\|^2_{B^{-1}}\).

Algorithm

1. Pick the control and its map to the initial state: \(x_0\) itself, the coefficients \(X\) with \(x_0 = x_b + \Phi X\) (basis), or \(\theta\) with \(x_0 = \psi(\theta)\) (parameterization).
2. Roll the forward model out from \(x_0\) for the window length.
3. Add the prior term (\(B^{-1}\), basis.prior_inv or the parameterization penalty) to the masked observation misfit summed over the trajectory. Masked entries are marginalised out of \(R\), so a correlated \(R\) is not biased by the gaps.
4. Minimise the full nonlinear cost with minimiser (no linearisation or inner loop) and decode the minimiser to \(x_0^*\).

Attributes:

Name	Type	Description
forward	Any	pipekit_cycle.ForwardModel supplying step(state, dt) -> state.
obs_op	Any	Observation operator (linear or nonlinear). Supplies __call__ and linearize; __call__ may accept a mask kwarg (MaskedIdentity / InterpObs style).
prior_mean	Float[Array, ' N']	Background \(x_b\).
prior_cov_op	AbstractLinearOperator	\(B\). Used only when basis is None.
obs_cov_op	AbstractLinearOperator | None	\(R\).
basis	Any	Optional ReducedBasis. When supplied, the control is the coefficient vector \(X\) (basis.nbasis entries) and \(B^{-1}\) is replaced by basis.prior_inv — the same as parameterization=LinearParameterization(basis, prior_mean).
parameterization	StateParameterization | None	Optional StateParameterization; the control is \(\theta\), \(x_0 = \psi(\theta)\) and \(\mathrm{pen}(\theta)\) replaces the background term. Exclusive with basis.
minimiser	AbstractMinimiser | Minimiser	optimistix.AbstractMinimiser for the outer optimisation (default optx.BFGS(rtol=1e-6, atol=1e-6)), or a Minimiser such as OptaxMinimiser.
minimiser_adjoint	AbstractAdjoint	optimistix.AbstractAdjoint for differentiating through the minimum. Default optx.ImplicitAdjoint.
max_steps	int	Iteration cap on the outer solver when minimiser is an optimistix solver; a vardax Minimiser keeps its own.
cov_solver	AbstractLinearSolver	lineax solver for the covariance solves in the cost (\(B^{-1}\), \(R^{-1}\)). Default lx.CG(atol=1e-6, rtol=1e-6), which suits large matrix-free operators but is inexact on an ill-conditioned covariance (for example a tiny \(R\)); pass lx.AutoLinearSolver(well_posed=True) or lx.Cholesky() for an exact solve when the operator can be materialised.
forward_adjoint	Any	As in StrongFourDVar: a diffrax.AbstractAdjoint for a forward implementing SupportsForwardAdjoint.
checkpoint	bool | int	As in StrongFourDVar: False, True or a block length k for jax.checkpoint over the rollout.

Examples:

Trivial dynamics, one timestep, \(B = R = I\) and no basis: the cost is that of 3DVar, so \(x_0^* = y / 2\).

>>> import jax, jax.numpy as jnp, lineax as lx, vardax
>>> class Identity:
...     dt = 1.0
...
...     def step(self, x, dt):
...         return x
>>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((4,), jnp.float32))
>>> common = dict(
...     forward=Identity(),
...     obs_op=vardax.MaskedIdentity(),
...     prior_mean=jnp.zeros(4),
...     prior_cov_op=eye,
...     obs_cov_op=eye,
... )
>>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 4)), mask=jnp.ones((1, 1, 4)))
>>> xa = vardax.ReducedBasisFourDVar(**common)(batch)
>>> xa.shape
(1, 4)
>>> bool(jnp.allclose(xa, 0.5, atol=1e-3))
True

With a basis the increment is confined to \(\mathrm{span}(\Phi)\). A single unit-variance constant mode gives \(X^* = \tfrac{\sum_i y_i}{1 + N}\) and \(x_0^* = X^* \mathbf{1}\):

>>> basis = vardax.linear_basis(jnp.ones((4, 1)))
>>> xa = vardax.ReducedBasisFourDVar(**common, basis=basis)(batch)
>>> bool(jnp.allclose(xa, 4.0 / 5.0, atol=1e-3))
True

Composed with a Fourier basis from fourier_basis, the analysis is still a state-space array:

>>> fb = vardax.fourier_basis([jnp.linspace(0.0, 1.0, 4)], 2)
>>> xb = vardax.ReducedBasisFourDVar(**common, basis=fb)(batch)
>>> xb.shape
(1, 4)

basis and parameterization are exclusive:

>>> param = vardax.LinearParameterization(basis, jnp.zeros(4))
>>> vardax.ReducedBasisFourDVar(
...     **common, basis=basis, parameterization=param
... )
Traceback (most recent call last):
...
ConfigurationError: Pass basis= or parameterization=, not both; ...

Source code in src/vardax/_src/models/reduced_basis_fourdvar.py

class ReducedBasisFourDVar(eqx.Module):
    r"""Reduced-basis 4DVar.

    $$J(X) = \tfrac12 X^\top Q^{-1} X + J_o(x_b + \Phi X)$$

    $$J_o(x_0) = \tfrac12 \sum_t \|m_t \odot (y_t - H(M^t(x_0)))\|^2_{R^{-1}}$$

    Without a basis the control is $x_0$ itself and the first term is
    $\tfrac12 \|x_0 - x_b\|^2_{B^{-1}}$.

    Algorithm:
        1. Pick the control and its map to the initial state: $x_0$ itself,
           the coefficients $X$ with $x_0 = x_b + \Phi X$ (``basis``), or
           $\theta$ with $x_0 = \psi(\theta)$ (``parameterization``).
        2. Roll the forward model out from $x_0$ for the window length.
        3. Add the prior term ($B^{-1}$, ``basis.prior_inv`` or the
           parameterization penalty) to the masked observation misfit
           summed over the trajectory. Masked entries are marginalised out
           of $R$, so a correlated $R$ is not biased by the gaps.
        4. Minimise the full nonlinear cost with ``minimiser`` (no
           linearisation or inner loop) and decode the minimiser to
           $x_0^*$.

    Attributes:
        forward: ``pipekit_cycle.ForwardModel`` supplying
            ``step(state, dt) -> state``.
        obs_op: Observation operator (linear or nonlinear). Supplies
            ``__call__`` and ``linearize``; ``__call__`` may accept a
            ``mask`` kwarg (`MaskedIdentity` / `InterpObs` style).
        prior_mean: Background $x_b$.
        prior_cov_op: $B$. Used only when ``basis is None``.
        obs_cov_op: $R$.
        basis: Optional `ReducedBasis`. When supplied, the control is
            the coefficient vector $X$ (``basis.nbasis`` entries)
            and $B^{-1}$ is replaced by ``basis.prior_inv`` — the same as
            ``parameterization=LinearParameterization(basis, prior_mean)``.
        parameterization: Optional
            [`StateParameterization`][vardax.StateParameterization]; the
            control is $\theta$, $x_0 = \psi(\theta)$ and
            $\mathrm{pen}(\theta)$ replaces the background term. Exclusive
            with ``basis``.
        minimiser: ``optimistix.AbstractMinimiser`` for the outer
            optimisation (default ``optx.BFGS(rtol=1e-6, atol=1e-6)``), or
            a [`Minimiser`][vardax.Minimiser] such as
            [`OptaxMinimiser`][vardax.OptaxMinimiser].
        minimiser_adjoint: ``optimistix.AbstractAdjoint`` for
            differentiating through the minimum. Default
            ``optx.ImplicitAdjoint``.
        max_steps: Iteration cap on the outer solver when ``minimiser``
            is an ``optimistix`` solver; a vardax `Minimiser` keeps its own.
        cov_solver: ``lineax`` solver for the covariance solves in the
            cost ($B^{-1}$, $R^{-1}$). Default
            ``lx.CG(atol=1e-6, rtol=1e-6)``, which suits large
            matrix-free operators but is inexact on an ill-conditioned
            covariance (for example a tiny $R$); pass
            ``lx.AutoLinearSolver(well_posed=True)`` or ``lx.Cholesky()``
            for an exact solve when the operator can be materialised.
        forward_adjoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            a ``diffrax.AbstractAdjoint`` for a forward implementing
            [`SupportsForwardAdjoint`][vardax.SupportsForwardAdjoint].
        checkpoint: As in [`StrongFourDVar`][vardax.StrongFourDVar]:
            ``False``, ``True`` or a block length ``k`` for
            ``jax.checkpoint`` over the rollout.

    Examples:
        Trivial dynamics, one timestep, $B = R = I$ and no basis: the
        cost is that of 3DVar, so $x_0^* = y / 2$.

        >>> import jax, jax.numpy as jnp, lineax as lx, vardax
        >>> class Identity:
        ...     dt = 1.0
        ...
        ...     def step(self, x, dt):
        ...         return x
        >>> eye = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((4,), jnp.float32))
        >>> common = dict(
        ...     forward=Identity(),
        ...     obs_op=vardax.MaskedIdentity(),
        ...     prior_mean=jnp.zeros(4),
        ...     prior_cov_op=eye,
        ...     obs_cov_op=eye,
        ... )
        >>> batch = vardax.Batch1D(input=jnp.ones((1, 1, 4)), mask=jnp.ones((1, 1, 4)))
        >>> xa = vardax.ReducedBasisFourDVar(**common)(batch)
        >>> xa.shape
        (1, 4)
        >>> bool(jnp.allclose(xa, 0.5, atol=1e-3))
        True

        With a basis the increment is confined to $\mathrm{span}(\Phi)$.
        A single unit-variance constant mode gives
        $X^* = \tfrac{\sum_i y_i}{1 + N}$ and $x_0^* = X^* \mathbf{1}$:

        >>> basis = vardax.linear_basis(jnp.ones((4, 1)))
        >>> xa = vardax.ReducedBasisFourDVar(**common, basis=basis)(batch)
        >>> bool(jnp.allclose(xa, 4.0 / 5.0, atol=1e-3))
        True

        Composed with a Fourier basis from [`fourier_basis`][vardax.fourier_basis],
        the analysis is still a state-space array:

        >>> fb = vardax.fourier_basis([jnp.linspace(0.0, 1.0, 4)], 2)
        >>> xb = vardax.ReducedBasisFourDVar(**common, basis=fb)(batch)
        >>> xb.shape
        (1, 4)

        ``basis`` and ``parameterization`` are exclusive:

        >>> param = vardax.LinearParameterization(basis, jnp.zeros(4))
        >>> vardax.ReducedBasisFourDVar(
        ...     **common, basis=basis, parameterization=param
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: Pass basis= or parameterization=, not both; ...
    """

    forward: Any
    obs_op: Any
    prior_mean: Float[Array, " N"]
    prior_cov_op: lx.AbstractLinearOperator
    obs_cov_op: lx.AbstractLinearOperator | None
    basis: Any
    parameterization: StateParameterization | None
    minimiser: optx.AbstractMinimiser | Minimiser = eqx.field(static=True)
    minimiser_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    cov_solver: lx.AbstractLinearSolver = eqx.field(static=True)
    max_steps: int = eqx.field(static=True, default=200)
    forward_adjoint: Any = eqx.field(static=True, default=None)
    checkpoint: bool | int = eqx.field(static=True, default=False)

    def __init__(
        self,
        forward: Any,
        obs_op: Any,
        prior_mean: Float[Array, " N"],
        prior_cov_op: lx.AbstractLinearOperator,
        obs_cov_op: lx.AbstractLinearOperator | None,
        *,
        basis: Any = None,
        parameterization: StateParameterization | None = None,
        minimiser: optx.AbstractMinimiser | Minimiser | None = None,
        minimiser_adjoint: optx.AbstractAdjoint | None = None,
        max_steps: int = 200,
        forward_adjoint: Any = None,
        checkpoint: bool | int = False,
        cov_solver: lx.AbstractLinearSolver | None = None,
    ) -> None:
        # fail at construction, not on the first traced call
        checkpoint_every(checkpoint)
        resolve_forward(forward, forward_adjoint)
        self.forward = forward
        self.obs_op = obs_op
        self.prior_mean = prior_mean
        self.prior_cov_op = prior_cov_op
        self.obs_cov_op = obs_cov_op
        if basis is not None and parameterization is not None:
            raise ConfigurationError(
                "Pass basis= or parameterization=, not both; "
                "LinearParameterization(basis, prior_mean) is the basis path."
            )
        self.basis = basis
        self.parameterization = parameterization
        self.minimiser = minimiser or optx.BFGS(rtol=1e-6, atol=1e-6)
        self.minimiser_adjoint = minimiser_adjoint or optx.ImplicitAdjoint()
        self.max_steps = max_steps
        self.cov_solver = cov_solver or lx.CG(atol=1e-6, rtol=1e-6)
        self.forward_adjoint = forward_adjoint
        self.checkpoint = checkpoint

    def _rollout(
        self,
        initial_state: Float[Array, " N"],
        n_steps: int,
    ) -> Float[Array, "T_plus_1 N"]:
        """Step the forward ``n_steps`` times, returning the ``T+1``-long trajectory."""
        forward = resolve_forward(self.forward, self.forward_adjoint)
        return rollout(forward, initial_state, n_steps, checkpoint=self.checkpoint)

    def __call__(self, batch: Batch1D | PointObsBatch) -> Float[Array, "*batch N"]:
        r"""Reduced-basis 4DVar analysis: minimise over ``x_0`` (or ``X``).

        A [`PointObsBatch`][vardax.PointObsBatch] needs
        ``obs_op=GridPointObs(...)``: observation $i$ is compared with the
        trajectory state at ``time_index[i]`` over a window of ``n_times``
        states, with error variance $\sigma_i^2$ (``obs_cov_op`` must
        be ``None`` or an identity).

        Args:
            batch: A gridded [`Batch1D`][vardax.Batch1D] whose time axis
                holds the window, or a
                [`PointObsBatch`][vardax.PointObsBatch].

        Returns:
            The analysed initial state $x_0^*$ of every window.

        Raises:
            ObservationError: For a point batch with a non-trivial
                ``obs_cov_op`` or an out-of-range ``time_index``.
            MissingInputError: For a gridded batch when ``obs_cov_op`` is
                ``None``, or a point batch with ``time_index`` but no
                ``n_times``.
            InputTypeError: For a point batch whose ``obs_op`` is not a
                ``GridPointObs``.
        """
        if isinstance(batch, PointObsBatch):
            check_point_obs_cov(self.obs_cov_op)
            n_steps = window_steps(batch)

            def _one_points(window: PointObsBatch) -> Float[Array, " N"]:
                def _obs_cost(
                    trajectory: Float[Array, "T_plus_1 N"],
                ) -> Float[Array, ""]:
                    return grid_point_obs_cost(self.obs_op, trajectory, window)

                return self._solve(_obs_cost, n_steps)

            return map_point_batch(_one_points, batch)

        obs_cov_op = require_obs_cov(self.obs_cov_op)
        n_steps = batch.input.shape[1] - 1

        def _one(
            window_obs: Float[Array, "T_plus_1 N"],
            window_mask: Float[Array, "T_plus_1 N"],
        ) -> Float[Array, " N"]:
            def _obs_cost(
                trajectory: Float[Array, "T_plus_1 N"],
            ) -> Float[Array, ""]:
                def _per_step(state, observation, step_mask):
                    predicted_obs = (
                        self.obs_op(state, mask=step_mask)
                        if _accepts_mask(self.obs_op)
                        else self.obs_op(state)
                    )
                    # r = m ⊙ (y − H(x))
                    residual = step_mask * (observation - predicted_obs)
                    # Solve against the marginalized covariance
                    # R~ = M R M + (I - M): for a 0/1 mask its inverse on
                    # the observed block equals (R_oo)^{-1}, so masked
                    # entries drop out of the likelihood exactly. Solving
                    # the full R against a zeroed residual would instead
                    # weight observed entries by the observed block of
                    # R^{-1}, biasing the analysis when R is correlated.
                    mask_op = lx.DiagonalLinearOperator(step_mask)
                    masked_obs_cov = lx.TaggedLinearOperator(
                        mask_op @ obs_cov_op @ mask_op
                        + lx.DiagonalLinearOperator(1.0 - step_mask),
                        lx.positive_semidefinite_tag,
                    )
                    # R̃⁻¹r
                    weighted_residual = lx.linear_solve(
                        masked_obs_cov,
                        residual,
                        solver=self.cov_solver,
                    ).value
                    # J_o,t = ½ rᵀ R̃⁻¹ r
                    return 0.5 * jnp.sum(residual * weighted_residual)

                per_step_costs = jax.vmap(_per_step)(
                    trajectory, window_obs, window_mask
                )
                return jnp.sum(per_step_costs)

            return self._solve(_obs_cost, n_steps)

        return jax.vmap(_one)(batch.input, batch.mask)

    def _solve(
        self, _obs_cost: Callable[[Array], Array], n_steps: int
    ) -> Float[Array, " N"]:
        """Minimise background (or penalty) plus the trajectory ``_obs_cost``."""
        param = self.parameterization
        if self.basis is not None:
            # The reduced-basis path, unchanged; it equals
            # LinearParameterization(basis, prior_mean) mathematically.
            def cost(coefficients: Float[Array, " M"], _args: Any) -> Float[Array, ""]:
                # J_b = ½ Xᵀ P⁻¹ X
                background_cost = 0.5 * jnp.sum(
                    coefficients * self.basis.prior_inv(coefficients)
                )
                basis_increment = self.basis.operg(0.0, coefficients)
                # x₀ = x_b + Φ X
                initial_state = self.prior_mean + as_shape(
                    basis_increment, self.prior_mean.shape
                )
                return background_cost + _obs_cost(
                    self._rollout(initial_state, n_steps=n_steps)
                )

            initial_control = jnp.zeros(
                (self.basis.nbasis,), dtype=self.prior_mean.dtype
            )
        elif param is not None:

            def cost(theta: Any, _args: Any) -> Float[Array, ""]:
                initial_state = as_shape(param.decode(theta), self.prior_mean.shape)
                return param.penalty(theta) + _obs_cost(
                    self._rollout(initial_state, n_steps=n_steps)
                )

            initial_control = param.init_control()
        else:

            def cost(initial_state: Float[Array, " N"], _args: Any) -> Float[Array, ""]:
                # δx = x₀ − x_b
                increment = initial_state - self.prior_mean
                # B⁻¹δx
                weighted_increment = lx.linear_solve(
                    self.prior_cov_op,
                    increment,
                    solver=self.cov_solver,
                ).value
                # J_b = ½ δxᵀ B⁻¹ δx
                background_cost = 0.5 * jnp.sum(increment * weighted_increment)
                return background_cost + _obs_cost(
                    self._rollout(initial_state, n_steps=n_steps)
                )

            initial_control = self.prior_mean

        control_star = minimise(
            self.minimiser,
            cost,
            initial_control,
            adjoint=self.minimiser_adjoint,
            max_steps=self.max_steps,
        )
        if self.basis is not None:
            increment = self.basis.operg(0.0, control_star)
            # x₀* = x_b + Φ X*
            return self.prior_mean + as_shape(increment, self.prior_mean.shape)
        if param is not None:
            return as_shape(param.decode(control_star), self.prior_mean.shape)
        return control_star

    def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
        """Analyse one window and return its [`Posterior`][vardax.Posterior].

        Runs the analysis, then ``adapter(analysis, self, batch)``; the
        default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
        with this model's $B$ and $R$.

        Args:
            batch: A ``Batch1D`` holding one window.
            adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
                e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

        Returns:
            The window's posterior: mean and lazy covariance.

        Raises:
            ShapeError: If ``batch`` holds more than one window.
            UnsupportedInputError: For a point batch, or for a model with
                ``basis=`` or ``parameterization=``, whose prior is not
                $B$.
        """
        from vardax._src.posterior._precision import model_posterior

        return model_posterior(self, batch, adapter)

    def as_analysis_step(self) -> _ReducedBasisFourDVarAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            An analysis step that accepts single arrays or windowed lists,
            swaps in the supplied forecast and operators, and returns the
            analysis (rolled to the window end for a smoother).
        """
        return _ReducedBasisFourDVarAnalysisStep(self)

posterior

posterior(batch: Batch1D, adapter: Any = None) -> Any

Analyse one window and return its Posterior.

Runs the analysis, then adapter(analysis, self, batch); the default adapter is LaplaceCovariance with this model's \(B\) and \(R\).

Parameters:

Name	Type	Description	Default
batch	Batch1D	A Batch1D holding one window.	required
adapter	Any	Optional PosteriorAdapter, e.g. GaussNewtonHessian.	None

Returns:

Type	Description
Any	The window's posterior: mean and lazy covariance.

Raises:

Type	Description
ShapeError	If batch holds more than one window.
UnsupportedInputError	For a point batch, or for a model with basis= or parameterization=, whose prior is not \(B\).

Source code in src/vardax/_src/models/reduced_basis_fourdvar.py

def posterior(self, batch: Batch1D, adapter: Any = None) -> Any:
    """Analyse one window and return its [`Posterior`][vardax.Posterior].

    Runs the analysis, then ``adapter(analysis, self, batch)``; the
    default adapter is [`LaplaceCovariance`][vardax.LaplaceCovariance]
    with this model's $B$ and $R$.

    Args:
        batch: A ``Batch1D`` holding one window.
        adapter: Optional [`PosteriorAdapter`][vardax.PosteriorAdapter],
            e.g. [`GaussNewtonHessian`][vardax.GaussNewtonHessian].

    Returns:
        The window's posterior: mean and lazy covariance.

    Raises:
        ShapeError: If ``batch`` holds more than one window.
        UnsupportedInputError: For a point batch, or for a model with
            ``basis=`` or ``parameterization=``, whose prior is not
            $B$.
    """
    from vardax._src.posterior._precision import model_posterior

    return model_posterior(self, batch, adapter)

as_analysis_step

as_analysis_step() -> _ReducedBasisFourDVarAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_ReducedBasisFourDVarAnalysisStep	An analysis step that accepts single arrays or windowed lists,
_ReducedBasisFourDVarAnalysisStep	swaps in the supplied forecast and operators, and returns the
_ReducedBasisFourDVarAnalysisStep	analysis (rolled to the window end for a smoother).

Source code in src/vardax/_src/models/reduced_basis_fourdvar.py

def as_analysis_step(self) -> _ReducedBasisFourDVarAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        An analysis step that accepts single arrays or windowed lists,
        swaps in the supplied forecast and operators, and returns the
        analysis (rolled to the window end for a smoother).
    """
    return _ReducedBasisFourDVarAnalysisStep(self)

Neural-field 4DVar

SpaceTimeFieldVar fits one continuous field \(u(t, s) = f_\theta(t, s)\) over the whole window to point observations, with an optional PDE residual as a soft constraint and weight decay in place of \(B\); no rollout. It returns a SpaceTimeFieldResult, queryable anywhere, and ensemble(...) gives a deep-ensemble spread.

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

SpaceTimeFieldVar

Bases: Module

Neural-field 4D-Var: fit one continuous field \(u(t, s) = f_\theta(t, s)\).

Minimises, over the network weights \(\theta\),

\[J(\theta) = J_o + J_\mathcal{N} + \tfrac{\alpha}{2} \lVert\theta\rVert^2\]

\[J_o(\theta) = \tfrac12 \sum_k m_k (y_k - f_\theta(t_k, s_k))^2 / \sigma_k^2\]

\[J_\mathcal{N} = \tfrac{\lambda}{2 N_c} \sum_j \| \mathcal N[f_\theta](c_j) \|^2\]

with the observations of a PointObsBatch (\(m_k = 0\) on padding rows, which contribute nothing, NaN included), the physics term of an optional ResidualPenalty and weight decay \(\alpha\). There is no background covariance and no rollout: the whole window is the field, which can be queried and differentiated anywhere (Oh 2025).

Coordinates put time first: the field maps a point (t, *s) of shape (1 + D,) to (K,) (or a scalar), the observation coordinates are (N_obs, 1 + D) in the same order, and so are the collocation points, as the residuals of vardax.examples expect. The observed quantity is the field's single output (FieldPointObs); a multi-channel field needs its own wrapper.

Scale of \(\lambda\): the observation term is a sum over the \(N_\text{obs}\) valid observations, each weighted by \(1/\sigma_k^2\), while the residual term is a mean over the collocation points. A unit change of the field everywhere moves the first by about \(N_\text{obs}/(2\sigma^2)\) and the second by about \(\lambda/2\) times the squared residual it causes, so \(\lambda\) trades off against \(N_\text{obs}/(2\sigma^2)\) and must grow with it to keep the same balance (the advection test uses \(\lambda = 10^3\) for 100 observations with \(\sigma = 0.05\), where \(N_\text{obs}/(2\sigma^2) = 2 \times 10^4\)).

The fit runs in one eqx.filter_jit-compiled function of the model, the observations and the key, so repeated analyses with new observations (or a new weight_decay or penalty weight, both array leaves) of the same shapes do not recompile.

The penalty's collocation points are fixed during a fit: a ResidualPenalty with a sampler draws them once, from key, at the start of the analysis, so the cost the minimiser sees is deterministic (its patience test, best-iterate tracking and implicit adjoint all assume that). Draw more points, or refit with new keys, for a better estimate of the expected residual.

Algorithm

1. Split the field's leaves into the control \(\theta\) (by filter_spec) and the frozen rest.
2. Draw the penalty's collocation points once from key when it has a sampler.
3. Evaluate the field at the observation coordinates and form the masked, \(\sigma\)-weighted observation cost.
4. Add the weight decay \(\tfrac{\alpha}{2}\lVert\theta\rVert^2\) and, if present, the mean squared PDE residual at the collocation points.
5. Minimise over \(\theta\) with minimiser in one compiled function and return the fitted field with the cost, gradient norm and step count.

Attributes:

Name	Type	Description
field	Any	The initial network, a single-example eqx.Module (1 + D,) -> (K,) such as geonnax.SIREN or geonnax.NeuralFieldMLP. Its parameters are the initial control.
residual	ResidualPenalty | None	Optional ResidualPenalty (the physics term); None fits the observations alone.
weight_decay	Float[Array, '']	\(\alpha \ge 0\), stored as a 0-d array (a leaf: it can be changed without recompiling and differentiated).
minimiser	OptaxMinimiser | OptimistixMinimiser	An OptaxMinimiser (the default, Adam with learning rate \(10^{-3}\)) or an OptimistixMinimiser.
filter_spec	Any	Which leaves of field are the control, as in NeuralFieldParameterization: exclude frozen arrays (random Fourier frequencies) so the weight decay does not shrink them. Default eqx.is_inexact_array.
t0	float	Start of the window, \(t_0\) (default 0.0).

Examples:

A field linear in fixed features, \(f_w(t, x) = w_0 + w_1 x\), with no physics is Bayesian linear regression: the analysis is the ridge solution, here \(w^\star = (0, 1)\) to within the small weight decay.

>>> import equinox as eqx, jax, jax.numpy as jnp, vardax
>>> class Line(eqx.Module):
...     w: jax.Array
...
...     def __call__(self, s):
...         return (self.w[0] + self.w[1] * s[1])[None]
>>> obs = vardax.PointObsBatch(
...     coords=jnp.array([[0.0, -1.0], [0.5, 0.0], [1.0, 1.0], [0.0, 9.0]]),
...     values=jnp.array([-1.0, 0.0, 1.0, jnp.nan]),  # last row: padding
...     sigma=jnp.ones(4),
...     valid=jnp.array([True, True, True, False]),
... )
>>> ana = vardax.SpaceTimeFieldVar(
...     field=Line(jnp.zeros(2)),
...     weight_decay=1e-6,
...     minimiser=vardax.OptimistixMinimiser(),
... )
>>> result = ana(obs)
>>> result.initial_condition(jnp.array([[0.5], [2.0]])).round(3)
Array([0.5, 2. ], dtype=float32)

A deep ensemble refits from random initialisations; the result's array leaves carry a member axis:

>>> def init(key):
...     return Line(jax.random.normal(key, (2,)))
>>> members = ana.ensemble(obs, init, 3, key=jax.random.PRNGKey(0))
>>> members.cost.shape
(3,)

The cycle adapter returns the analysis on a grid directly:

>>> step = ana.as_analysis_step(grid=jnp.array([[0.5], [2.0]]))
>>> step(None, obs).shape
(2,)

A negative weight decay is rejected:

>>> vardax.SpaceTimeFieldVar(
...     field=Line(jnp.zeros(2)), weight_decay=-1.0
... )
Traceback (most recent call last):
...
ConfigurationError: weight_decay must be >= 0, got -1.0.

References

Oh, Y. (2025). On the effect of neural field reparameterization for 4DVAR. arXiv:2509.21751.

Source code in src/vardax/_src/models/spacetime_field.py

class SpaceTimeFieldVar(eqx.Module):
    r"""Neural-field 4D-Var: fit one continuous field $u(t, s) = f_\theta(t, s)$.

    Minimises, over the network weights $\theta$,

    $$J(\theta) = J_o + J_\mathcal{N} + \tfrac{\alpha}{2} \lVert\theta\rVert^2$$

    $$J_o(\theta) = \tfrac12 \sum_k m_k (y_k - f_\theta(t_k, s_k))^2 / \sigma_k^2$$

    $$J_\mathcal{N} = \tfrac{\lambda}{2 N_c} \sum_j \| \mathcal N[f_\theta](c_j) \|^2$$

    with the observations of a [`PointObsBatch`][vardax.PointObsBatch]
    ($m_k = 0$ on padding rows, which contribute nothing, ``NaN``
    included), the physics term of an optional
    [`ResidualPenalty`][vardax.ResidualPenalty] and weight decay $\alpha$.
    There is no background covariance and no rollout: the whole window
    is the field, which can be queried and differentiated anywhere
    (Oh 2025).

    Coordinates put **time first**: the field maps a point
    ``(t, *s)`` of shape ``(1 + D,)`` to ``(K,)`` (or a scalar), the
    observation coordinates are ``(N_obs, 1 + D)`` in the same order, and
    so are the collocation points, as the residuals of
    `vardax.examples` expect. The observed quantity is the field's single
    output (`FieldPointObs`); a multi-channel field needs its own wrapper.

    Scale of $\lambda$: the observation term is a *sum* over the
    $N_\text{obs}$ valid observations, each weighted by $1/\sigma_k^2$,
    while the residual term is a *mean* over the collocation points. A
    unit change of the field everywhere moves the first by about
    $N_\text{obs}/(2\sigma^2)$ and the second by about $\lambda/2$ times
    the squared residual it causes, so $\lambda$ trades off against
    $N_\text{obs}/(2\sigma^2)$ and must grow with it to keep the same
    balance (the advection test uses $\lambda = 10^3$ for 100
    observations with $\sigma = 0.05$, where
    $N_\text{obs}/(2\sigma^2) = 2 \times 10^4$).

    The fit runs in one ``eqx.filter_jit``-compiled function of the model,
    the observations and the key, so repeated analyses with new
    observations (or a new ``weight_decay`` or penalty weight, both array
    leaves) of the same shapes do not recompile.

    The penalty's collocation points are fixed during a fit: a
    ``ResidualPenalty`` with a ``sampler`` draws them once, from ``key``,
    at the start of the analysis, so the cost the minimiser sees is
    deterministic (its patience test, best-iterate tracking and implicit
    adjoint all assume that). Draw more points, or refit with new keys,
    for a better estimate of the expected residual.

    Algorithm:
        1. Split the field's leaves into the control $\theta$ (by
           ``filter_spec``) and the frozen rest.
        2. Draw the penalty's collocation points once from ``key`` when it
           has a ``sampler``.
        3. Evaluate the field at the observation coordinates and form the
           masked, $\sigma$-weighted observation cost.
        4. Add the weight decay $\tfrac{\alpha}{2}\lVert\theta\rVert^2$ and,
           if present, the mean squared PDE residual at the collocation
           points.
        5. Minimise over $\theta$ with ``minimiser`` in one compiled
           function and return the fitted field with the cost, gradient
           norm and step count.

    Attributes:
        field: The initial network, a single-example ``eqx.Module``
            ``(1 + D,) -> (K,)`` such as `geonnax.SIREN` or
            `geonnax.NeuralFieldMLP`. Its parameters are the initial
            control.
        residual: Optional [`ResidualPenalty`][vardax.ResidualPenalty]
            (the physics term); ``None`` fits the observations alone.
        weight_decay: $\alpha \ge 0$, stored as a 0-d array (a leaf: it
            can be changed without recompiling and differentiated).
        minimiser: An [`OptaxMinimiser`][vardax.OptaxMinimiser] (the
            default, Adam with learning rate $10^{-3}$) or an
            [`OptimistixMinimiser`][vardax.OptimistixMinimiser].
        filter_spec: Which leaves of ``field`` are the control, as in
            [`NeuralFieldParameterization`][vardax.NeuralFieldParameterization]:
            exclude frozen arrays (random Fourier frequencies) so the weight
            decay does not shrink them. Default ``eqx.is_inexact_array``.
        t0: Start of the window, $t_0$ (default ``0.0``).

    Examples:
        A field linear in fixed features, $f_w(t, x) = w_0 + w_1 x$, with no
        physics is Bayesian linear regression: the analysis is the ridge
        solution, here $w^\star = (0, 1)$ to within the small weight decay.

        >>> import equinox as eqx, jax, jax.numpy as jnp, vardax
        >>> class Line(eqx.Module):
        ...     w: jax.Array
        ...
        ...     def __call__(self, s):
        ...         return (self.w[0] + self.w[1] * s[1])[None]
        >>> obs = vardax.PointObsBatch(
        ...     coords=jnp.array([[0.0, -1.0], [0.5, 0.0], [1.0, 1.0], [0.0, 9.0]]),
        ...     values=jnp.array([-1.0, 0.0, 1.0, jnp.nan]),  # last row: padding
        ...     sigma=jnp.ones(4),
        ...     valid=jnp.array([True, True, True, False]),
        ... )
        >>> ana = vardax.SpaceTimeFieldVar(
        ...     field=Line(jnp.zeros(2)),
        ...     weight_decay=1e-6,
        ...     minimiser=vardax.OptimistixMinimiser(),
        ... )
        >>> result = ana(obs)
        >>> result.initial_condition(jnp.array([[0.5], [2.0]])).round(3)
        Array([0.5, 2. ], dtype=float32)

        A deep ensemble refits from random initialisations; the result's
        array leaves carry a member axis:

        >>> def init(key):
        ...     return Line(jax.random.normal(key, (2,)))
        >>> members = ana.ensemble(obs, init, 3, key=jax.random.PRNGKey(0))
        >>> members.cost.shape
        (3,)

        The cycle adapter returns the analysis on a grid directly:

        >>> step = ana.as_analysis_step(grid=jnp.array([[0.5], [2.0]]))
        >>> step(None, obs).shape
        (2,)

        A negative weight decay is rejected:

        >>> vardax.SpaceTimeFieldVar(  # doctest: +IGNORE_EXCEPTION_DETAIL
        ...     field=Line(jnp.zeros(2)), weight_decay=-1.0
        ... )
        Traceback (most recent call last):
        ...
        ConfigurationError: weight_decay must be >= 0, got -1.0.

    References:
        Oh, Y. (2025). On the effect of neural field reparameterization for
        4DVAR. arXiv:2509.21751.
    """

    field: Any
    residual: ResidualPenalty | None
    weight_decay: Float[Array, ""] = eqx.field(converter=jnp.asarray)
    minimiser: OptaxMinimiser | OptimistixMinimiser = eqx.field(static=True)
    filter_spec: Any = eqx.field(static=True)
    t0: float = eqx.field(static=True)

    def __init__(
        self,
        field: Any,
        residual: ResidualPenalty | None = None,
        weight_decay: float | Float[Array, ""] = 0.0,
        minimiser: OptaxMinimiser | OptimistixMinimiser | None = None,
        *,
        filter_spec: Any = eqx.is_inexact_array,
        t0: float = 0.0,
    ) -> None:
        if jnp.ndim(weight_decay) != 0:
            raise ShapeError(
                f"weight_decay must be a scalar; got shape {jnp.shape(weight_decay)}."
            )
        try:
            negative = float(weight_decay) < 0
        except jax.errors.ConcretizationTypeError:
            negative = False  # built under a transform: the sign is not known
        if negative:
            raise ConfigurationError(f"weight_decay must be >= 0, got {weight_decay}.")
        if residual is not None and not isinstance(residual, ResidualPenalty):
            raise InputTypeError(
                "residual must be a vardax.ResidualPenalty or None; got "
                f"{type(residual).__name__}."
            )
        if minimiser is None:
            minimiser = OptaxMinimiser()
        if not isinstance(minimiser, OptaxMinimiser | OptimistixMinimiser):
            raise InputTypeError(
                "minimiser must be a vardax.OptaxMinimiser or "
                f"vardax.OptimistixMinimiser; got {type(minimiser).__name__}."
            )
        self.field = field
        self.residual = residual
        self.weight_decay = weight_decay
        self.minimiser = minimiser
        self.filter_spec = filter_spec
        self.t0 = t0

    @property
    def _parameterization(self) -> NeuralFieldParameterization:
        # Only the control split, `as_field` and the penalty are used: the
        # space-time field is evaluated at the observation and collocation
        # points, never on a grid, so the grid is empty. The weight is 1
        # (penalty ½‖θ‖²): alpha multiplies it in the cost, as an array leaf.
        return NeuralFieldParameterization(
            self.field,
            jnp.zeros((0, 1)),
            1.0,
            filter_spec=self.filter_spec,
        )

    def _fixed_residual(self, key: PRNGKeyArray | None) -> ResidualPenalty | None:
        """The residual penalty with its collocation points drawn and fixed."""
        residual = self.residual
        if residual is None or residual.sampler is None:
            return residual
        if key is None:
            raise MissingInputError(
                "The residual penalty has a sampler: pass key= to draw its "
                "collocation points."
            )
        return ResidualPenalty(
            residual_fn=residual.residual_fn,
            collocation=residual.sampler(key),
            weight=residual.weight,
        )

    def __call__(
        self, obs: PointObsBatch, *, key: PRNGKeyArray | None = None
    ) -> SpaceTimeFieldResult:
        """Fit the field to one sample of point observations.

        Args:
            obs: The observations, one sample, with ``coords`` ordered time
                first (``time_index`` and ``n_times`` are not used). For a
                batch of samples, ``eqx.filter_vmap`` over this call.
            key: Draws the collocation points when the residual penalty
                has a ``sampler``; unused otherwise.

        Returns:
            A [`SpaceTimeFieldResult`][vardax.SpaceTimeFieldResult].

        Raises:
            InputTypeError: If ``obs`` is not a `PointObsBatch`.
            ShapeError: If ``obs`` carries a batch axis.
            MissingInputError: If the penalty needs a key and none is
                given.
        """
        if not isinstance(obs, PointObsBatch):
            raise InputTypeError(
                f"obs must be a vardax.PointObsBatch; got {type(obs).__name__}."
            )
        if obs.values.ndim != 1:
            raise ShapeError(
                "SpaceTimeFieldVar fits one sample: obs.values must be "
                f"(N_obs,), got shape {obs.values.shape}; use eqx.filter_vmap "
                "for a batch."
            )
        return _fit(self, obs, key)

    def ensemble(
        self,
        obs: PointObsBatch,
        field_init: Any,
        n_members: int,
        *,
        key: PRNGKeyArray,
    ) -> SpaceTimeFieldResult:
        r"""A deep ensemble: refit from ``n_members`` random initialisations.

        Each member re-initialises the field with ``field_init`` and draws
        its own collocation points (when the penalty has a sampler), then
        fits as `__call__` does; the members run under ``eqx.filter_vmap``.
        The spread of the members is a heuristic uncertainty: small where
        the observations pin the field down, large away from them. A
        $\theta$-space Laplace push-forward $J \Sigma_\theta J^\top$
        ($J = \partial u / \partial\theta$) is not implemented.

        Args:
            obs: The observations, one sample.
            field_init: ``key -> field``, building a network with the
                same structure as ``self.field`` (e.g.
                ``lambda k: geonnax.SIREN.init(2, 32, 1, depth=3, key=k)``).
            n_members: Number of members, at least 1.
            key: Split into one initialisation key and one collocation key
                per member.

        Returns:
            A [`SpaceTimeFieldResult`][vardax.SpaceTimeFieldResult] whose
            array leaves carry a leading axis of size ``n_members``.

        Raises:
            ConfigurationError: If ``n_members < 1``.
        """
        if n_members < 1:
            raise ConfigurationError(f"n_members must be >= 1, got {n_members}.")

        def member(member_key: PRNGKeyArray) -> SpaceTimeFieldResult:
            init_key, collocation_key = jax.random.split(member_key)
            model = eqx.tree_at(lambda tree: tree.field, self, field_init(init_key))
            return model(obs, key=collocation_key)

        return eqx.filter_vmap(member)(jax.random.split(key, n_members))

    def as_analysis_step(
        self,
        grid: Float[Array, "N D"] | None = None,
        *,
        key: PRNGKeyArray | None = None,
    ) -> _SpaceTimeFieldAnalysisStep:
        r"""Adapt to ``pipekit_cycle.AnalysisStep``.

        The step fits the field to the cycle's observations. With a
        ``grid`` it returns the analysis at the window start on it,
        $u(t_0, s_j)$ (`SpaceTimeFieldResult.initial_condition`), the
        state a ``ForwardModel`` steps on; without one it returns the
        [`SpaceTimeFieldResult`][vardax.SpaceTimeFieldResult] itself.

        The collocation points are drawn **once**, here, from ``key``,
        and reused by every cycle. ``pipekit_cycle.DACycle`` passes no
        time: the times are the first column of each cycle's
        ``PointObsBatch``, so they must be **window-relative**, in the
        same window $[t_0, t_1]$ every cycle as the collocation points.
        The step checks that (when the times are concrete): $t_0$ and the
        valid observation times must lie in the collocation points' time
        range, widened by 10 % on each side for the gaps a random draw
        leaves at the ends.

        Args:
            grid: Spatial points $s_j$ of the returned state, one per row,
                or ``None``.
            key: Collocation key; required when the penalty has a
                ``sampler``.

        Returns:
            The adapter.

        Raises:
            GridError: If ``grid`` is not 2-D.
            MissingInputError: If the penalty has a sampler and ``key`` is
                ``None``.
            ConfigurationError: If $t_0$ is outside the collocation
                window.
        """
        if grid is not None:
            grid = jnp.asarray(grid)
            if grid.ndim != 2:
                raise GridError(f"grid must be 2-D (N, D); got shape {grid.shape}.")
        if self.residual is not None and self.residual.sampler is not None:
            if key is None:
                raise MissingInputError(
                    "The residual penalty has a sampler: pass key= to "
                    "as_analysis_step to draw the collocation points, which "
                    "every cycle then reuses."
                )
        model = eqx.tree_at(
            lambda tree: tree.residual,
            self,
            self._fixed_residual(key),
            is_leaf=lambda leaf: leaf is None,
        )
        window = _collocation_window(model.residual)
        if window is not None and not window[0] <= self.t0 <= window[1]:
            raise ConfigurationError(
                f"t0 = {self.t0} is outside the collocation time window "
                f"[{window[0]:.4g}, {window[1]:.4g}]: the initial condition "
                "would be unconstrained by the physics."
            )
        return _SpaceTimeFieldAnalysisStep(model, grid)

ensemble

ensemble(
    obs: PointObsBatch,
    field_init: Any,
    n_members: int,
    *,
    key: PRNGKeyArray,
) -> SpaceTimeFieldResult

A deep ensemble: refit from n_members random initialisations.

Each member re-initialises the field with field_init and draws its own collocation points (when the penalty has a sampler), then fits as __call__ does; the members run under eqx.filter_vmap. The spread of the members is a heuristic uncertainty: small where the observations pin the field down, large away from them. A \(\theta\)-space Laplace push-forward \(J \Sigma_\theta J^\top\) (\(J = \partial u / \partial\theta\)) is not implemented.

Parameters:

Name	Type	Description	Default
obs	PointObsBatch	The observations, one sample.	required
field_init	Any	key -> field, building a network with the same structure as self.field (e.g. lambda k: geonnax.SIREN.init(2, 32, 1, depth=3, key=k)).	required
n_members	int	Number of members, at least 1.	required
key	PRNGKeyArray	Split into one initialisation key and one collocation key per member.	required

Returns:

Type	Description
SpaceTimeFieldResult	A SpaceTimeFieldResult whose
SpaceTimeFieldResult	array leaves carry a leading axis of size n_members.

Raises:

Type	Description
ConfigurationError	If n_members < 1.

Source code in src/vardax/_src/models/spacetime_field.py

def ensemble(
    self,
    obs: PointObsBatch,
    field_init: Any,
    n_members: int,
    *,
    key: PRNGKeyArray,
) -> SpaceTimeFieldResult:
    r"""A deep ensemble: refit from ``n_members`` random initialisations.

    Each member re-initialises the field with ``field_init`` and draws
    its own collocation points (when the penalty has a sampler), then
    fits as `__call__` does; the members run under ``eqx.filter_vmap``.
    The spread of the members is a heuristic uncertainty: small where
    the observations pin the field down, large away from them. A
    $\theta$-space Laplace push-forward $J \Sigma_\theta J^\top$
    ($J = \partial u / \partial\theta$) is not implemented.

    Args:
        obs: The observations, one sample.
        field_init: ``key -> field``, building a network with the
            same structure as ``self.field`` (e.g.
            ``lambda k: geonnax.SIREN.init(2, 32, 1, depth=3, key=k)``).
        n_members: Number of members, at least 1.
        key: Split into one initialisation key and one collocation key
            per member.

    Returns:
        A [`SpaceTimeFieldResult`][vardax.SpaceTimeFieldResult] whose
        array leaves carry a leading axis of size ``n_members``.

    Raises:
        ConfigurationError: If ``n_members < 1``.
    """
    if n_members < 1:
        raise ConfigurationError(f"n_members must be >= 1, got {n_members}.")

    def member(member_key: PRNGKeyArray) -> SpaceTimeFieldResult:
        init_key, collocation_key = jax.random.split(member_key)
        model = eqx.tree_at(lambda tree: tree.field, self, field_init(init_key))
        return model(obs, key=collocation_key)

    return eqx.filter_vmap(member)(jax.random.split(key, n_members))

as_analysis_step

as_analysis_step(
    grid: Float[Array, "N D"] | None = None,
    *,
    key: PRNGKeyArray | None = None,
) -> _SpaceTimeFieldAnalysisStep

Adapt to pipekit_cycle.AnalysisStep.

The step fits the field to the cycle's observations. With a grid it returns the analysis at the window start on it, \(u(t_0, s_j)\) (SpaceTimeFieldResult.initial_condition), the state a ForwardModel steps on; without one it returns the SpaceTimeFieldResult itself.

The collocation points are drawn once, here, from key, and reused by every cycle. pipekit_cycle.DACycle passes no time: the times are the first column of each cycle's PointObsBatch, so they must be window-relative, in the same window \([t_0, t_1]\) every cycle as the collocation points. The step checks that (when the times are concrete): \(t_0\) and the valid observation times must lie in the collocation points' time range, widened by 10 % on each side for the gaps a random draw leaves at the ends.

Parameters:

Name	Type	Description	Default
grid	Float[Array, 'N D'] | None	Spatial points \(s_j\) of the returned state, one per row, or None.	None
key	PRNGKeyArray | None	Collocation key; required when the penalty has a sampler.	None

Returns:

Type	Description
_SpaceTimeFieldAnalysisStep	The adapter.

Raises:

Type	Description
GridError	If grid is not 2-D.
MissingInputError	If the penalty has a sampler and key is None.
ConfigurationError	If \(t_0\) is outside the collocation window.

Source code in src/vardax/_src/models/spacetime_field.py

def as_analysis_step(
    self,
    grid: Float[Array, "N D"] | None = None,
    *,
    key: PRNGKeyArray | None = None,
) -> _SpaceTimeFieldAnalysisStep:
    r"""Adapt to ``pipekit_cycle.AnalysisStep``.

    The step fits the field to the cycle's observations. With a
    ``grid`` it returns the analysis at the window start on it,
    $u(t_0, s_j)$ (`SpaceTimeFieldResult.initial_condition`), the
    state a ``ForwardModel`` steps on; without one it returns the
    [`SpaceTimeFieldResult`][vardax.SpaceTimeFieldResult] itself.

    The collocation points are drawn **once**, here, from ``key``,
    and reused by every cycle. ``pipekit_cycle.DACycle`` passes no
    time: the times are the first column of each cycle's
    ``PointObsBatch``, so they must be **window-relative**, in the
    same window $[t_0, t_1]$ every cycle as the collocation points.
    The step checks that (when the times are concrete): $t_0$ and the
    valid observation times must lie in the collocation points' time
    range, widened by 10 % on each side for the gaps a random draw
    leaves at the ends.

    Args:
        grid: Spatial points $s_j$ of the returned state, one per row,
            or ``None``.
        key: Collocation key; required when the penalty has a
            ``sampler``.

    Returns:
        The adapter.

    Raises:
        GridError: If ``grid`` is not 2-D.
        MissingInputError: If the penalty has a sampler and ``key`` is
            ``None``.
        ConfigurationError: If $t_0$ is outside the collocation
            window.
    """
    if grid is not None:
        grid = jnp.asarray(grid)
        if grid.ndim != 2:
            raise GridError(f"grid must be 2-D (N, D); got shape {grid.shape}.")
    if self.residual is not None and self.residual.sampler is not None:
        if key is None:
            raise MissingInputError(
                "The residual penalty has a sampler: pass key= to "
                "as_analysis_step to draw the collocation points, which "
                "every cycle then reuses."
            )
    model = eqx.tree_at(
        lambda tree: tree.residual,
        self,
        self._fixed_residual(key),
        is_leaf=lambda leaf: leaf is None,
    )
    window = _collocation_window(model.residual)
    if window is not None and not window[0] <= self.t0 <= window[1]:
        raise ConfigurationError(
            f"t0 = {self.t0} is outside the collocation time window "
            f"[{window[0]:.4g}, {window[1]:.4g}]: the initial condition "
            "would be unconstrained by the physics."
        )
    return _SpaceTimeFieldAnalysisStep(model, grid)

SpaceTimeFieldResult

Bases: Module

A fitted space-time field and its minimisation diagnostics.

Attributes:

Name	Type	Description
field	Any	The analysis \(f_{\theta^\star}\), a single-example callable of one space-time point (t, *s) of shape (1 + D,).
t0	float	Start of the window, \(t_0\), used by initial_condition.
cost	Float[Array, '']	\(J(\theta^\star)\).
grad_norm	Float[Array, '']	\(\lVert\nabla J(\theta^\star)\rVert_2\); a first-order optimiser stopped by its step budget leaves it non-zero.
num_steps	Int[Array, '']	Optimiser steps taken.

An ensemble returns one result whose array leaves carry a leading member axis; evaluate it with eqx.filter_vmap(lambda r: r.initial_condition(grid))(result).

Examples:

Fit a line to three points, then query the fitted field:

>>> import equinox as eqx, jax, jax.numpy as jnp, vardax
>>> class Line(eqx.Module):
...     w: jax.Array
...
...     def __call__(self, s):
...         return (self.w[0] + self.w[1] * s[1])[None]
>>> obs = vardax.PointObsBatch(
...     coords=jnp.array([[0.0, -1.0], [0.5, 0.0], [1.0, 1.0]]),
...     values=jnp.array([-1.0, 0.0, 1.0]),
...     sigma=jnp.ones(3),
...     valid=jnp.ones(3, bool),
... )
>>> ana = vardax.SpaceTimeFieldVar(
...     field=Line(jnp.zeros(2)),
...     weight_decay=1e-6,
...     minimiser=vardax.OptimistixMinimiser(),
... )
>>> result = ana(obs)
>>> result.evaluate(jnp.array([[0.0, 2.0], [0.0, 3.0]])).round(2)
Array([2., 3.], dtype=float32)
>>> bool(result.grad_norm < 1e-3)
True

The analysis at the window start on a spatial grid is the initial condition for a forecast, and a non-2-D grid is rejected:

>>> result.initial_condition(jnp.array([[0.5], [2.0]])).shape
(2,)
>>> result.initial_condition(
...     jnp.array([0.5, 2.0])
... )
Traceback (most recent call last):
...
GridError: grid must be 2-D (N, D); got shape (2,).

Source code in src/vardax/_src/models/spacetime_field.py

class SpaceTimeFieldResult(eqx.Module):
    r"""A fitted space-time field and its minimisation diagnostics.

    Attributes:
        field: The analysis $f_{\theta^\star}$, a single-example callable
            of one space-time point ``(t, *s)`` of shape ``(1 + D,)``.
        t0: Start of the window, $t_0$, used by `initial_condition`.
        cost: $J(\theta^\star)$.
        grad_norm: $\lVert\nabla J(\theta^\star)\rVert_2$; a first-order
            optimiser stopped by its step budget leaves it non-zero.
        num_steps: Optimiser steps taken.

    An [`ensemble`][vardax.SpaceTimeFieldVar.ensemble] returns one result
    whose array leaves carry a leading member axis; evaluate it with
    ``eqx.filter_vmap(lambda r: r.initial_condition(grid))(result)``.

    Examples:
        Fit a line to three points, then query the fitted field:

        >>> import equinox as eqx, jax, jax.numpy as jnp, vardax
        >>> class Line(eqx.Module):
        ...     w: jax.Array
        ...
        ...     def __call__(self, s):
        ...         return (self.w[0] + self.w[1] * s[1])[None]
        >>> obs = vardax.PointObsBatch(
        ...     coords=jnp.array([[0.0, -1.0], [0.5, 0.0], [1.0, 1.0]]),
        ...     values=jnp.array([-1.0, 0.0, 1.0]),
        ...     sigma=jnp.ones(3),
        ...     valid=jnp.ones(3, bool),
        ... )
        >>> ana = vardax.SpaceTimeFieldVar(
        ...     field=Line(jnp.zeros(2)),
        ...     weight_decay=1e-6,
        ...     minimiser=vardax.OptimistixMinimiser(),
        ... )
        >>> result = ana(obs)
        >>> result.evaluate(jnp.array([[0.0, 2.0], [0.0, 3.0]])).round(2)
        Array([2., 3.], dtype=float32)
        >>> bool(result.grad_norm < 1e-3)
        True

        The analysis at the window start on a spatial grid is the initial
        condition for a forecast, and a non-2-D grid is rejected:

        >>> result.initial_condition(jnp.array([[0.5], [2.0]])).shape
        (2,)
        >>> result.initial_condition(
        ...     jnp.array([0.5, 2.0])
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        GridError: grid must be 2-D (N, D); got shape (2,).
    """

    field: Any
    t0: float = eqx.field(static=True)
    cost: Float[Array, ""]
    grad_norm: Float[Array, ""]
    num_steps: Int[Array, ""]

    def evaluate(self, coords: Float[Array, "N P"]) -> Float[Array, "N ..."]:
        """Evaluate the field at space-time points.

        Args:
            coords: Space-time points ``(t, *s)``, one per row.

        Returns:
            One value per point for a single-channel field, else one
            row of ``K`` channels per point.
        """
        return _evaluate(self.field, jnp.asarray(coords))

    def initial_condition(self, grid: Real[Array, "N D"]) -> Float[Array, "N ..."]:
        r"""The analysis at the window start, $u(t_0, s_j)$, on a grid.

        This is the initial state for a forecast with a ``ForwardModel``.

        Args:
            grid: Spatial points $s_j$, one per row.

        Returns:
            One value per grid point for a single-channel field, else one
            row of ``K`` channels per point.

        Raises:
            GridError: If ``grid`` is not 2-D.
        """
        grid = jnp.asarray(grid)
        if grid.ndim != 2:
            raise GridError(f"grid must be 2-D (N, D); got shape {grid.shape}.")
        if not jnp.issubdtype(grid.dtype, jnp.inexact):
            # An integer grid would truncate t0 in the time column.
            grid = grid.astype(jnp.result_type(float))
        time = jnp.full((grid.shape[0],), self.t0, dtype=grid.dtype)
        points = einx.id("n, n d -> n (1 + d)", time, grid)
        return _evaluate(self.field, points)

evaluate

evaluate(
    coords: Float[Array, "N P"],
) -> Float[Array, "N ..."]

Evaluate the field at space-time points.

Parameters:

Name	Type	Description	Default
coords	Float[Array, 'N P']	Space-time points (t, *s), one per row.	required

Returns:

Type	Description
Float[Array, 'N ...']	One value per point for a single-channel field, else one
Float[Array, 'N ...']	row of K channels per point.

Source code in src/vardax/_src/models/spacetime_field.py

def evaluate(self, coords: Float[Array, "N P"]) -> Float[Array, "N ..."]:
    """Evaluate the field at space-time points.

    Args:
        coords: Space-time points ``(t, *s)``, one per row.

    Returns:
        One value per point for a single-channel field, else one
        row of ``K`` channels per point.
    """
    return _evaluate(self.field, jnp.asarray(coords))

initial_condition

initial_condition(
    grid: Real[Array, "N D"],
) -> Float[Array, "N ..."]

The analysis at the window start, \(u(t_0, s_j)\), on a grid.

This is the initial state for a forecast with a ForwardModel.

Parameters:

Name	Type	Description	Default
grid	Real[Array, 'N D']	Spatial points \(s_j\), one per row.	required

Returns:

Type	Description
Float[Array, 'N ...']	One value per grid point for a single-channel field, else one
Float[Array, 'N ...']	row of K channels per point.

Raises:

Type	Description
GridError	If grid is not 2-D.

Source code in src/vardax/_src/models/spacetime_field.py

def initial_condition(self, grid: Real[Array, "N D"]) -> Float[Array, "N ..."]:
    r"""The analysis at the window start, $u(t_0, s_j)$, on a grid.

    This is the initial state for a forecast with a ``ForwardModel``.

    Args:
        grid: Spatial points $s_j$, one per row.

    Returns:
        One value per grid point for a single-channel field, else one
        row of ``K`` channels per point.

    Raises:
        GridError: If ``grid`` is not 2-D.
    """
    grid = jnp.asarray(grid)
    if grid.ndim != 2:
        raise GridError(f"grid must be 2-D (N, D); got shape {grid.shape}.")
    if not jnp.issubdtype(grid.dtype, jnp.inexact):
        # An integer grid would truncate t0 in the time column.
        grid = grid.astype(jnp.result_type(float))
    time = jnp.full((grid.shape[0],), self.t0, dtype=grid.dtype)
    points = einx.id("n, n d -> n (1 + d)", time, grid)
    return _evaluate(self.field, points)

Learned solvers — 4DVarNet

End-to-end-trainable 4DVar: the variational cost is kept explicit, but the inner-loop descent direction is produced by a ConvLSTM gradient modulator instead of a hand-tuned optimiser. The 1D variant operates on Batch1D (e.g. Lorenz-96 trajectories); the 2D variant on Batch2D / Batch2DMultivar fields (e.g. SSH reconstruction). The inner-loop iteration functions live on the Costs, Priors & Solvers page; training utilities on Training & Adjoints.

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

FourDVarNet1D

Bases: Module

End-to-end 4DVarNet model for 1-D spatiotemporal reconstruction.

Minimises the variational cost

\[ J(x) = \|\mathbf{m} \odot (x - y)\|^2 + \lambda \|x - \varphi(x)\|^2 \]

using n_solver_steps learned gradient steps modulated by a ConvLSTM, with the differentiation strategy selected by solver_adjoint.

Algorithm

1. Start from the masked observations, \(x_0 = m \odot y\), with a zero ConvLSTM state.
2. For each of n_solver_steps iterations, compute \(\nabla J(x_k)\), pass it through the ConvLSTM gradient modulator \(\mathcal{G}\) and update \(x_{k+1} = x_k - \alpha\, \mathcal{G}(\cdot)\).
3. Return \(x_K\). A KStepAdjoint runs the one-step solver and an ImplicitAdjoint the fixed-point projection instead; any other adjoint unrolls with ordinary backprop.

Parameters:

Name	Type	Description	Default
state_dim	int	Spatial size \(N\) of the state.	required
n_time	int	Number of time steps \(T\) in a window.	required
latent_dim	int	Dimensionality of the prior's latent code.	32
hidden_dim	int	Number of hidden channels of the ConvLSTM.	64
n_solver_steps	int	Number of solver iterations to unroll.	15
alpha	float	Gradient step-size \(\alpha\).	0.2
prior_weight	float	Weight \(\lambda\) for the prior cost term.	1.0
solver_adjoint	AbstractAdjoint | None	optimistix.AbstractAdjoint selecting the differentiation strategy. Defaults to RecursiveCheckpointAdjoint (standard backprop). Use OneStepAdjoint() for O(1)-memory training (Bolte et al. 2023) or ImplicitAdjoint() for fixed-point projection.	None
key	PRNGKeyArray	PRNG key for the prior and gradient-modulator weights.	required

Attributes:

Name	Type	Description
prior	BilinAEPrior1D	BilinAEPrior1D learned prior.
grad_mod	ConvLSTMGradMod1D	ConvLSTMGradMod1D learned gradient modulator.

Examples:

A tiny model reconstructs a batch with the shape of its input:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     key=jax.random.PRNGKey(0),
... )
>>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
>>> model(batch).shape
(2, 3, 4)

Selecting the one-step adjoint changes only the gradient path, not the forward value:

>>> fast = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     solver_adjoint=vdx.OneStepAdjoint(),
...     key=jax.random.PRNGKey(0),
... )
>>> bool(jnp.allclose(fast(batch), model(batch), atol=1e-5))
True

Compose it with train_step for one optimiser update:

>>> import optax
>>> target = vdx.Batch1D(
...     input=batch.input, mask=batch.mask, target=jnp.zeros((2, 3, 4))
... )
>>> optimizer = optax.adam(1e-3)
>>> opt_state = optimizer.init(model)
>>> model, opt_state, loss = vdx.train_step(model, target, optimizer, opt_state)
>>> loss.shape
()

Source code in src/vardax/_src/model.py

class FourDVarNet1D(eqx.Module):
    r"""End-to-end 4DVarNet model for 1-D spatiotemporal reconstruction.

    Minimises the variational cost

    $$
    J(x) = \|\mathbf{m} \odot (x - y)\|^2 + \lambda \|x - \varphi(x)\|^2
    $$

    using ``n_solver_steps`` learned gradient steps modulated by a ConvLSTM,
    with the differentiation strategy selected by ``solver_adjoint``.

    Algorithm:
        1. Start from the masked observations, $x_0 = m \odot y$, with a zero
           ConvLSTM state.
        2. For each of ``n_solver_steps`` iterations, compute
           $\nabla J(x_k)$, pass it through the ConvLSTM gradient modulator
           $\mathcal{G}$ and update $x_{k+1} = x_k - \alpha\, \mathcal{G}(\cdot)$.
        3. Return $x_K$.  A ``KStepAdjoint`` runs the one-step solver and an
           ``ImplicitAdjoint`` the fixed-point projection instead; any other
           adjoint unrolls with ordinary backprop.

    Args:
        state_dim: Spatial size $N$ of the state.
        n_time: Number of time steps $T$ in a window.
        latent_dim: Dimensionality of the prior's latent code.
        hidden_dim: Number of hidden channels of the ConvLSTM.
        n_solver_steps: Number of solver iterations to unroll.
        alpha: Gradient step-size $\alpha$.
        prior_weight: Weight $\lambda$ for the prior cost term.
        solver_adjoint: ``optimistix.AbstractAdjoint`` selecting the
            differentiation strategy. Defaults to
            ``RecursiveCheckpointAdjoint`` (standard backprop). Use
            ``OneStepAdjoint()`` for O(1)-memory training (Bolte et al.
            2023) or ``ImplicitAdjoint()`` for fixed-point projection.
        key: PRNG key for the prior and gradient-modulator weights.

    Attributes:
        prior: BilinAEPrior1D learned prior.
        grad_mod: ConvLSTMGradMod1D learned gradient modulator.

    Examples:
        A tiny model reconstructs a batch with the shape of its input:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
        >>> model(batch).shape
        (2, 3, 4)

        Selecting the one-step adjoint changes only the gradient path, not
        the forward value:

        >>> fast = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     solver_adjoint=vdx.OneStepAdjoint(),
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> bool(jnp.allclose(fast(batch), model(batch), atol=1e-5))
        True

        Compose it with [`train_step`][vardax.train_step] for one optimiser
        update:

        >>> import optax
        >>> target = vdx.Batch1D(
        ...     input=batch.input, mask=batch.mask, target=jnp.zeros((2, 3, 4))
        ... )
        >>> optimizer = optax.adam(1e-3)
        >>> opt_state = optimizer.init(model)
        >>> model, opt_state, loss = vdx.train_step(model, target, optimizer, opt_state)
        >>> loss.shape
        ()
    """

    n_solver_steps: int = eqx.field(static=True)
    alpha: float = eqx.field(static=True)
    prior_weight: float = eqx.field(static=True)
    solver_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    prior: BilinAEPrior1D
    grad_mod: ConvLSTMGradMod1D

    def __init__(
        self,
        state_dim: int,
        n_time: int,
        latent_dim: int = 32,
        hidden_dim: int = 64,
        n_solver_steps: int = 15,
        alpha: float = 0.2,
        prior_weight: float = 1.0,
        solver_adjoint: optx.AbstractAdjoint | None = None,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.n_solver_steps = n_solver_steps
        self.alpha = alpha
        self.prior_weight = prior_weight
        self.solver_adjoint = solver_adjoint or _default_adjoint()
        key_prior, key_grad = jax.random.split(key)
        self.prior = BilinAEPrior1D(
            state_dim=state_dim,
            latent_dim=latent_dim,
            n_time=n_time,
            key=key_prior,
        )
        self.grad_mod = ConvLSTMGradMod1D(
            state_channels=n_time,
            hidden_dim=hidden_dim,
            key=key_grad,
        )

    def __call__(self, batch: Batch1D) -> Float[Array, "B T N"]:
        """Run the solver and return the final state estimate.

        Dispatches on ``solver_adjoint`` to pick the differentiation path.

        Args:
            batch: Observed batch $(y, m)$.

        Returns:
            Reconstructed state $x_K$ with the shape of ``batch.input``.
        """
        if isinstance(self.solver_adjoint, KStepAdjoint):
            return self._call_one_step(batch)
        if isinstance(self.solver_adjoint, optx.ImplicitAdjoint):
            return self._call_implicit(batch)
        # Default: optimistix.RecursiveCheckpointAdjoint() and anything
        # else falls through to the unrolled backprop path.
        return self._call_unrolled(batch)

    def _call_unrolled(self, batch: Batch1D) -> Float[Array, "B T N"]:
        batch_size, _, n_space = batch.input.shape
        # u₀ = m ⊙ y
        state = batch.input * batch.mask
        lstm = LSTMState1D.zeros(batch_size, self.grad_mod.hidden_dim, n_space)

        for _ in range(self.n_solver_steps):

            def cost_fn(current_state):
                obs_diff = batch.mask * (current_state - batch.input)
                # J_obs = ‖m ⊙ (u − y)‖²
                obs_cost = jnp.sum(obs_diff**2)
                # J_prior = λ ‖u − φ(u)‖²
                prior_cost = self.prior_weight * jnp.sum(
                    (current_state - self.prior(current_state)) ** 2
                )
                return obs_cost + prior_cost

            grad = jax.grad(cost_fn)(state)
            update, lstm = self.grad_mod(grad, state, lstm)
            # u ← u − α · update
            state = state - self.alpha * update

        return state

    def _call_one_step(self, batch: Batch1D) -> Float[Array, "B T N"]:
        from .solver import one_step_solve_4dvarnet_1d

        return one_step_solve_4dvarnet_1d(
            batch,
            self.prior,
            self.grad_mod,
            n_steps=self.n_solver_steps,
            hidden_dim=self.grad_mod.hidden_dim,
            alpha=self.alpha,
            prior_weight=self.prior_weight,
            k=getattr(self.solver_adjoint, "k", 1),
        )

    def _call_implicit(self, batch: Batch1D) -> Float[Array, "B T N"]:
        # Uses the fixed-point projection solver (prior only; grad_mod
        # and alpha are not used in this mode).
        from .solver import solve_4dvarnet_1d_fixedpoint

        return solve_4dvarnet_1d_fixedpoint(
            batch, self.prior, n_fp_steps=self.n_solver_steps
        )

    def as_analysis_step(self) -> _FourDVarNet1DAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
            masks the non-finite entries of ``obs`` and runs this model on them.

        Examples:
            Non-finite observations become unobserved entries:

            >>> import jax
            >>> import jax.numpy as jnp
            >>> import vardax as vdx
            >>> model = vdx.FourDVarNet1D(
            ...     state_dim=4,
            ...     n_time=3,
            ...     latent_dim=2,
            ...     hidden_dim=3,
            ...     n_solver_steps=1,
            ...     key=jax.random.PRNGKey(0),
            ... )
            >>> step = model.as_analysis_step()
            >>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
            >>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
            (1, 3, 4)
        """
        return _FourDVarNet1DAnalysisStep(self)

as_analysis_step

as_analysis_step() -> _FourDVarNet1DAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_FourDVarNet1DAnalysisStep	A callable (forecast, obs, *, obs_op, obs_err_cov) -> analysis that
_FourDVarNet1DAnalysisStep	masks the non-finite entries of obs and runs this model on them.

Examples:

Non-finite observations become unobserved entries:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=1,
...     key=jax.random.PRNGKey(0),
... )
>>> step = model.as_analysis_step()
>>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
>>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
(1, 3, 4)

Source code in src/vardax/_src/model.py

def as_analysis_step(self) -> _FourDVarNet1DAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
        masks the non-finite entries of ``obs`` and runs this model on them.

    Examples:
        Non-finite observations become unobserved entries:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=1,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> step = model.as_analysis_step()
        >>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
        >>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
        (1, 3, 4)
    """
    return _FourDVarNet1DAnalysisStep(self)

FourDVarNet2D

Bases: Module

End-to-end 4DVarNet model for 2-D spatiotemporal reconstruction.

Minimises the same variational cost as FourDVarNet1D,

\[ J(x) = \|\mathbf{m} \odot (x - y)\|^2 + \lambda \|x - \varphi(x)\|^2 \]

over fields with two spatial axes.

Algorithm

1. Start from the masked observations, \(x_0 = m \odot y\), with a zero ConvLSTM state.
2. For each of n_solver_steps iterations, compute \(\nabla J(x_k)\), pass it through the ConvLSTM gradient modulator \(\mathcal{G}\) and update \(x_{k+1} = x_k - \alpha\, \mathcal{G}(\cdot)\).
3. Return \(x_K\). A KStepAdjoint runs the one-step solver; any other supported adjoint unrolls with ordinary backprop.

Parameters:

Name	Type	Description	Default
n_time	int	Number of time steps \(T\) in a window.	required
height	int	Spatial height \(H\).	required
width	int	Spatial width \(W\).	required
latent_dim	int	Dimensionality of the prior's latent code.	64
hidden_dim	int	Number of hidden channels of the ConvLSTM.	64
n_solver_steps	int	Number of solver iterations to unroll.	15
alpha	float	Gradient step-size \(\alpha\).	0.2
prior_weight	float	Weight \(\lambda\) for the prior cost term.	1.0
solver_adjoint	AbstractAdjoint | None	optimistix.AbstractAdjoint selecting the differentiation strategy. Defaults to RecursiveCheckpointAdjoint (standard backprop). ImplicitAdjoint is not yet implemented for 2-D models.	None
key	PRNGKeyArray	PRNG key for the prior and gradient-modulator weights.	required

Attributes:

Name	Type	Description
prior	BilinAEPrior2D	BilinAEPrior2D learned prior.
grad_mod	ConvLSTMGradMod2D	ConvLSTMGradMod2D learned gradient modulator.

Examples:

A tiny model reconstructs a batch with the shape of its input:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> model = vdx.FourDVarNet2D(
...     n_time=3,
...     height=4,
...     width=5,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     key=jax.random.PRNGKey(0),
... )
>>> batch = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
... )
>>> model(batch).shape
(1, 3, 4, 5)

Selecting the one-step adjoint changes only the gradient path, not the forward value:

>>> fast = vdx.FourDVarNet2D(
...     n_time=3,
...     height=4,
...     width=5,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     solver_adjoint=vdx.OneStepAdjoint(),
...     key=jax.random.PRNGKey(0),
... )
>>> bool(jnp.allclose(fast(batch), model(batch), atol=1e-5))
True

Source code in src/vardax/_src/model.py

class FourDVarNet2D(eqx.Module):
    r"""End-to-end 4DVarNet model for 2-D spatiotemporal reconstruction.

    Minimises the same variational cost as [`FourDVarNet1D`][vardax.FourDVarNet1D],

    $$
    J(x) = \|\mathbf{m} \odot (x - y)\|^2 + \lambda \|x - \varphi(x)\|^2
    $$

    over fields with two spatial axes.

    Algorithm:
        1. Start from the masked observations, $x_0 = m \odot y$, with a zero
           ConvLSTM state.
        2. For each of ``n_solver_steps`` iterations, compute
           $\nabla J(x_k)$, pass it through the ConvLSTM gradient modulator
           $\mathcal{G}$ and update $x_{k+1} = x_k - \alpha\, \mathcal{G}(\cdot)$.
        3. Return $x_K$.  A ``KStepAdjoint`` runs the one-step solver; any
           other supported adjoint unrolls with ordinary backprop.

    Args:
        n_time: Number of time steps $T$ in a window.
        height: Spatial height $H$.
        width: Spatial width $W$.
        latent_dim: Dimensionality of the prior's latent code.
        hidden_dim: Number of hidden channels of the ConvLSTM.
        n_solver_steps: Number of solver iterations to unroll.
        alpha: Gradient step-size $\alpha$.
        prior_weight: Weight $\lambda$ for the prior cost term.
        solver_adjoint: ``optimistix.AbstractAdjoint`` selecting the
            differentiation strategy. Defaults to
            ``RecursiveCheckpointAdjoint`` (standard backprop).
            ``ImplicitAdjoint`` is not yet implemented for 2-D models.
        key: PRNG key for the prior and gradient-modulator weights.

    Attributes:
        prior: BilinAEPrior2D learned prior.
        grad_mod: ConvLSTMGradMod2D learned gradient modulator.

    Examples:
        A tiny model reconstructs a batch with the shape of its input:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet2D(
        ...     n_time=3,
        ...     height=4,
        ...     width=5,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> batch = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
        ... )
        >>> model(batch).shape
        (1, 3, 4, 5)

        Selecting the one-step adjoint changes only the gradient path, not
        the forward value:

        >>> fast = vdx.FourDVarNet2D(
        ...     n_time=3,
        ...     height=4,
        ...     width=5,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     solver_adjoint=vdx.OneStepAdjoint(),
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> bool(jnp.allclose(fast(batch), model(batch), atol=1e-5))
        True
    """

    n_solver_steps: int = eqx.field(static=True)
    alpha: float = eqx.field(static=True)
    prior_weight: float = eqx.field(static=True)
    solver_adjoint: optx.AbstractAdjoint = eqx.field(static=True)
    prior: BilinAEPrior2D
    grad_mod: ConvLSTMGradMod2D

    def __init__(
        self,
        n_time: int,
        height: int,
        width: int,
        latent_dim: int = 64,
        hidden_dim: int = 64,
        n_solver_steps: int = 15,
        alpha: float = 0.2,
        prior_weight: float = 1.0,
        solver_adjoint: optx.AbstractAdjoint | None = None,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.n_solver_steps = n_solver_steps
        self.alpha = alpha
        self.prior_weight = prior_weight
        self.solver_adjoint = solver_adjoint or _default_adjoint()
        key_prior, key_grad = jax.random.split(key)
        self.prior = BilinAEPrior2D(
            latent_dim=latent_dim,
            n_time=n_time,
            height=height,
            width=width,
            key=key_prior,
        )
        self.grad_mod = ConvLSTMGradMod2D(
            state_channels=n_time,
            hidden_dim=hidden_dim,
            key=key_grad,
        )

    def __call__(self, batch: Batch2D) -> Float[Array, "B T H W"]:
        """Run the solver and return the final state estimate.

        Dispatches on ``solver_adjoint`` to pick the differentiation path.

        Args:
            batch: Observed batch $(y, m)$.

        Returns:
            Reconstructed state $x_K$ with the shape of ``batch.input``.
        """
        if isinstance(self.solver_adjoint, KStepAdjoint):
            return self._call_one_step(batch)
        if isinstance(self.solver_adjoint, optx.ImplicitAdjoint):
            return self._call_implicit(batch)
        return self._call_unrolled(batch)

    def _call_unrolled(self, batch: Batch2D) -> Float[Array, "B T H W"]:
        batch_size, _, height, width = batch.input.shape
        # u₀ = m ⊙ y
        state = batch.input * batch.mask
        lstm = LSTMState2D.zeros(batch_size, self.grad_mod.hidden_dim, height, width)

        for _ in range(self.n_solver_steps):

            def cost_fn(current_state):
                obs_diff = batch.mask * (current_state - batch.input)
                # J_obs = ‖m ⊙ (u − y)‖²
                obs_cost = jnp.sum(obs_diff**2)
                # J_prior = λ ‖u − φ(u)‖²
                prior_cost = self.prior_weight * jnp.sum(
                    (current_state - self.prior(current_state)) ** 2
                )
                return obs_cost + prior_cost

            grad = jax.grad(cost_fn)(state)
            update, lstm = self.grad_mod(grad, state, lstm)
            # u ← u − α · update
            state = state - self.alpha * update

        return state

    def _call_one_step(self, batch: Batch2D) -> Float[Array, "B T H W"]:
        from .solver import one_step_solve_4dvarnet_2d

        return one_step_solve_4dvarnet_2d(
            batch,
            self.prior,
            self.grad_mod,
            n_steps=self.n_solver_steps,
            hidden_dim=self.grad_mod.hidden_dim,
            alpha=self.alpha,
            prior_weight=self.prior_weight,
            k=getattr(self.solver_adjoint, "k", 1),
        )

    def _call_implicit(self, batch: Batch2D) -> Float[Array, "B T H W"]:
        raise NotImplementedError(
            "ImplicitAdjoint for FourDVarNet2D is not yet implemented; "
            "use RecursiveCheckpointAdjoint or OneStepAdjoint."
        )

    def as_analysis_step(self) -> _FourDVarNet2DAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
            masks the non-finite entries of ``obs`` and runs this model on them.
        """
        return _FourDVarNet2DAnalysisStep(self)

as_analysis_step

as_analysis_step() -> _FourDVarNet2DAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_FourDVarNet2DAnalysisStep	A callable (forecast, obs, *, obs_op, obs_err_cov) -> analysis that
_FourDVarNet2DAnalysisStep	masks the non-finite entries of obs and runs this model on them.

Source code in src/vardax/_src/model.py

def as_analysis_step(self) -> _FourDVarNet2DAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
        masks the non-finite entries of ``obs`` and runs this model on them.
    """
    return _FourDVarNet2DAnalysisStep(self)