Protocols & Types¶

vardax components are coupled structurally: anything with the right methods conforms, with no base class to inherit. Three of the Protocols — AnalysisStep, ForwardModel, and ObservationOperator — are re-exported from pipekit-cycle and define the seam along which vardax plugs into assimilation cycles; the rest are vardax-specific and define the seams inside a variational method (prior, cost, gradient modulator, posterior adapter, minimiser). All are runtime-checkable, so isinstance checks work at the boundaries.

pipekit-cycle protocols¶

The orchestration contract. ForwardModel propagates state between cycle times, ObservationOperator maps state to observation space (with linearize for the tangent-linear), and AnalysisStep is what every vardax model's .as_analysis_step() returns.

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(...)

AnalysisStep ¶

Bases: Protocol

Combine forecast state with observations to produce analysis state.

Implementations: ensemble Kalman analyses (EnKF, ETKF, LETKF), variational solvers (3D/4D-Var), particle filters, smoothers. Algorithm libraries supply concrete classes.

Members

call(forecast, obs, *, obs_op, obs_err_cov): Return the analysis state given the forecast, the observations, the observation operator, and the observation-error covariance.

Source code in pipekit_cycle/protocols.py

@runtime_checkable
class AnalysisStep(Protocol):
    """Combine forecast state with observations to produce analysis state.

    Implementations: ensemble Kalman analyses (EnKF, ETKF, LETKF),
    variational solvers (3D/4D-Var), particle filters, smoothers.
    Algorithm libraries supply concrete classes.

    Members:
        __call__(forecast, obs, *, obs_op, obs_err_cov): Return the
            analysis state given the forecast, the observations, the
            observation operator, and the observation-error covariance.
    """

    def __call__(
        self,
        forecast: Any,
        obs: Any,
        *,
        obs_op: ObservationOperator,
        obs_err_cov: Any,
    ) -> Any: ...

ForwardModel ¶

Bases: Protocol

Advance a model state forward in time by dt.

Implementations: domain forward models (chemistry transport, ocean state, plume dispersion), neural emulators, hybrid physics + ML hybrids. Algorithm libraries provide adapters that satisfy this protocol structurally.

Members

step(state, dt): Return the state advanced by dt. dt: Default integration step. state_signature: Optional pipekit.Signature describing the shape / dtype of the state carrier. None if the model doesn't track named dimensions.

Source code in pipekit_cycle/protocols.py

@runtime_checkable
class ForwardModel(Protocol):
    """Advance a model state forward in time by ``dt``.

    Implementations: domain forward models (chemistry transport,
    ocean state, plume dispersion), neural emulators, hybrid
    physics + ML hybrids. Algorithm libraries provide adapters
    that satisfy this protocol structurally.

    Members:
        step(state, dt): Return the state advanced by ``dt``.
        dt: Default integration step.
        state_signature: Optional `pipekit.Signature` describing the
            shape / dtype of the state carrier. ``None`` if the model
            doesn't track named dimensions.
    """

    def step(self, state: Any, dt: float) -> Any: ...

    @property
    def dt(self) -> float: ...

    @property
    def state_signature(self) -> Any: ...

ObservationOperator ¶

Bases: Protocol

Map model state → predicted observations.

The H operator in classical data-assimilation notation: H(x) produces "what would the observations look like if the state were x?". The innovation in DA is then obs - H(forecast).

Members

call(state): Return predicted observations for state. linearize(state): Optional tangent-linear operator at state (returns a callable / matrix). Implementations that don't expose a linearisation may raise NotImplementedError.

Source code in pipekit_cycle/protocols.py

@runtime_checkable
class ObservationOperator(Protocol):
    """Map model state → predicted observations.

    The H operator in classical data-assimilation notation:
    ``H(x)`` produces "what would the observations look like if the
    state were ``x``?". The innovation in DA is then ``obs - H(forecast)``.

    Members:
        __call__(state): Return predicted observations for ``state``.
        linearize(state): Optional tangent-linear operator at
            ``state`` (returns a callable / matrix). Implementations
            that don't expose a linearisation may raise
            ``NotImplementedError``.
    """

    def __call__(self, state: Any) -> Any: ...

    def linearize(self, state: Any) -> Any: ...

vardax protocols¶

The internal seams: implement these to swap in custom priors, cost functions, gradient modulators, posterior adapters, or inner-loop minimisers without touching the model classes.

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(...)

Prior ¶

Bases: Protocol

Prior model: maps state to its regularised reconstruction.

For an autoencoder prior \(\varphi_\theta\), the variational cost includes ||x - φ(x)||^2. For a dynamical prior wrapping a ForwardModel, φ(x) is the forward integration. For the identity prior, φ(x) = x.

Members

__call__(x) -> x_prior — apply the prior model.

Examples:

A callable that shrinks the state towards zero conforms.

>>> import jax.numpy as jnp, vardax
>>> class Shrink:
...     def __call__(self, x):
...         return 0.9 * x
>>> isinstance(Shrink(), vardax.Prior)
True

The built-in identity prior conforms too.

>>> isinstance(vardax.IdentityPrior(), vardax.Prior)
True
>>> float(vardax.IdentityPrior()(jnp.ones(3)).sum())
3.0

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

@runtime_checkable
class Prior(Protocol):
    r"""Prior model: maps state to its regularised reconstruction.

    For an autoencoder prior $\varphi_\theta$, the variational
    cost includes ``||x - φ(x)||^2``. For a dynamical prior wrapping a
    `ForwardModel`, ``φ(x)`` is the forward integration. For the
    identity prior, ``φ(x) = x``.

    Members:
        ``__call__(x) -> x_prior`` — apply the prior model.

    Examples:
        A callable that shrinks the state towards zero conforms.

        >>> import jax.numpy as jnp, vardax
        >>> class Shrink:
        ...     def __call__(self, x):
        ...         return 0.9 * x
        >>> isinstance(Shrink(), vardax.Prior)
        True

        The built-in identity prior conforms too.

        >>> isinstance(vardax.IdentityPrior(), vardax.Prior)
        True
        >>> float(vardax.IdentityPrior()(jnp.ones(3)).sum())
        3.0
    """

    def __call__(self, x: Float[Array, ...]) -> Float[Array, ...]: ...

TemporalPrior ¶

Bases: Protocol

Time-dependent prior: dynamics-aware regularisation over a window.

Unlike the static Prior seam (\(\varphi: x \mapsto x_\text{prior}\)), a temporal prior needs the time coordinates of the assimilation window, \(\varphi: (x, t_s) \mapsto x_\text{pred}\), with a loss scoring dynamical consistency of a state sequence (Decision D18). Implementations: DynIncrements, DynTrajectory.

A temporal prior is adapted to contexts expecting the one-argument Prior seam by binding the time grid: prior.bind(ts) returns a Prior-conforming callable.

Members

__call__(x, ts) -> x_pred — propagate through the dynamics. loss(x, ts, x_gt=None, params=None) -> scalar — dynamical residual of the sequence x against x_gt.

Examples:

>>> import jax.numpy as jnp, vardax
>>> class Persistence:
...     def __call__(self, x, ts):
...         return x
...
...     def loss(self, x, ts, x_gt=None, params=None):
...         return jnp.mean((x[1:] - x[:-1]) ** 2)
>>> prior = Persistence()
>>> isinstance(prior, vardax.TemporalPrior)
True
>>> float(prior.loss(jnp.ones((4, 2)), jnp.arange(4.0)))
0.0

An object with only __call__ is not a temporal prior.

>>> isinstance(vardax.IdentityPrior(), vardax.TemporalPrior)
False

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

@runtime_checkable
class TemporalPrior(Protocol):
    r"""Time-dependent prior: dynamics-aware regularisation over a window.

    Unlike the static [`Prior`][vardax.Prior] seam
    ($\varphi: x \mapsto x_\text{prior}$), a temporal prior needs the time
    coordinates of the assimilation window,
    $\varphi: (x, t_s) \mapsto x_\text{pred}$, with a
    ``loss`` scoring dynamical consistency of a state sequence (Decision
    D18). Implementations: [`DynIncrements`][vardax.DynIncrements],
    [`DynTrajectory`][vardax.DynTrajectory].

    A temporal prior is adapted to contexts expecting the one-argument
    `Prior` seam by binding the time grid:
    ``prior.bind(ts)`` returns a `Prior`-conforming callable.

    Members:
        ``__call__(x, ts) -> x_pred`` — propagate through the dynamics.
        ``loss(x, ts, x_gt=None, params=None) -> scalar`` — dynamical
        residual of the sequence ``x`` against ``x_gt``.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> class Persistence:
        ...     def __call__(self, x, ts):
        ...         return x
        ...
        ...     def loss(self, x, ts, x_gt=None, params=None):
        ...         return jnp.mean((x[1:] - x[:-1]) ** 2)
        >>> prior = Persistence()
        >>> isinstance(prior, vardax.TemporalPrior)
        True
        >>> float(prior.loss(jnp.ones((4, 2)), jnp.arange(4.0)))
        0.0

        An object with only ``__call__`` is not a temporal prior.

        >>> isinstance(vardax.IdentityPrior(), vardax.TemporalPrior)
        False
    """

    def __call__(
        self,
        x: Float[Array, ...],
        ts: Float[Array, ...],
    ) -> Float[Array, ...]: ...

    def loss(
        self,
        x: Float[Array, ...],
        ts: Float[Array, ...],
        x_gt: Float[Array, ...] | None = None,
        params: Any = None,
    ) -> Float[Array, ""]: ...

CostFunction ¶

Bases: Protocol

Variational cost function J(x, batch, **kwargs) -> scalar.

Implementations include vardax.costs.ThreeDVarCost, StrongConstraintCost, WeakConstraintCost, IncrementalCost, and FourDVarNetCost.

Members

__call__(x, batch, **kwargs) -> scalar

Examples:

A quadratic cost around the batch.

>>> import jax.numpy as jnp, vardax
>>> class Quadratic:
...     def __call__(self, x, batch, **kwargs):
...         return 0.5 * jnp.sum((x - batch) ** 2)
>>> cost = Quadratic()
>>> isinstance(cost, vardax.CostFunction)
True
>>> float(cost(jnp.ones(3), jnp.zeros(3)))
1.5

It is the argument a Minimiser works on.

>>> x = vardax.OptimistixMinimiser()(cost, jnp.zeros(3), jnp.ones(3))
>>> bool(jnp.allclose(x, 1.0, atol=1e-5))
True

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

@runtime_checkable
class CostFunction(Protocol):
    """Variational cost function ``J(x, batch, **kwargs) -> scalar``.

    Implementations include `vardax.costs.ThreeDVarCost`,
    `StrongConstraintCost`, `WeakConstraintCost`, `IncrementalCost`,
    and `FourDVarNetCost`.

    Members:
        ``__call__(x, batch, **kwargs) -> scalar``

    Examples:
        A quadratic cost around the batch.

        >>> import jax.numpy as jnp, vardax
        >>> class Quadratic:
        ...     def __call__(self, x, batch, **kwargs):
        ...         return 0.5 * jnp.sum((x - batch) ** 2)
        >>> cost = Quadratic()
        >>> isinstance(cost, vardax.CostFunction)
        True
        >>> float(cost(jnp.ones(3), jnp.zeros(3)))
        1.5

        It is the argument a [`Minimiser`][vardax.Minimiser] works on.

        >>> x = vardax.OptimistixMinimiser()(cost, jnp.zeros(3), jnp.ones(3))
        >>> bool(jnp.allclose(x, 1.0, atol=1e-5))
        True
    """

    def __call__(
        self,
        x: Float[Array, ...],
        batch: Any,
        **kwargs: Any,
    ) -> Float[Array, ""]: ...

GradModulator ¶

Bases: Protocol

Learned gradient modulator for the FourDVarNet inner solver.

Takes the current variational-cost gradient and the modulator's own carry state, returns a state update and the new carry. Used only by FourDVarNet; the classical analysis methods use optimistix.AbstractMinimiser instead.

Members

__call__(grad, state, carry) -> (update, new_carry)

Examples:

A modulator that scales the gradient and carries nothing.

>>> import jax.numpy as jnp, vardax
>>> class Scale:
...     def __call__(self, grad, state, carry):
...         return -0.1 * grad, carry
>>> isinstance(Scale(), vardax.GradModulator)
True
>>> update, carry = Scale()(jnp.ones(3), jnp.zeros(3), None)
>>> update.shape, carry is None
((3,), True)

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

@runtime_checkable
class GradModulator(Protocol):
    """Learned gradient modulator for the FourDVarNet inner solver.

    Takes the current variational-cost gradient and the modulator's
    own carry state, returns a state update and the new carry. Used
    only by `FourDVarNet`; the classical analysis methods use
    ``optimistix.AbstractMinimiser`` instead.

    Members:
        ``__call__(grad, state, carry) -> (update, new_carry)``

    Examples:
        A modulator that scales the gradient and carries nothing.

        >>> import jax.numpy as jnp, vardax
        >>> class Scale:
        ...     def __call__(self, grad, state, carry):
        ...         return -0.1 * grad, carry
        >>> isinstance(Scale(), vardax.GradModulator)
        True
        >>> update, carry = Scale()(jnp.ones(3), jnp.zeros(3), None)
        >>> update.shape, carry is None
        ((3,), True)
    """

    def __call__(
        self,
        grad: Float[Array, ...],
        state: Float[Array, ...],
        carry: Any,
    ) -> tuple[Float[Array, ...], Any]: ...

PosteriorAdapter ¶

Bases: Protocol

Turns an analysis output into a Posterior container.

Implementations: LaplaceCovariance, GaussNewtonHessian, EnsembleCovariance. Each computes the posterior covariance via a different approximation; the contract returned is the same.

Members

__call__(analysis, model, batch) -> Posterior

Examples:

>>> import jax.numpy as jnp, vardax
>>> class Passthrough:
...     def __call__(self, analysis, model, batch):
...         return analysis
>>> isinstance(Passthrough(), vardax.PosteriorAdapter)
True

The built-in adapters conform as well.

>>> isinstance(vardax.LaplaceCovariance(), vardax.PosteriorAdapter)
True

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

@runtime_checkable
class PosteriorAdapter(Protocol):
    """Turns an analysis output into a `Posterior` container.

    Implementations: `LaplaceCovariance`, `GaussNewtonHessian`,
    `EnsembleCovariance`. Each computes the posterior covariance via a
    different approximation; the contract returned is the same.

    Members:
        ``__call__(analysis, model, batch) -> Posterior``

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> class Passthrough:
        ...     def __call__(self, analysis, model, batch):
        ...         return analysis
        >>> isinstance(Passthrough(), vardax.PosteriorAdapter)
        True

        The built-in adapters conform as well.

        >>> isinstance(vardax.LaplaceCovariance(), vardax.PosteriorAdapter)
        True
    """

    def __call__(
        self,
        analysis: Float[Array, ...],
        model: AnalysisStep,
        batch: Any,
    ) -> Any: ...

Minimiser ¶

Bases: Protocol

Minimises a cost cost_fn(x, batch) from an initial guess.

A Minimiser knows how to minimise a CostFunction from an initial guess x0 (any pytree) against a batch. Implementations: OptimistixMinimiser (optimistix solvers such as BFGS, GaussNewton, NonlinearCG) and OptaxMinimiser (first-order optax optimisers with early stopping, for large controls). The classical models accept either as minimiser=.

Members

__call__(cost_fn, x0, batch) -> x_star

Examples:

A minimiser that takes one gradient step.

>>> import jax, jax.numpy as jnp, vardax
>>> class OneStep:
...     def __call__(self, cost_fn, x0, batch):
...         return x0 - jax.grad(cost_fn)(x0, batch)
>>> isinstance(OneStep(), vardax.Minimiser)
True
>>> cost = lambda x, m: 0.5 * jnp.sum((x - m) ** 2)
>>> x = OneStep()(cost, jnp.zeros(2), jnp.ones(2))
>>> bool(jnp.allclose(x, 1.0))
True

The built-in minimisers conform.

>>> isinstance(vardax.OptaxMinimiser(), vardax.Minimiser)
True

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

@runtime_checkable
class Minimiser(Protocol):
    """Minimises a cost ``cost_fn(x, batch)`` from an initial guess.

    A `Minimiser` knows how to minimise a `CostFunction` from an
    initial guess ``x0`` (any pytree) against a batch. Implementations:
    [`OptimistixMinimiser`][vardax.OptimistixMinimiser] (optimistix
    solvers such as BFGS, GaussNewton, NonlinearCG) and
    [`OptaxMinimiser`][vardax.OptaxMinimiser] (first-order optax
    optimisers with early stopping, for large controls). The classical
    models accept either as ``minimiser=``.

    Members:
        ``__call__(cost_fn, x0, batch) -> x_star``

    Examples:
        A minimiser that takes one gradient step.

        >>> import jax, jax.numpy as jnp, vardax
        >>> class OneStep:
        ...     def __call__(self, cost_fn, x0, batch):
        ...         return x0 - jax.grad(cost_fn)(x0, batch)
        >>> isinstance(OneStep(), vardax.Minimiser)
        True
        >>> cost = lambda x, m: 0.5 * jnp.sum((x - m) ** 2)
        >>> x = OneStep()(cost, jnp.zeros(2), jnp.ones(2))
        >>> bool(jnp.allclose(x, 1.0))
        True

        The built-in minimisers conform.

        >>> isinstance(vardax.OptaxMinimiser(), vardax.Minimiser)
        True
    """

    def __call__(
        self,
        cost_fn: CostFunction,
        x0: Float[Array, ...],
        batch: Any,
    ) -> Float[Array, ...]: ...

ReducedBasis ¶

Bases: Protocol

Reduced-basis control-vector operator.

Maps a coefficient vector \(X \in \mathbb{R}^M\) to a state increment \(\Phi(t) X\) on the grid, with a separable Gaussian background term \(\tfrac{1}{2} X^\top Q^{-1} X\).

Implementations: LinearBasis (any linear family — built via the rbf_basis / fourier_basis / wavelet_basis / eof_basis / linear_basis constructors) and CompositeBasis. Used by variational analysis steps to parameterise the increment in a low-dimensional space.

This protocol is defined in vardax for now; the plan is to upstream it into pipekit_cycle.protocols so other algorithm libraries can satisfy the same contract.

Members

operg(t, coefficients, state=None) -> increment prior_inv(coefficients) -> Q⁻¹ coefficients nbasis: int property

Examples:

The basis built by linear_basis conforms.

>>> import jax.numpy as jnp, vardax
>>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=1.0)
>>> isinstance(basis, vardax.ReducedBasis)
True
>>> basis.nbasis
2

prior_inv applies \(Q^{-1}\) to a coefficient vector.

>>> basis.prior_inv(jnp.ones(2)).shape
(2,)

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

@runtime_checkable
class ReducedBasis(Protocol):
    r"""Reduced-basis control-vector operator.

    Maps a coefficient vector $X \in \mathbb{R}^M$ to a state increment
    $\Phi(t) X$ on the grid, with a separable Gaussian background term
    $\tfrac{1}{2} X^\top Q^{-1} X$.

    Implementations: `LinearBasis` (any linear family — built via the
    `rbf_basis` / `fourier_basis` / `wavelet_basis` / `eof_basis` /
    `linear_basis` constructors) and `CompositeBasis`. Used by
    variational analysis steps to parameterise the increment in a
    low-dimensional space.

    This protocol is defined in vardax for now; the plan is to upstream
    it into ``pipekit_cycle.protocols`` so other algorithm libraries
    can satisfy the same contract.

    Members:
        ``operg(t, coefficients, state=None) -> increment``
        ``prior_inv(coefficients) -> Q⁻¹ coefficients``
        ``nbasis: int`` property

    Examples:
        The basis built by ``linear_basis`` conforms.

        >>> import jax.numpy as jnp, vardax
        >>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=1.0)
        >>> isinstance(basis, vardax.ReducedBasis)
        True
        >>> basis.nbasis
        2

        ``prior_inv`` applies $Q^{-1}$ to a coefficient vector.

        >>> basis.prior_inv(jnp.ones(2)).shape
        (2,)
    """

    def operg(self, t: float, coefficients: Any, state: Any = None) -> Any: ...

    def prior_inv(self, coefficients: Any) -> Any: ...

    @property
    def nbasis(self) -> int: ...

StateParameterization ¶

Bases: Protocol

Control-to-state map with its own penalty.

Generalises ReducedBasis from a linear \(x = x_b + \Phi X\) to any \(x = \psi(\theta)\) with \(\theta\) a pytree. A model given parameterization= minimises over \(\theta\), with \(\mathrm{pen}(\theta)\) in place of the background term: \(J(\theta) = \mathrm{pen}(\theta) + J_o(\psi(\theta))\).

Implementations: LinearParameterization (adapts any ReducedBasis) and NeuralFieldParameterization (a coordinate network on the grid).

Members

decode(theta, t=None) -> state — the state \(\psi(\theta)\), at time t for a time-dependent parameterization. penalty(theta) -> scalar — \(\mathrm{pen}(\theta)\), e.g. \(\tfrac12 X^\top Q^{-1} X\). init_control(*, key=None) -> theta — the starting control; key is for parameterizations whose start is random.

Examples:

>>> import jax.numpy as jnp, vardax
>>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=1.0)
>>> param = vardax.LinearParameterization(basis, background=jnp.zeros(4))
>>> isinstance(param, vardax.StateParameterization)
True
>>> param.decode(param.init_control()).shape
(4,)

A minimal custom parameterization that uses the state itself as the control.

>>> class Direct:
...     def decode(self, theta, t=None):
...         return theta
...
...     def penalty(self, theta):
...         return 0.5 * jnp.sum(theta**2)
...
...     def init_control(self, *, key=None):
...         return jnp.zeros(4)
>>> isinstance(Direct(), vardax.StateParameterization)
True

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

@runtime_checkable
class StateParameterization(Protocol):
    r"""Control-to-state map with its own penalty.

    Generalises [`ReducedBasis`][vardax.ReducedBasis] from a linear
    $x = x_b + \Phi X$ to any $x = \psi(\theta)$ with $\theta$ a pytree.
    A model given ``parameterization=`` minimises over $\theta$, with
    $\mathrm{pen}(\theta)$ in place of the background term:
    $J(\theta) = \mathrm{pen}(\theta) + J_o(\psi(\theta))$.

    Implementations: [`LinearParameterization`][vardax.LinearParameterization]
    (adapts any `ReducedBasis`) and
    [`NeuralFieldParameterization`][vardax.NeuralFieldParameterization]
    (a coordinate network on the grid).

    Members:
        ``decode(theta, t=None) -> state`` — the state $\psi(\theta)$, at
        time ``t`` for a time-dependent parameterization.
        ``penalty(theta) -> scalar`` — $\mathrm{pen}(\theta)$, e.g.
        $\tfrac12 X^\top Q^{-1} X$.
        ``init_control(*, key=None) -> theta`` — the starting control;
        ``key`` is for parameterizations whose start is random.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=1.0)
        >>> param = vardax.LinearParameterization(basis, background=jnp.zeros(4))
        >>> isinstance(param, vardax.StateParameterization)
        True
        >>> param.decode(param.init_control()).shape
        (4,)

        A minimal custom parameterization that uses the state itself as
        the control.

        >>> class Direct:
        ...     def decode(self, theta, t=None):
        ...         return theta
        ...
        ...     def penalty(self, theta):
        ...         return 0.5 * jnp.sum(theta**2)
        ...
        ...     def init_control(self, *, key=None):
        ...         return jnp.zeros(4)
        >>> isinstance(Direct(), vardax.StateParameterization)
        True
    """

    def decode(self, theta: Any, t: float | None = None) -> Float[Array, ...]: ...

    def penalty(self, theta: Any) -> Float[Array, ""]: ...

    def init_control(self, *, key: Any = None) -> Any: ...

SupportsForwardAdjoint ¶

Bases: Protocol

A ForwardModel whose integrator can take a diffrax adjoint.

The 4DVar models' forward_adjoint= slot calls forward.with_adjoint(adjoint) and rolls out the returned forward. A forward that integrates with diffrax.diffeqsolve implements it by storing the adjoint and passing it as diffeqsolve(..., adjoint=adjoint). A per-step forward cannot, and a 4DVar model given forward_adjoint= with such a forward raises UnsupportedInputError; bound its memory with checkpoint= instead.

Members

dt — step length. step(state, dt) -> state — the ForwardModel contract. with_adjoint(adjoint) -> forward — a copy that integrates with adjoint.

Examples:

>>> import dataclasses
>>> import vardax
>>> @dataclasses.dataclass(frozen=True)
... class Decay:
...     dt: float = 1.0
...     adjoint: object = None
...
...     def step(self, x, dt):
...         return 0.5 * x
...
...     def with_adjoint(self, adjoint):
...         return dataclasses.replace(self, adjoint=adjoint)
>>> isinstance(Decay(), vardax.SupportsForwardAdjoint)
True
>>> Decay().with_adjoint("backsolve").adjoint
'backsolve'

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

@runtime_checkable
class SupportsForwardAdjoint(Protocol):
    r"""A ``ForwardModel`` whose integrator can take a ``diffrax`` adjoint.

    The 4DVar models' ``forward_adjoint=`` slot calls
    ``forward.with_adjoint(adjoint)`` and rolls out the returned forward.
    A forward that integrates with ``diffrax.diffeqsolve`` implements it
    by storing the adjoint and passing it as ``diffeqsolve(...,
    adjoint=adjoint)``. A per-step forward cannot, and a 4DVar model given
    ``forward_adjoint=`` with such a forward raises
    [`UnsupportedInputError`][vardax.UnsupportedInputError]; bound its
    memory with ``checkpoint=`` instead.

    Members:
        ``dt`` — step length.
        ``step(state, dt) -> state`` — the ``ForwardModel`` contract.
        ``with_adjoint(adjoint) -> forward`` — a copy that integrates with
        ``adjoint``.

    Examples:
        >>> import dataclasses
        >>> import vardax
        >>> @dataclasses.dataclass(frozen=True)
        ... class Decay:
        ...     dt: float = 1.0
        ...     adjoint: object = None
        ...
        ...     def step(self, x, dt):
        ...         return 0.5 * x
        ...
        ...     def with_adjoint(self, adjoint):
        ...         return dataclasses.replace(self, adjoint=adjoint)
        >>> isinstance(Decay(), vardax.SupportsForwardAdjoint)
        True
        >>> Decay().with_adjoint("backsolve").adjoint
        'backsolve'
    """

    dt: float

    def step(self, state: Any, dt: float) -> Any: ...

    def with_adjoint(self, adjoint: Any) -> Any: ...

Batch & state types¶

The typed carriers that flow through the solvers and training loops: 1D and 2D (single- and multi-variable) observation batches, and the recurrent-state containers of the ConvLSTM gradient modulators.

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(...)

Batch1D ¶

Bases: Module

Batch of 1-D spatiotemporal data.

Holds the observed field \(y\), the observation mask \(m\) and, optionally, the ground truth \(x^\mathrm{gt}\), for samples of a 1-D state over time.

Attributes:

Name	Type	Description
`input`	`Float[Array, ' B *sample']`	Observed (masked) input field \(y\).
`mask`	`Float[Array, ' B *sample']`	Binary observation mask \(m\).
`target`	`Float[Array, ' B *sample'] \| None`	Ground-truth field \(x^\mathrm{gt}\). Optional; absent in operational use where only observations are available.

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch1D(input=jnp.zeros((2, 5, 8)), mask=jnp.ones((2, 5, 8)))
>>> batch.input.shape, batch.target is None
((2, 5, 8), True)

The batch is a pytree, so it passes through JAX transforms.

>>> import jax
>>> len(jax.tree_util.tree_leaves(batch))
2

It feeds the cost functions.

>>> float(vardax.variational_cost(batch.input, batch, vardax.IdentityPrior()))
0.0

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

class Batch1D(eqx.Module):
    r"""Batch of 1-D spatiotemporal data.

    Holds the observed field $y$, the observation mask $m$ and, optionally,
    the ground truth $x^\mathrm{gt}$, for samples of a 1-D state over time.

    Attributes:
        input: Observed (masked) input field $y$.
        mask: Binary observation mask $m$.
        target: Ground-truth field $x^\mathrm{gt}$. Optional; absent in
            operational use where only observations are available.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch1D(input=jnp.zeros((2, 5, 8)), mask=jnp.ones((2, 5, 8)))
        >>> batch.input.shape, batch.target is None
        ((2, 5, 8), True)

        The batch is a pytree, so it passes through JAX transforms.

        >>> import jax
        >>> len(jax.tree_util.tree_leaves(batch))
        2

        It feeds the cost functions.

        >>> float(vardax.variational_cost(batch.input, batch, vardax.IdentityPrior()))
        0.0
    """

    # 4DVarNet batches are (B, T, N); OI and 3D-Var put (B, M) observations
    # through the same type, so only the batch axis is fixed
    input: Float[Array, " B *sample"]
    mask: Float[Array, " B *sample"]
    target: Float[Array, " B *sample"] | None = None

Batch2D ¶

Bases: Module

Batch of 2-D spatiotemporal data.

Holds the observed field \(y\), the observation mask \(m\) and, optionally, the ground truth \(x^\mathrm{gt}\), for samples of a 2-D state over time.

Attributes:

Name	Type	Description
`input`	`Float[Array, 'B T H W']`	Observed (masked) input field \(y\).
`mask`	`Float[Array, 'B T H W']`	Binary observation mask \(m\).
`target`	`Float[Array, 'B T H W'] \| None`	Ground-truth field \(x^\mathrm{gt}\). Optional.

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch2D(
...     input=jnp.zeros((1, 3, 4, 4)),
...     mask=jnp.ones((1, 3, 4, 4)),
...     target=jnp.zeros((1, 3, 4, 4)),
... )
>>> batch.input.shape
(1, 3, 4, 4)

It feeds the 2-D observation cost.

>>> float(vardax.obs_cost_2d(batch.input + 1.0, batch.input, batch.mask))
1.0

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

class Batch2D(eqx.Module):
    r"""Batch of 2-D spatiotemporal data.

    Holds the observed field $y$, the observation mask $m$ and, optionally,
    the ground truth $x^\mathrm{gt}$, for samples of a 2-D state over time.

    Attributes:
        input: Observed (masked) input field $y$.
        mask: Binary observation mask $m$.
        target: Ground-truth field $x^\mathrm{gt}$. Optional.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch2D(
        ...     input=jnp.zeros((1, 3, 4, 4)),
        ...     mask=jnp.ones((1, 3, 4, 4)),
        ...     target=jnp.zeros((1, 3, 4, 4)),
        ... )
        >>> batch.input.shape
        (1, 3, 4, 4)

        It feeds the 2-D observation cost.

        >>> float(vardax.obs_cost_2d(batch.input + 1.0, batch.input, batch.mask))
        1.0
    """

    input: Float[Array, "B T H W"]
    mask: Float[Array, "B T H W"]
    target: Float[Array, "B T H W"] | None = None

Batch2DMultivar ¶

Bases: Module

Batch of 2-D multivariate spatiotemporal data.

Like Batch2D with an extra channel axis for the variables.

Attributes:

Name	Type	Description
`input`	`Float[Array, 'B T C H W']`	Observed (masked) input field \(y\).
`mask`	`Float[Array, 'B T C H W']`	Binary observation mask \(m\).
`target`	`Float[Array, 'B T C H W'] \| None`	Ground-truth field \(x^\mathrm{gt}\). Optional.

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch2DMultivar(
...     input=jnp.zeros((1, 3, 2, 4, 4)), mask=jnp.ones((1, 3, 2, 4, 4))
... )
>>> batch.mask.shape
(1, 3, 2, 4, 4)

The ground truth is optional.

>>> batch.target is None
True

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

class Batch2DMultivar(eqx.Module):
    r"""Batch of 2-D multivariate spatiotemporal data.

    Like [`Batch2D`][vardax.Batch2D] with an extra channel axis for the
    variables.

    Attributes:
        input: Observed (masked) input field $y$.
        mask: Binary observation mask $m$.
        target: Ground-truth field $x^\mathrm{gt}$. Optional.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch2DMultivar(
        ...     input=jnp.zeros((1, 3, 2, 4, 4)), mask=jnp.ones((1, 3, 2, 4, 4))
        ... )
        >>> batch.mask.shape
        (1, 3, 2, 4, 4)

        The ground truth is optional.

        >>> batch.target is None
        True
    """

    input: Float[Array, "B T C H W"]
    mask: Float[Array, "B T C H W"]
    target: Float[Array, "B T C H W"] | None = None

PointObsBatch ¶

Bases: Module

Scattered point observations, padded to a fixed count.

Observation \(i\) has coordinates \(s_i\), value \(y_i\), error standard deviation \(\sigma_i\) and a validity flag \(m_i\); padding rows have valid=False and contribute exactly nothing to the cost or its gradient, whatever their other fields hold (NaN included). The observation cost is \(J_o = \tfrac12 \sum_i m_i (y_i - \mathcal H(x)(s_i, t_i))^2 / \sigma_i^2\) (point_obs_cost).

Every array field may carry a leading batch axis; the models then return one analysis per sample, and without it a single analysis.

Attributes:

Name	Type	Description
`coords`	`Float[Array, '*batch N_obs D']`	Coordinates \(s_i\), e.g. `(x,)`, `(y, x)` or `(lon, lat, t)`.
`values`	`Float[Array, '*batch N_obs']`	Observed values \(y_i\).
`sigma`	`Float[Array, '*batch N_obs']`	Per-observation error standard deviations \(\sigma_i\).
`valid`	`Bool[Array, '*batch N_obs']`	Validity flags \(m_i\), `False` on padding rows.
`time_index`	`Int[Array, '*batch N_obs'] \| None`	Window step of each observation, for gridded-state models (the trajectory state it is compared with). `None`: every observation is at step 0.
`n_times`	`int \| None`	Number of states in the assimilation window (rollout steps + 1) for the 4D-Var models; required with `time_index`. Static. `None` means 1.

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.PointObsBatch(
...     coords=jnp.array([[0.5], [2.0], [0.0]]),
...     values=jnp.array([1.0, 2.0, 0.0]),
...     sigma=jnp.ones(3),
...     valid=jnp.array([True, True, False]),
...     time_index=jnp.array([0, 3, 0]),
...     n_times=4,
... )
>>> batch.values.shape
(3,)

Padding rows contribute nothing, even when they hold NaN.

>>> padded = vardax.PointObsBatch(
...     coords=jnp.zeros((2, 1)),
...     values=jnp.array([1.0, jnp.nan]),
...     sigma=jnp.ones(2),
...     valid=jnp.array([True, False]),
... )
>>> int(padded.valid.sum())
1

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

class PointObsBatch(eqx.Module):
    r"""Scattered point observations, padded to a fixed count.

    Observation $i$ has coordinates $s_i$, value $y_i$, error standard
    deviation $\sigma_i$ and a validity flag $m_i$; padding rows have
    ``valid=False`` and contribute exactly nothing to the cost or its
    gradient, whatever their other fields hold (``NaN`` included). The
    observation cost is
    $J_o = \tfrac12 \sum_i m_i (y_i - \mathcal H(x)(s_i, t_i))^2 / \sigma_i^2$
    ([`point_obs_cost`][vardax.point_obs_cost]).

    Every array field may carry a leading batch axis; the models then
    return one analysis per sample, and without it a single analysis.

    Attributes:
        coords: Coordinates $s_i$, e.g. ``(x,)``, ``(y, x)`` or
            ``(lon, lat, t)``.
        values: Observed values $y_i$.
        sigma: Per-observation error standard deviations $\sigma_i$.
        valid: Validity flags $m_i$, ``False`` on padding rows.
        time_index: Window step of each observation, for gridded-state
            models (the trajectory state it is compared with). ``None``:
            every observation is at step 0.
        n_times: Number of states in the assimilation window (rollout
            steps + 1) for the 4D-Var models; required with
            ``time_index``. Static. ``None`` means 1.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.PointObsBatch(
        ...     coords=jnp.array([[0.5], [2.0], [0.0]]),
        ...     values=jnp.array([1.0, 2.0, 0.0]),
        ...     sigma=jnp.ones(3),
        ...     valid=jnp.array([True, True, False]),
        ...     time_index=jnp.array([0, 3, 0]),
        ...     n_times=4,
        ... )
        >>> batch.values.shape
        (3,)

        Padding rows contribute nothing, even when they hold ``NaN``.

        >>> padded = vardax.PointObsBatch(
        ...     coords=jnp.zeros((2, 1)),
        ...     values=jnp.array([1.0, jnp.nan]),
        ...     sigma=jnp.ones(2),
        ...     valid=jnp.array([True, False]),
        ... )
        >>> int(padded.valid.sum())
        1
    """

    coords: Float[Array, "*batch N_obs D"]
    values: Float[Array, "*batch N_obs"]
    sigma: Float[Array, "*batch N_obs"]
    valid: Bool[Array, "*batch N_obs"]
    time_index: Int[Array, "*batch N_obs"] | None = None
    n_times: int | None = eqx.field(static=True, default=None)

LSTMState1D ¶

Bases: Module

Hidden state for a 1-D ConvLSTM gradient modulator.

Attributes:

Name	Type	Description
`hidden`	`Float[Array, 'B H_dim N']`	Hidden state tensor.
`cell`	`Float[Array, 'B H_dim N']`	Cell state tensor.

Examples:

>>> import vardax
>>> state = vardax.LSTMState1D.zeros(batch_size=2, hidden_dim=3, seq_len=5)
>>> state.hidden.shape, state.cell.shape
((2, 3, 5), (2, 3, 5))

Both tensors start at zero.

>>> float(abs(state.hidden).sum() + abs(state.cell).sum())
0.0

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

class LSTMState1D(eqx.Module):
    """Hidden state for a 1-D ConvLSTM gradient modulator.

    Attributes:
        hidden: Hidden state tensor.
        cell: Cell state tensor.

    Examples:
        >>> import vardax
        >>> state = vardax.LSTMState1D.zeros(batch_size=2, hidden_dim=3, seq_len=5)
        >>> state.hidden.shape, state.cell.shape
        ((2, 3, 5), (2, 3, 5))

        Both tensors start at zero.

        >>> float(abs(state.hidden).sum() + abs(state.cell).sum())
        0.0
    """

    hidden: Float[Array, "B H_dim N"]
    cell: Float[Array, "B H_dim N"]

    @classmethod
    def zeros(
        cls,
        batch_size: int,
        hidden_dim: int,
        seq_len: int,
    ) -> LSTMState1D:
        """Create a zero-initialised LSTM state.

        Args:
            batch_size: Number of samples $B$.
            hidden_dim: Hidden channel count.
            seq_len: Length $N$ of the 1-D axis.

        Returns:
            A state with zero hidden and cell tensors.
        """
        return cls(
            hidden=jnp.zeros((batch_size, hidden_dim, seq_len)),
            cell=jnp.zeros((batch_size, hidden_dim, seq_len)),
        )

zeros `classmethod` ¶

zeros(
    batch_size: int, hidden_dim: int, seq_len: int
) -> LSTMState1D

Create a zero-initialised LSTM state.

Parameters:

Name	Type	Description	Default
`batch_size`	`int`	Number of samples \(B\).	required
`hidden_dim`	`int`	Hidden channel count.	required
`seq_len`	`int`	Length \(N\) of the 1-D axis.	required

Returns:

Type	Description
`LSTMState1D`	A state with zero hidden and cell tensors.

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

@classmethod
def zeros(
    cls,
    batch_size: int,
    hidden_dim: int,
    seq_len: int,
) -> LSTMState1D:
    """Create a zero-initialised LSTM state.

    Args:
        batch_size: Number of samples $B$.
        hidden_dim: Hidden channel count.
        seq_len: Length $N$ of the 1-D axis.

    Returns:
        A state with zero hidden and cell tensors.
    """
    return cls(
        hidden=jnp.zeros((batch_size, hidden_dim, seq_len)),
        cell=jnp.zeros((batch_size, hidden_dim, seq_len)),
    )

LSTMState2D ¶

Bases: Module

Hidden state for a 2-D ConvLSTM gradient modulator.

Attributes:

Name	Type	Description
`hidden`	`Float[Array, 'B H_dim H W']`	Hidden state tensor.
`cell`	`Float[Array, 'B H_dim H W']`	Cell state tensor.

Examples:

>>> import vardax
>>> state = vardax.LSTMState2D.zeros(1, 4, height=3, width=2)
>>> state.hidden.shape
(1, 4, 3, 2)

It is a pytree with two leaves.

>>> import jax
>>> len(jax.tree_util.tree_leaves(state))
2

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

class LSTMState2D(eqx.Module):
    """Hidden state for a 2-D ConvLSTM gradient modulator.

    Attributes:
        hidden: Hidden state tensor.
        cell: Cell state tensor.

    Examples:
        >>> import vardax
        >>> state = vardax.LSTMState2D.zeros(1, 4, height=3, width=2)
        >>> state.hidden.shape
        (1, 4, 3, 2)

        It is a pytree with two leaves.

        >>> import jax
        >>> len(jax.tree_util.tree_leaves(state))
        2
    """

    hidden: Float[Array, "B H_dim H W"]
    cell: Float[Array, "B H_dim H W"]

    @classmethod
    def zeros(
        cls,
        batch_size: int,
        hidden_dim: int,
        height: int,
        width: int,
    ) -> LSTMState2D:
        """Create a zero-initialised LSTM state.

        Args:
            batch_size: Number of samples $B$.
            hidden_dim: Hidden channel count.
            height: Grid height $H$.
            width: Grid width $W$.

        Returns:
            A state with zero hidden and cell tensors.
        """
        return cls(
            hidden=jnp.zeros((batch_size, hidden_dim, height, width)),
            cell=jnp.zeros((batch_size, hidden_dim, height, width)),
        )

zeros `classmethod` ¶

zeros(
    batch_size: int,
    hidden_dim: int,
    height: int,
    width: int,
) -> LSTMState2D

Create a zero-initialised LSTM state.

Parameters:

Name	Type	Description	Default
`batch_size`	`int`	Number of samples \(B\).	required
`hidden_dim`	`int`	Hidden channel count.	required
`height`	`int`	Grid height \(H\).	required
`width`	`int`	Grid width \(W\).	required

Returns:

Type	Description
`LSTMState2D`	A state with zero hidden and cell tensors.

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

@classmethod
def zeros(
    cls,
    batch_size: int,
    hidden_dim: int,
    height: int,
    width: int,
) -> LSTMState2D:
    """Create a zero-initialised LSTM state.

    Args:
        batch_size: Number of samples $B$.
        hidden_dim: Hidden channel count.
        height: Grid height $H$.
        width: Grid width $W$.

    Returns:
        A state with zero hidden and cell tensors.
    """
    return cls(
        hidden=jnp.zeros((batch_size, hidden_dim, height, width)),
        cell=jnp.zeros((batch_size, hidden_dim, height, width)),
    )

Protocols & Types¶

pipekit-cycle protocols¶

AnalysisStep ¶

ForwardModel ¶

ObservationOperator ¶

vardax protocols¶

Prior ¶

TemporalPrior ¶

CostFunction ¶

GradModulator ¶

PosteriorAdapter ¶

Minimiser ¶

ReducedBasis ¶

StateParameterization ¶

SupportsForwardAdjoint ¶

Batch & state types¶

Batch1D ¶

Batch2D ¶

Batch2DMultivar ¶

PointObsBatch ¶

LSTMState1D ¶

zeros classmethod ¶

LSTMState2D ¶

zeros classmethod ¶

zeros `classmethod` ¶

zeros `classmethod` ¶