Costs, Priors & Solvers

The functional core under the model classes: pure cost functions that score a candidate state against observations and prior, prior modules that supply the regularisation term, and the inner-loop solver functions that drive the 4DVarNet iteration. The model classes are thin, stateful-looking wrappers over these pieces — drop down to this layer when building custom methods or instrumenting the optimisation.

Cost functions

The variational cost \(J(x) = J_\text{obs}(x) + J_\text{prior}(x)\) and its gradient, with the observation and prior terms also available separately (decomposed_loss returns them unsummed for logging). The _1d / _2d suffixes match the Batch1D / Batch2D carriers. See 3DVar and strong-constraint 4DVar for the math each term implements.

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

variational_cost

variational_cost(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> Float[Array, ""]

Compute the variational cost \(U(x)\).

\[ U(x) = \alpha_{obs} J_o + \alpha_{prior} J_p \]

with \(J_o = \operatorname{mean}\big((m \odot (x - y))^2\big)\) and \(J_p = \operatorname{mean}\big((x - \varphi(x))^2\big)\).

Algorithm

1. Form the masked departure \(m \odot (x - y)\) from the batch.
2. Form the prior departure \(x - \varphi(x)\) with prior_fn.
3. Return the weighted sum of the two mean squares.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	Current state estimate.	required
batch	Batch1D	Observed data batch with input (y) and mask (m).	required
prior_fn	Callable[..., Any]	Callable x -> x_prior.	required
alpha_obs	float	Weight for the observation term (default 0.5).	0.5
alpha_prior	float	Weight for the prior term (default 0.5).	0.5

Returns:

Type	Description
Float[Array, '']	Scalar cost value.

Examples:

With the trivial IdentityPrior the prior term vanishes, leaving the weighted observation MSE.

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
>>> x = jnp.ones((1, 2, 4))
>>> float(vardax.variational_cost(x, batch, vardax.IdentityPrior()))
0.5

Setting alpha_obs=0 removes the observation term.

>>> cost = vardax.variational_cost(
...     x, batch, vardax.IdentityPrior(), alpha_obs=0.0
... )
>>> float(cost)
0.0

The gradient has the shape of the state.

>>> import jax
>>> g = jax.grad(vardax.variational_cost)(x, batch, vardax.IdentityPrior())
>>> g.shape
(1, 2, 4)

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

def variational_cost(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> Float[Array, ""]:
    r"""Compute the variational cost $U(x)$.

    $$ U(x) = \alpha_{obs} J_o + \alpha_{prior} J_p $$

    with $J_o = \operatorname{mean}\big((m \odot (x - y))^2\big)$ and
    $J_p = \operatorname{mean}\big((x - \varphi(x))^2\big)$.

    Algorithm:
        1. Form the masked departure $m \odot (x - y)$ from the batch.
        2. Form the prior departure $x - \varphi(x)$ with ``prior_fn``.
        3. Return the weighted sum of the two mean squares.

    Args:
        x: Current state estimate.
        batch: Observed data batch with ``input`` (``y``) and ``mask``
            (``m``).
        prior_fn: Callable ``x -> x_prior``.
        alpha_obs: Weight for the observation term (default ``0.5``).
        alpha_prior: Weight for the prior term (default ``0.5``).

    Returns:
        Scalar cost value.

    Examples:
        With the trivial [`IdentityPrior`][vardax.IdentityPrior] the
        prior term vanishes, leaving the weighted observation MSE.

        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
        >>> x = jnp.ones((1, 2, 4))
        >>> float(vardax.variational_cost(x, batch, vardax.IdentityPrior()))
        0.5

        Setting ``alpha_obs=0`` removes the observation term.

        >>> cost = vardax.variational_cost(
        ...     x, batch, vardax.IdentityPrior(), alpha_obs=0.0
        ... )
        >>> float(cost)
        0.0

        The gradient has the shape of the state.

        >>> import jax
        >>> g = jax.grad(vardax.variational_cost)(x, batch, vardax.IdentityPrior())
        >>> g.shape
        (1, 2, 4)
    """
    obs_diff = batch.mask * (x - batch.input)
    # J_o = mean((m ⊙ (x − y))²)
    observation_term_cost = jnp.mean(obs_diff**2)
    # J_p = mean((x − φ(x))²)
    prior_term_cost = jnp.mean((x - prior_fn(x)) ** 2)
    return alpha_obs * observation_term_cost + alpha_prior * prior_term_cost

variational_cost_grad

variational_cost_grad(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> Float[Array, ...]

Gradient of variational_cost w.r.t. x.

\[ g = \nabla_x U(x) \]

Algorithm

1. Differentiate variational_cost with respect to its first argument x.
2. Return the resulting array.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	Current state estimate.	required
batch	Batch1D	Observed data batch.	required
prior_fn	Callable[..., Any]	Callable x -> x_prior.	required
alpha_obs	float	Weight for the observation term.	0.5
alpha_prior	float	Weight for the prior term.	0.5

Returns:

Type	Description
Float[Array, ...]	The gradient \(\nabla_x U(x)\), shaped like x.

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
>>> x = jnp.ones((1, 2, 4))
>>> g = vardax.variational_cost_grad(x, batch, vardax.IdentityPrior())
>>> g.shape
(1, 2, 4)

It matches jax.grad of the cost.

>>> import jax
>>> ref = jax.grad(vardax.variational_cost)(x, batch, vardax.IdentityPrior())
>>> bool(jnp.allclose(g, ref))
True

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

def variational_cost_grad(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> Float[Array, ...]:
    r"""Gradient of [`variational_cost`][vardax.variational_cost] w.r.t. ``x``.

    $$ g = \nabla_x U(x) $$

    Algorithm:
        1. Differentiate [`variational_cost`][vardax.variational_cost]
           with respect to its first argument ``x``.
        2. Return the resulting array.

    Args:
        x: Current state estimate.
        batch: Observed data batch.
        prior_fn: Callable ``x -> x_prior``.
        alpha_obs: Weight for the observation term.
        alpha_prior: Weight for the prior term.

    Returns:
        The gradient $\nabla_x U(x)$, shaped like ``x``.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
        >>> x = jnp.ones((1, 2, 4))
        >>> g = vardax.variational_cost_grad(x, batch, vardax.IdentityPrior())
        >>> g.shape
        (1, 2, 4)

        It matches ``jax.grad`` of the cost.

        >>> import jax
        >>> ref = jax.grad(vardax.variational_cost)(x, batch, vardax.IdentityPrior())
        >>> bool(jnp.allclose(g, ref))
        True
    """
    return jax.grad(variational_cost)(x, batch, prior_fn, alpha_obs, alpha_prior)

obs_cost_1d

obs_cost_1d(
    state: Float[Array, "B T N"],
    obs: Float[Array, "B T N"],
    mask: Float[Array, "B T N"],
    nan_to_num: bool = False,
) -> Float[Array, ""]

Observation cost for 1-D data.

Computes the masked mean-squared error between the state and observations:

\[ J_{obs} = \operatorname{mean}\big((m \odot (x - y))^2\big) \]

with state \(x\), observations \(y\) and observation mask \(m\) (1 where observed, 0 elsewhere). The mean runs over all elements, so it is the masked squared error divided by the total element count, not by the number of observed locations.

Algorithm

1. Optionally replace NaN entries of obs by zero.
2. Form the masked departure \(m \odot (x - y)\).
3. Return the mean of its square.

Parameters:

Name	Type	Description	Default
state	Float[Array, 'B T N']	Current state estimate \(x\).	required
obs	Float[Array, 'B T N']	Observations \(y\).	required
mask	Float[Array, 'B T N']	Binary observation mask \(m\). A value of 1 indicates an observed location.	required
nan_to_num	bool	When True, replace NaN entries in obs with zero (via :func:jax.numpy.nan_to_num) before masking. Use this for real-world geophysical products (e.g. satellite altimetry) where gaps are stored as NaN rather than an explicit binary mask. Without it, a NaN observation poisons both the cost and its gradient even at masked-out locations, since 0 * NaN is NaN in JAX. Ported from mfourdvar's ObsOperator.loss.	False

Returns:

Type	Description
Float[Array, '']	Scalar observation cost.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import obs_cost_1d
>>> state = jnp.ones((1, 1, 4))
>>> obs = jnp.zeros((1, 1, 4))
>>> mask = jnp.ones((1, 1, 4))
>>> float(obs_cost_1d(state, obs, mask))
1.0

Masked-out locations do not contribute, even when the observation there is NaN (with nan_to_num=True).

>>> obs = jnp.array([[[0.0, jnp.nan, 0.0, 0.0]]])
>>> mask = jnp.array([[[1.0, 0.0, 1.0, 1.0]]])
>>> float(obs_cost_1d(state, obs, mask, nan_to_num=True))
0.75

The cost is differentiable in the state.

>>> import jax
>>> jax.grad(obs_cost_1d)(state, jnp.zeros((1, 1, 4)), mask).shape
(1, 1, 4)

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

def obs_cost_1d(
    state: Float[Array, "B T N"],
    obs: Float[Array, "B T N"],
    mask: Float[Array, "B T N"],
    nan_to_num: bool = False,
) -> Float[Array, ""]:
    r"""Observation cost for 1-D data.

    Computes the masked mean-squared error between the state and observations:

    $$ J_{obs} = \operatorname{mean}\big((m \odot (x - y))^2\big) $$

    with state $x$, observations $y$ and observation mask $m$ (``1`` where
    observed, ``0`` elsewhere). The mean runs over all elements, so it is
    the masked squared error divided by the total element count, not by
    the number of observed locations.

    Algorithm:
        1. Optionally replace ``NaN`` entries of ``obs`` by zero.
        2. Form the masked departure $m \odot (x - y)$.
        3. Return the mean of its square.

    Args:
        state: Current state estimate $x$.
        obs: Observations $y$.
        mask: Binary observation mask $m$.
            A value of ``1`` indicates an observed location.
        nan_to_num: When ``True``, replace ``NaN`` entries in ``obs`` with
            zero (via :func:`jax.numpy.nan_to_num`) *before* masking. Use this
            for real-world geophysical products (e.g. satellite altimetry)
            where gaps are stored as ``NaN`` rather than an explicit binary
            mask. Without it, a ``NaN`` observation poisons both the cost and
            its gradient even at masked-out locations, since ``0 * NaN`` is
            ``NaN`` in JAX. Ported from ``mfourdvar``'s ``ObsOperator.loss``.

    Returns:
        Scalar observation cost.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import obs_cost_1d
        >>> state = jnp.ones((1, 1, 4))
        >>> obs = jnp.zeros((1, 1, 4))
        >>> mask = jnp.ones((1, 1, 4))
        >>> float(obs_cost_1d(state, obs, mask))
        1.0

        Masked-out locations do not contribute, even when the observation
        there is ``NaN`` (with ``nan_to_num=True``).

        >>> obs = jnp.array([[[0.0, jnp.nan, 0.0, 0.0]]])
        >>> mask = jnp.array([[[1.0, 0.0, 1.0, 1.0]]])
        >>> float(obs_cost_1d(state, obs, mask, nan_to_num=True))
        0.75

        The cost is differentiable in the state.

        >>> import jax
        >>> jax.grad(obs_cost_1d)(state, jnp.zeros((1, 1, 4)), mask).shape
        (1, 1, 4)
    """
    if nan_to_num:
        obs = jnp.nan_to_num(obs)
    diff = mask * (state - obs)
    return jnp.mean(diff**2)

obs_cost_2d

obs_cost_2d(
    state: Float[Array, "B T H W"],
    obs: Float[Array, "B T H W"],
    mask: Float[Array, "B T H W"],
    nan_to_num: bool = False,
) -> Float[Array, ""]

Observation cost for 2-D data.

Computes the masked mean-squared error between the state and observations:

\[ J_{obs} = \operatorname{mean}\big((m \odot (x - y))^2\big) \]

with state \(x\), observations \(y\) and observation mask \(m\) (1 where observed, 0 elsewhere), as in obs_cost_1d.

Algorithm

1. Optionally replace NaN entries of obs by zero.
2. Form the masked departure \(m \odot (x - y)\).
3. Return the mean of its square.

Parameters:

Name	Type	Description	Default
state	Float[Array, 'B T H W']	Current state estimate \(x\).	required
obs	Float[Array, 'B T H W']	Observations \(y\).	required
mask	Float[Array, 'B T H W']	Binary observation mask \(m\).	required
nan_to_num	bool	When True, replace NaN entries in obs with zero before masking (see obs_cost_1d for the rationale).	False

Returns:

Type	Description
Float[Array, '']	Scalar observation cost.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import obs_cost_2d
>>> state = jnp.ones((1, 1, 2, 2))
>>> obs = jnp.zeros((1, 1, 2, 2))
>>> float(obs_cost_2d(state, obs, jnp.ones((1, 1, 2, 2))))
1.0

With nothing observed the cost is zero.

>>> float(obs_cost_2d(state, obs, jnp.zeros((1, 1, 2, 2))))
0.0

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

def obs_cost_2d(
    state: Float[Array, "B T H W"],
    obs: Float[Array, "B T H W"],
    mask: Float[Array, "B T H W"],
    nan_to_num: bool = False,
) -> Float[Array, ""]:
    r"""Observation cost for 2-D data.

    Computes the masked mean-squared error between the state and observations:

    $$ J_{obs} = \operatorname{mean}\big((m \odot (x - y))^2\big) $$

    with state $x$, observations $y$ and observation mask $m$ (``1`` where
    observed, ``0`` elsewhere), as in [`obs_cost_1d`][vardax.obs_cost_1d].

    Algorithm:
        1. Optionally replace ``NaN`` entries of ``obs`` by zero.
        2. Form the masked departure $m \odot (x - y)$.
        3. Return the mean of its square.

    Args:
        state: Current state estimate $x$.
        obs: Observations $y$.
        mask: Binary observation mask $m$.
        nan_to_num: When ``True``, replace ``NaN`` entries in ``obs`` with
            zero before masking (see [`obs_cost_1d`][vardax.obs_cost_1d]
            for the rationale).

    Returns:
        Scalar observation cost.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import obs_cost_2d
        >>> state = jnp.ones((1, 1, 2, 2))
        >>> obs = jnp.zeros((1, 1, 2, 2))
        >>> float(obs_cost_2d(state, obs, jnp.ones((1, 1, 2, 2))))
        1.0

        With nothing observed the cost is zero.

        >>> float(obs_cost_2d(state, obs, jnp.zeros((1, 1, 2, 2))))
        0.0
    """
    if nan_to_num:
        obs = jnp.nan_to_num(obs)
    diff = mask * (state - obs)
    return jnp.mean(diff**2)

prior_cost

prior_cost(
    state: Float[Array, ...],
    prior_reconstruction: Float[Array, ...],
) -> Float[Array, ""]

Prior cost based on learned autoencoder reconstruction.

Computes the mean-squared error between the state and its reconstruction through the learned prior (autoencoder):

\[ J_{prior} = \operatorname{mean}\big((x - \varphi(x))^2\big) \]

Algorithm

1. Take the reconstruction \(\varphi(x)\) of the state.
2. Return the mean squared difference \(x - \varphi(x)\).

Parameters:

Name	Type	Description	Default
state	Float[Array, ...]	Current state estimate \(x\).	required
prior_reconstruction	Float[Array, ...]	Autoencoder reconstruction \(\varphi(x)\) of the state.	required

Returns:

Type	Description
Float[Array, '']	Scalar prior cost.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import prior_cost
>>> x = jnp.ones((2, 3))
>>> float(prior_cost(x, x))
0.0

A reconstruction that is off by two everywhere costs four.

>>> float(prior_cost(x, x + 2.0))
4.0

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

def prior_cost(
    state: Float[Array, ...],
    prior_reconstruction: Float[Array, ...],
) -> Float[Array, ""]:
    r"""Prior cost based on learned autoencoder reconstruction.

    Computes the mean-squared error between the state and its reconstruction
    through the learned prior (autoencoder):

    $$ J_{prior} = \operatorname{mean}\big((x - \varphi(x))^2\big) $$

    Algorithm:
        1. Take the reconstruction $\varphi(x)$ of the state.
        2. Return the mean squared difference $x - \varphi(x)$.

    Args:
        state: Current state estimate $x$.
        prior_reconstruction: Autoencoder reconstruction $\varphi(x)$ of
            the state.

    Returns:
        Scalar prior cost.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import prior_cost
        >>> x = jnp.ones((2, 3))
        >>> float(prior_cost(x, x))
        0.0

        A reconstruction that is off by two everywhere costs four.

        >>> float(prior_cost(x, x + 2.0))
        4.0
    """
    return jnp.mean((state - prior_reconstruction) ** 2)

decomposed_loss

decomposed_loss(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> dict[str, Float[Array, ""]]

Compute the decomposed variational loss.

Returns individual observation and prior components alongside the total, matching the ModelLoss pattern from the legacy codebase.

\[ U = \alpha_{obs} J_o + \alpha_{prior} J_p \]

with \(J_o = \operatorname{mean}\big((m \odot (x - y))^2\big)\) and \(J_p = \operatorname{mean}\big((x - \varphi(x))^2\big)\) as in variational_cost.

Algorithm

1. Compute the weighted observation term.
2. Compute the weighted prior term.
3. Return both and their sum.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	Current state estimate.	required
batch	Batch1D	Observed data batch.	required
prior_fn	Callable[..., Any]	Callable x -> x_prior.	required
alpha_obs	float	Weight for the observation term.	0.5
alpha_prior	float	Weight for the prior term.	0.5

Returns:

Type	Description
dict[str, Float[Array, '']]	Dictionary with keys "obs", "prior", and "total".

Examples:

>>> import jax.numpy as jnp, vardax
>>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
>>> x = jnp.ones((1, 2, 4))
>>> parts = vardax.decomposed_loss(x, batch, vardax.IdentityPrior())
>>> sorted(parts)
['obs', 'prior', 'total']
>>> float(parts["obs"]), float(parts["prior"])
(0.5, 0.0)

The total equals variational_cost.

>>> total = vardax.variational_cost(x, batch, vardax.IdentityPrior())
>>> bool(jnp.allclose(parts["total"], total))
True

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

def decomposed_loss(
    x: Float[Array, ...],
    batch: Batch1D,
    prior_fn: Callable[..., Any],
    alpha_obs: float = 0.5,
    alpha_prior: float = 0.5,
) -> dict[str, Float[Array, ""]]:
    r"""Compute the decomposed variational loss.

    Returns individual observation and prior components alongside the
    total, matching the ``ModelLoss`` pattern from the legacy codebase.

    $$ U = \alpha_{obs} J_o + \alpha_{prior} J_p $$

    with $J_o = \operatorname{mean}\big((m \odot (x - y))^2\big)$ and
    $J_p = \operatorname{mean}\big((x - \varphi(x))^2\big)$ as in
    [`variational_cost`][vardax.variational_cost].

    Algorithm:
        1. Compute the weighted observation term.
        2. Compute the weighted prior term.
        3. Return both and their sum.

    Args:
        x: Current state estimate.
        batch: Observed data batch.
        prior_fn: Callable ``x -> x_prior``.
        alpha_obs: Weight for the observation term.
        alpha_prior: Weight for the prior term.

    Returns:
        Dictionary with keys ``"obs"``, ``"prior"``, and ``"total"``.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> batch = vardax.Batch1D(input=jnp.zeros((1, 2, 4)), mask=jnp.ones((1, 2, 4)))
        >>> x = jnp.ones((1, 2, 4))
        >>> parts = vardax.decomposed_loss(x, batch, vardax.IdentityPrior())
        >>> sorted(parts)
        ['obs', 'prior', 'total']
        >>> float(parts["obs"]), float(parts["prior"])
        (0.5, 0.0)

        The total equals [`variational_cost`][vardax.variational_cost].

        >>> total = vardax.variational_cost(x, batch, vardax.IdentityPrior())
        >>> bool(jnp.allclose(parts["total"], total))
        True
    """
    obs_diff = batch.mask * (x - batch.input)
    obs = alpha_obs * jnp.mean(obs_diff**2)
    prior = alpha_prior * jnp.mean((x - prior_fn(x)) ** 2)
    return {"obs": obs, "prior": prior, "total": obs + prior}

strong_variational_cost

strong_variational_cost(
    initial_state: Float[Array, ...],
    ts: Float[Array, " T"],
    batch: Batch1D,
    forward_fn: Callable[..., Any],
    *,
    background: Float[Array, ...] | None = None,
    alpha_obs: float = 0.5,
    alpha_bg: float = 0.5,
    nan_to_num: bool = False,
) -> Float[Array, ""]

Strong-constraint variational cost \(U(x_0)\).

\[ U(x_0) = \alpha_{obs} J_o + \alpha_{bg} J_b \]

with \(J_o = \operatorname{mean}\big((m \odot (\varphi(x_0) - y))^2\big)\) and \(J_b = \operatorname{mean}\big((x_0 - x_b)^2\big)\).

The dynamical model is enforced as a hard constraint: the initial state initial_state is propagated through the dynamics forward_fn(initial_state, ts) and only the resulting trajectory is scored against the observations. This differs from the weak / soft-constraint variational_cost, where the model appears as an additive prior penalty and the whole state field is free.

Ported from mfourdvar's StrongVarCost, adapted to a functional form that mirrors variational_cost.

Algorithm

1. Propagate \(x_0\) through forward_fn(initial_state, ts) to get the trajectory \(\varphi(x_0)\).
2. Score the trajectory against the observations with the mask.
3. Add the background term, defaulting to zero when background is omitted.
4. Return the weighted sum.

Parameters:

Name	Type	Description	Default
initial_state	Float[Array, ...]	Initial state \(x_0\) (the control variable), whatever forward_fn expects.	required
ts	Float[Array, ' T']	Time coordinates \(t\) passed to forward_fn.	required
batch	Batch1D	Observed data batch supplying input (y) and mask (m), matching forward_fn's trajectory output.	required
forward_fn	Callable[..., Any]	Callable (initial_state, ts) -> trajectory producing a state trajectory the same shape as batch.input. A DynTrajectory (optionally vmap-ed over the batch axis) is the canonical choice.	required
background	Float[Array, ...] | None	Background state for the background term. Defaults to initial_state (i.e. no background penalty).	None
alpha_obs	float	Weight for the observation term (default 0.5).	0.5
alpha_bg	float	Weight for the background term (default 0.5).	0.5
nan_to_num	bool	When True, replace NaN entries in the observations with zero before masking (see obs_cost_1d).	False

Returns:

Type	Description
Float[Array, '']	Scalar strong-constraint variational cost.

Examples:

A persistence model repeats the initial state over the window.

>>> import jax.numpy as jnp, vardax
>>> ts = jnp.arange(3.0)
>>> forward = lambda x0, ts: jnp.broadcast_to(x0[:, None], (1, 3, 4))
>>> batch = vardax.Batch1D(input=jnp.zeros((1, 3, 4)), mask=jnp.ones((1, 3, 4)))
>>> x0 = jnp.ones((1, 4))
>>> float(vardax.strong_variational_cost(x0, ts, batch, forward))
0.5

A background away from initial_state adds the background term.

>>> cost = vardax.strong_variational_cost(
...     x0, ts, batch, forward, background=x0 - 1.0
... )
>>> float(cost)
1.0

The cost is differentiable in the control.

>>> import jax
>>> g = jax.grad(vardax.strong_variational_cost)(x0, ts, batch, forward)
>>> g.shape
(1, 4)

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

def strong_variational_cost(
    initial_state: Float[Array, ...],
    ts: Float[Array, " T"],
    batch: Batch1D,
    forward_fn: Callable[..., Any],
    *,
    background: Float[Array, ...] | None = None,
    alpha_obs: float = 0.5,
    alpha_bg: float = 0.5,
    nan_to_num: bool = False,
) -> Float[Array, ""]:
    r"""Strong-constraint variational cost $U(x_0)$.

    $$ U(x_0) = \alpha_{obs} J_o + \alpha_{bg} J_b $$

    with $J_o = \operatorname{mean}\big((m \odot (\varphi(x_0) - y))^2\big)$
    and $J_b = \operatorname{mean}\big((x_0 - x_b)^2\big)$.

    The dynamical model is enforced as a **hard** constraint: the initial
    state ``initial_state`` is propagated through the dynamics
    ``forward_fn(initial_state, ts)``
    and only the resulting trajectory is scored against the observations.
    This differs from the weak / soft-constraint
    [`variational_cost`][vardax.variational_cost], where the model appears
    as an *additive* prior penalty and the whole state field is free.

    Ported from ``mfourdvar``'s ``StrongVarCost``, adapted to a functional
    form that mirrors [`variational_cost`][vardax.variational_cost].

    Algorithm:
        1. Propagate $x_0$ through ``forward_fn(initial_state, ts)`` to get the
           trajectory $\varphi(x_0)$.
        2. Score the trajectory against the observations with the mask.
        3. Add the background term, defaulting to zero when ``background`` is
           omitted.
        4. Return the weighted sum.

    Args:
        initial_state: Initial state $x_0$ (the control variable), whatever
            ``forward_fn`` expects.
        ts: Time coordinates $t$ passed to ``forward_fn``.
        batch: Observed data batch supplying ``input`` (``y``) and ``mask``
            (``m``), matching ``forward_fn``'s trajectory output.
        forward_fn: Callable ``(initial_state, ts) -> trajectory`` producing a state
            trajectory the same shape as ``batch.input``. A
            [`DynTrajectory`][vardax.DynTrajectory] (optionally ``vmap``-ed
            over the batch axis) is the canonical choice.
        background: Background state for the background term. Defaults to
            ``initial_state``
            (i.e. no background penalty).
        alpha_obs: Weight for the observation term (default ``0.5``).
        alpha_bg: Weight for the background term (default ``0.5``).
        nan_to_num: When ``True``, replace ``NaN`` entries in the
            observations with zero before masking (see
            [`obs_cost_1d`][vardax.obs_cost_1d]).

    Returns:
        Scalar strong-constraint variational cost.

    Examples:
        A persistence model repeats the initial state over the window.

        >>> import jax.numpy as jnp, vardax
        >>> ts = jnp.arange(3.0)
        >>> forward = lambda x0, ts: jnp.broadcast_to(x0[:, None], (1, 3, 4))
        >>> batch = vardax.Batch1D(input=jnp.zeros((1, 3, 4)), mask=jnp.ones((1, 3, 4)))
        >>> x0 = jnp.ones((1, 4))
        >>> float(vardax.strong_variational_cost(x0, ts, batch, forward))
        0.5

        A background away from ``initial_state`` adds the background term.

        >>> cost = vardax.strong_variational_cost(
        ...     x0, ts, batch, forward, background=x0 - 1.0
        ... )
        >>> float(cost)
        1.0

        The cost is differentiable in the control.

        >>> import jax
        >>> g = jax.grad(vardax.strong_variational_cost)(x0, ts, batch, forward)
        >>> g.shape
        (1, 4)
    """
    if background is None:
        background = initial_state
    trajectory = forward_fn(initial_state, ts)
    obs = batch.input
    if nan_to_num:
        obs = jnp.nan_to_num(obs)
    # J_o = mean((m ⊙ (φ(x₀) − y))²)
    observation_term_cost = jnp.mean((batch.mask * (trajectory - obs)) ** 2)
    # J_b = mean((x₀ − x_b)²)
    background_term_cost = background_cost(initial_state, background)
    return alpha_obs * observation_term_cost + alpha_bg * background_term_cost

background_cost

background_cost(
    initial_state: Float[Array, ...],
    background: Float[Array, ...],
) -> Float[Array, ""]

Background cost \(\operatorname{mean}\big((x_0 - x_b)^2\big)\).

Penalises departure of the initial state \(x_0\) (the control variable in strong-constraint 4DVar) from the background / first guess \(x_b\). Uses the same mean-squared convention as the other functional costs in this module.

\[ J_b = \operatorname{mean}\big((x_0 - x_b)^2\big) \]

Algorithm

1. Form the departure \(x_0 - x_b\).
2. Return the mean of its square.

Parameters:

Name	Type	Description	Default
initial_state	Float[Array, ...]	Initial state estimate \(x_0\).	required
background	Float[Array, ...]	Background state \(x_b\).	required

Returns:

Type	Description
Float[Array, '']	Scalar background cost.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import background_cost
>>> x0 = jnp.ones(4)
>>> float(background_cost(x0, x0))
0.0

A uniform departure of 0.5 costs a quarter.

>>> float(background_cost(x0, x0 - 0.5))
0.25

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

def background_cost(
    initial_state: Float[Array, ...],
    background: Float[Array, ...],
) -> Float[Array, ""]:
    r"""Background cost $\operatorname{mean}\big((x_0 - x_b)^2\big)$.

    Penalises departure of the initial state $x_0$ (the control variable
    in strong-constraint 4DVar) from the background / first guess $x_b$.
    Uses the same mean-squared convention as the other functional costs in
    this module.

    $$ J_b = \operatorname{mean}\big((x_0 - x_b)^2\big) $$

    Algorithm:
        1. Form the departure $x_0 - x_b$.
        2. Return the mean of its square.

    Args:
        initial_state: Initial state estimate $x_0$.
        background: Background state $x_b$.

    Returns:
        Scalar background cost.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import background_cost
        >>> x0 = jnp.ones(4)
        >>> float(background_cost(x0, x0))
        0.0

        A uniform departure of 0.5 costs a quarter.

        >>> float(background_cost(x0, x0 - 0.5))
        0.25
    """
    return jnp.mean((initial_state - background) ** 2)

Physics residual penalty

ResidualPenalty adds a PDE as a weak constraint on a continuous field (a neural field, or any differentiable callable of one point): the residual \(\mathcal N[f]\) is evaluated by autodiff at collocation points, fixed or redrawn from a key on every call, and scored as \(\frac{\lambda}{2 N_c}\sum_j \lVert \mathcal N[f](c_j) \rVert^2\). Residuals are plain callables (field, s) -> (R,) built on geonnax.diffops; vardax ships two references in vardax.examples, 2-D advection–diffusion and 2-D incompressible Navier–Stokes in vorticity form (the field gives the streamfunction \(\psi\), with \(u = \nabla^\perp\psi\) and \(\omega = \Delta\psi\)).

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

ResidualPenalty

Bases: Module

PDE residual of a continuous field at collocation points, as a cost.

For a field \(f_\theta\) (a callable of one point, e.g. a neural field), a PDE operator \(\mathcal N\) supplied as residual_fn and collocation points \(c_1, \dots, c_{N_c}\),

\[ J_{\text{res}}(\theta) = \frac{\lambda}{2 N_c} \sum_{j=1}^{N_c} \big\lVert \mathcal N[f_\theta](c_j) \big\rVert^2 . \]

Added to an observation cost it imposes the physics as a weak constraint, without a time-stepping model. Residuals are callables built on geonnax.diffops (see vardax.examples for advection–diffusion and vorticity-form Navier–Stokes); vardax does not import a PDE library.

With a sampler, the points are redrawn from a key on every call. If each point is drawn from a distribution \(p\), the resampled cost is an unbiased estimate of \(\tfrac{\lambda}{2}\,\mathbb E_{c \sim p} \lVert \mathcal N[f_\theta](c) \rVert^2\), so its gradient is a stochastic gradient of that expected penalty.

The field must be smooth enough for the derivatives the residual takes (twice differentiable for advection–diffusion, four times for the streamfunction form of Navier–Stokes) at every collocation point. A raw jnp.linalg.norm or jnp.sqrt of a quantity that vanishes at a point gives NaN gradients there; guard it with the double-where idiom (jnp.where(sq > 0, jnp.sqrt(jnp.where(sq > 0, sq, 1.0)), 0.0)).

residual_fn and sampler are static fields, so they must be hashable (functions and closures are); collocation and weight are array leaves, so changing the weight (e.g. annealing \(\lambda\)) does not recompile a jitted cost.

Algorithm

1. Obtain the points \(c_j\), from collocation or by calling sampler(key).
2. Evaluate the residual \(\mathcal N[f_\theta](c_j)\) at every point.
3. Return \(\lambda / (2 N_c)\) times the sum of squared residuals.

Raises:

Type	Description
MissingInputError	If neither collocation nor sampler is given, or if a sampler is set and a call omits key.
ShapeError	If collocation is not (N_c, D) or weight is not a scalar.
ConfigurationError	If weight is negative.

Attributes:

Name	Type	Description
residual_fn	Callable[..., Float[Array, ' R']]	(field, s) -> (R,), the residual \(\mathcal N[f](s)\) at one point s of shape (D,), where field maps (D,) -> (K,). A scalar return is treated as (1,).
collocation	Float[Array, 'N D'] | None	Fixed collocation points of shape (N_c, D). May be None when sampler is given.
weight	Float[Array, '']	Penalty weight \(\lambda \ge 0\), stored as a 0-d array.
sampler	Callable[[PRNGKeyArray], Float[Array, 'N D']] | None	Optional key -> (N_c, D); when given, the points are resampled on every call and collocation is not used.

Examples:

A residual that is the field itself gives \(J = \tfrac{\lambda}{2} \lVert f \rVert^2\) for a constant field:

>>> import jax.numpy as jnp
>>> import vardax
>>> pen = vardax.ResidualPenalty(
...     residual_fn=lambda f, s: f(s),
...     collocation=jnp.zeros((5, 3)),
...     weight=0.5,
... )
>>> float(pen(lambda s: jnp.array([1.0, 2.0])))
1.25

With a sampler, pass a key on every call:

>>> import jax
>>> pen = vardax.ResidualPenalty(
...     residual_fn=lambda f, s: f(s),
...     collocation=None,
...     weight=1.0,
...     sampler=lambda key: jax.random.uniform(key, (64, 3)),
... )
>>> float(pen(lambda s: s[:1], key=jax.random.PRNGKey(0))) > 0
True

A penalty with neither collocation points nor a sampler is rejected.

>>> vardax.ResidualPenalty(
...     residual_fn=lambda f, s: f(s), collocation=None, weight=1.0
... )
Traceback (most recent call last):
...
MissingInputError: ResidualPenalty needs collocation points or a sampler.

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

class ResidualPenalty(eqx.Module):
    r"""PDE residual of a continuous field at collocation points, as a cost.

    For a field $f_\theta$ (a callable of one point, e.g. a neural field),
    a PDE operator $\mathcal N$ supplied as ``residual_fn`` and collocation
    points $c_1, \dots, c_{N_c}$,

    $$
    J_{\text{res}}(\theta) = \frac{\lambda}{2 N_c} \sum_{j=1}^{N_c}
        \big\lVert \mathcal N[f_\theta](c_j) \big\rVert^2 .
    $$

    Added to an observation cost it imposes the physics as a weak
    constraint, without a time-stepping model. Residuals are callables built
    on `geonnax.diffops` (see `vardax.examples` for advection–diffusion and
    vorticity-form Navier–Stokes); vardax does not import a PDE library.

    With a ``sampler``, the points are redrawn from a key on every call. If
    each point is drawn from a distribution $p$, the resampled cost is an
    unbiased estimate of
    $\tfrac{\lambda}{2}\,\mathbb E_{c \sim p} \lVert \mathcal N[f_\theta](c) \rVert^2$,
    so its gradient is a stochastic gradient of that expected penalty.

    The field must be smooth enough for the derivatives the residual takes
    (twice differentiable for advection–diffusion, four times for the
    streamfunction form of Navier–Stokes) at every collocation point. A raw
    `jnp.linalg.norm` or `jnp.sqrt` of a quantity that vanishes at a point
    gives NaN gradients there; guard it with the double-`where` idiom
    (`jnp.where(sq > 0, jnp.sqrt(jnp.where(sq > 0, sq, 1.0)), 0.0)`).

    ``residual_fn`` and ``sampler`` are static fields, so they must be
    hashable (functions and closures are); ``collocation`` and ``weight``
    are array leaves, so changing the weight (e.g. annealing $\lambda$)
    does not recompile a jitted cost.

    Algorithm:
        1. Obtain the points $c_j$, from ``collocation`` or by calling
           ``sampler(key)``.
        2. Evaluate the residual $\mathcal N[f_\theta](c_j)$ at every point.
        3. Return $\lambda / (2 N_c)$ times the sum of squared residuals.

    Raises:
        MissingInputError: If neither ``collocation`` nor ``sampler`` is
            given, or if a sampler is set and a call omits ``key``.
        ShapeError: If ``collocation`` is not ``(N_c, D)`` or ``weight`` is
            not a scalar.
        ConfigurationError: If ``weight`` is negative.

    Attributes:
        residual_fn: ``(field, s) -> (R,)``, the residual $\mathcal N[f](s)$
            at one point ``s`` of shape ``(D,)``, where ``field`` maps
            ``(D,) -> (K,)``. A scalar return is treated as ``(1,)``.
        collocation: Fixed collocation points of shape ``(N_c, D)``. May be
            ``None`` when ``sampler`` is given.
        weight: Penalty weight $\lambda \ge 0$, stored as a 0-d array.
        sampler: Optional ``key -> (N_c, D)``; when given, the points are
            resampled on every call and ``collocation`` is not used.

    Examples:
        A residual that is the field itself gives
        $J = \tfrac{\lambda}{2} \lVert f \rVert^2$ for a constant field:

        >>> import jax.numpy as jnp
        >>> import vardax
        >>> pen = vardax.ResidualPenalty(
        ...     residual_fn=lambda f, s: f(s),
        ...     collocation=jnp.zeros((5, 3)),
        ...     weight=0.5,
        ... )
        >>> float(pen(lambda s: jnp.array([1.0, 2.0])))
        1.25

        With a sampler, pass a key on every call:

        >>> import jax
        >>> pen = vardax.ResidualPenalty(
        ...     residual_fn=lambda f, s: f(s),
        ...     collocation=None,
        ...     weight=1.0,
        ...     sampler=lambda key: jax.random.uniform(key, (64, 3)),
        ... )
        >>> float(pen(lambda s: s[:1], key=jax.random.PRNGKey(0))) > 0
        True

        A penalty with neither collocation points nor a sampler is rejected.

        >>> vardax.ResidualPenalty(  # doctest: +IGNORE_EXCEPTION_DETAIL
        ...     residual_fn=lambda f, s: f(s), collocation=None, weight=1.0
        ... )
        Traceback (most recent call last):
        ...
        MissingInputError: ResidualPenalty needs collocation points or a sampler.
    """

    residual_fn: Callable[..., Float[Array, " R"]] = eqx.field(static=True)
    collocation: Float[Array, "N D"] | None
    weight: Float[Array, ""] = eqx.field(converter=jnp.asarray)
    sampler: Callable[[PRNGKeyArray], Float[Array, "N D"]] | None = eqx.field(
        static=True, default=None
    )

    def __check_init__(self) -> None:
        if self.collocation is None and self.sampler is None:
            raise MissingInputError(
                "ResidualPenalty needs collocation points or a sampler."
            )
        if self.collocation is not None and jnp.ndim(self.collocation) != 2:
            raise ShapeError(
                "collocation must have shape (N_c, D); got shape "
                f"{jnp.shape(self.collocation)}."
            )
        if jnp.ndim(self.weight) != 0:
            raise ShapeError(
                f"weight must be a scalar; got shape {jnp.shape(self.weight)}."
            )
        try:
            value = float(self.weight)
        except jax.errors.ConcretizationTypeError:
            return  # built under a transform: the sign is not known yet
        if value < 0:
            raise ConfigurationError(f"weight must be non-negative; got {value}.")

    def __call__(
        self,
        field: Callable[[Float[Array, " D"]], Float[Array, " K"]],
        *,
        key: PRNGKeyArray | None = None,
    ) -> Float[Array, ""]:
        """Evaluate the penalty for ``field``.

        Args:
            field: The field, a callable of one point ``(D,) -> (K,)``.
            key: PRNG key for the sampler. Required when ``sampler`` is set;
                ignored otherwise.

        Returns:
            The scalar penalty $J_{\text{res}}$.

        Raises:
            MissingInputError: If a sampler is set and ``key`` is ``None``.
            ShapeError: If the points are not of shape ``(N_c, D)``.
        """
        if self.sampler is not None:
            if key is None:
                raise MissingInputError("ResidualPenalty with a sampler needs a key.")
            points = self.sampler(key)
        else:
            points = self.collocation
        assert points is not None  # guaranteed by __check_init__
        if points.ndim != 2:
            raise ShapeError(
                f"collocation points must have shape (N_c, D); got {points.shape}."
            )

        # __check_init__ skips the sign check when the weight is traced
        weight = eqx.error_if(
            self.weight, self.weight < 0, "weight must be non-negative"
        )
        # N[f](c_j) for every point: (N_c, R)
        residuals = jax.vmap(
            lambda point: jnp.atleast_1d(self.residual_fn(field, point))
        )(points)
        # J = λ/(2 N_c) Σⱼ ‖N[f](c_j)‖²: ()
        return weight / (2 * points.shape[0]) * jnp.sum(residuals**2)

Public vardax.examples API surface — demo helpers and reference residuals.

Re-exports from vardax._src.examples (the notebook training loop, not part of the library API proper; see that module's docstring) and from vardax._src.utils.residuals (reference PDE residuals for ResidualPenalty).

advection_diffusion_2d

advection_diffusion_2d(
    velocity: tuple[float, float] | Float[Array, " 2"],
    diffusivity: float,
) -> Residual

Residual of 2-D advection–diffusion with constant coefficients.

For a scalar tracer \(u(t, x, y)\) advected by a constant velocity \(c = (c_x, c_y)\) and diffused with diffusivity \(\kappa \ge 0\),

\[ \mathcal N[u] = \partial_t u + c \cdot \nabla u - \kappa\, \Delta u , \]

with \(\nabla\) and \(\Delta\) acting on \((x, y)\) only.

Parameters:

Name	Type	Description	Default
velocity	tuple[float, float] | Float[Array, ' 2']	Constant advecting velocity \((c_x, c_y)\).	required
diffusivity	float	Diffusivity \(\kappa \ge 0\).	required

The field must be twice differentiable in \((x, y)\) and once in \(t\) at every collocation point; a raw jnp.sqrt or jnp.linalg.norm that vanishes there gives NaN derivatives, so guard it with the double-where idiom.

Returns:

Type	Description
Residual	residual_fn(field, s) -> (1,), where s = (t, x, y) and the
Residual	tracer is field(s)[0].

Raises:

Type	Description
ValueError	If velocity does not have two components or diffusivity is negative.

Examples:

A Gaussian bump translated by \(c\) and spread by diffusion solves the equation exactly:

>>> import jax.numpy as jnp
>>> from vardax.examples import advection_diffusion_2d
>>> c, kappa = jnp.array([1.0, 0.5]), 0.1
>>> def bump(s):
...     var = 0.04 + 2 * kappa * s[0]
...     r2 = jnp.sum((s[1:] - c * s[0]) ** 2)
...     return jnp.array([0.04 / var * jnp.exp(-r2 / (2 * var))])
>>> res = advection_diffusion_2d(c, kappa)
>>> r = res(bump, jnp.array([0.3, 0.25, 0.2]))
>>> r.shape, bool(jnp.abs(r[0]) < 1e-4)
((1,), True)

Source code in src/vardax/_src/utils/residuals.py

def advection_diffusion_2d(
    velocity: tuple[float, float] | Float[Array, " 2"],
    diffusivity: float,
) -> Residual:
    r"""Residual of 2-D advection–diffusion with constant coefficients.

    For a scalar tracer $u(t, x, y)$ advected by a constant velocity
    $c = (c_x, c_y)$ and diffused with diffusivity $\kappa \ge 0$,

    $$
    \mathcal N[u] = \partial_t u + c \cdot \nabla u - \kappa\, \Delta u ,
    $$

    with $\nabla$ and $\Delta$ acting on $(x, y)$ only.

    Args:
        velocity: Constant advecting velocity $(c_x, c_y)$.
        diffusivity: Diffusivity $\kappa \ge 0$.

    The field must be twice differentiable in $(x, y)$ and once in $t$ at
    every collocation point; a raw `jnp.sqrt` or `jnp.linalg.norm` that
    vanishes there gives NaN derivatives, so guard it with the double-`where`
    idiom.

    Returns:
        ``residual_fn(field, s) -> (1,)``, where ``s = (t, x, y)`` and the
        tracer is ``field(s)[0]``.

    Raises:
        ValueError: If ``velocity`` does not have two components or
            ``diffusivity`` is negative.

    Examples:
        A Gaussian bump translated by $c$ and spread by diffusion solves the
        equation exactly:

        >>> import jax.numpy as jnp
        >>> from vardax.examples import advection_diffusion_2d
        >>> c, kappa = jnp.array([1.0, 0.5]), 0.1
        >>> def bump(s):
        ...     var = 0.04 + 2 * kappa * s[0]
        ...     r2 = jnp.sum((s[1:] - c * s[0]) ** 2)
        ...     return jnp.array([0.04 / var * jnp.exp(-r2 / (2 * var))])
        >>> res = advection_diffusion_2d(c, kappa)
        >>> r = res(bump, jnp.array([0.3, 0.25, 0.2]))
        >>> r.shape, bool(jnp.abs(r[0]) < 1e-4)
        ((1,), True)
    """
    advection_velocity = jnp.asarray(velocity)
    if advection_velocity.shape != (2,):
        raise ShapeError(
            f"velocity must have two components; got shape {advection_velocity.shape}."
        )
    if diffusivity < 0:
        raise ConfigurationError(
            f"diffusivity must be non-negative; got {diffusivity}."
        )

    def residual(field: Field, point: Float[Array, " 3"]) -> Float[Array, " 1"]:
        def tracer(space_time_point: Float[Array, " 3"]) -> Float[Array, ""]:
            return field(space_time_point)[0]

        # ∇ₛu = (∂ₜu, ∂ₓu, ∂ᵧu): (3,)
        tracer_gradient = diffops.grad(tracer)(point)
        # Δu = ∂ₓₓu + ∂ᵧᵧu at fixed t: ()
        tracer_laplacian = diffops.laplacian(_at_time(tracer, point[0]))(point[1:])
        # N[u] = ∂ₜu + c·∇u - κΔu: ()
        pde_residual = (
            tracer_gradient[0]
            + jnp.dot(
                advection_velocity.astype(tracer_gradient.dtype), tracer_gradient[1:]
            )
            - diffusivity * tracer_laplacian
        )
        return einx.id(" -> 1", pde_residual)

    return residual

navier_stokes_vorticity_2d

navier_stokes_vorticity_2d(
    viscosity: float,
    forcing: Callable[
        [Float[Array, " 3"]], Float[Array, ""]
    ]
    | None = None,
) -> Residual

Residual of 2-D incompressible Navier–Stokes in vorticity form.

The field provides the streamfunction \(\psi(t, x, y)\) (its first output); velocity and vorticity are derived from it,

\[ u = \nabla^\perp \psi = (-\partial_y \psi,\ \partial_x \psi), \qquad \omega = \partial_x u_y - \partial_y u_x = \Delta \psi , \]

and the residual is the vorticity equation

\[ \mathcal N[\psi] = \partial_t \omega + u \cdot \nabla \omega - \nu\, \Delta \omega - F(t, x, y) . \]

Deriving \(u\) from \(\psi\) makes the flow divergence-free and the Poisson relation \(\Delta\psi = \omega\) hold exactly, so one residual component is enough. The price is fourth derivatives of the field (\(\Delta\omega = \Delta^2\psi\)), computed by nested automatic differentiation; for a field that returns \((\omega, \psi)\) separately, add \(\Delta\psi - \omega\) as a second residual instead.

The streamfunction must be four times differentiable in \((x, y)\) (and \(\Delta\psi\) once in \(t\)) at every collocation point; a raw jnp.sqrt or jnp.linalg.norm that vanishes there gives NaN derivatives, so guard it with the double-where idiom.

Parameters:

Name	Type	Description	Default
viscosity	float	Kinematic viscosity \(\nu \ge 0\).	required
forcing	Callable[[Float[Array, ' 3']], Float[Array, '']] | None	Optional vorticity forcing \(F(s)\), a scalar function of s = (t, x, y) (for Kolmogorov flow, \(F = -n f_0 \cos(n y)\) for the body force \(f_0 \sin(n y)\,\hat x\)). None means unforced.	None

Returns:

Type	Description
Residual	residual_fn(field, s) -> (1,), where s = (t, x, y) and the
Residual	streamfunction is field(s)[0].

Raises:

Type	Description
ValueError	If viscosity is negative.

Examples:

The decaying Taylor–Green vortex \(\psi = \sin x \sin y\, e^{-2\nu t}\) solves the unforced equation exactly:

>>> import jax.numpy as jnp
>>> from vardax.examples import navier_stokes_vorticity_2d
>>> nu = 0.1
>>> def tg(s):
...     return jnp.array(
...         [jnp.sin(s[1]) * jnp.sin(s[2]) * jnp.exp(-2 * nu * s[0])]
...     )
>>> res = navier_stokes_vorticity_2d(nu)
>>> r = res(tg, jnp.array([0.5, 0.3, 1.1]))
>>> r.shape, bool(jnp.abs(r[0]) < 1e-5)
((1,), True)

Source code in src/vardax/_src/utils/residuals.py

def navier_stokes_vorticity_2d(
    viscosity: float,
    forcing: Callable[[Float[Array, " 3"]], Float[Array, ""]] | None = None,
) -> Residual:
    r"""Residual of 2-D incompressible Navier–Stokes in vorticity form.

    The field provides the **streamfunction** $\psi(t, x, y)$ (its first
    output); velocity and vorticity are derived from it,

    $$
    u = \nabla^\perp \psi = (-\partial_y \psi,\ \partial_x \psi),
    \qquad \omega = \partial_x u_y - \partial_y u_x = \Delta \psi ,
    $$

    and the residual is the vorticity equation

    $$
    \mathcal N[\psi] = \partial_t \omega + u \cdot \nabla \omega
        - \nu\, \Delta \omega - F(t, x, y) .
    $$

    Deriving $u$ from $\psi$ makes the flow divergence-free and the Poisson
    relation $\Delta\psi = \omega$ hold exactly, so one residual component
    is enough. The price is fourth derivatives of the field ($\Delta\omega =
    \Delta^2\psi$), computed by nested automatic differentiation; for a field
    that returns $(\omega, \psi)$ separately, add $\Delta\psi - \omega$ as a
    second residual instead.

    The streamfunction must be four times differentiable in $(x, y)$ (and
    $\Delta\psi$ once in $t$) at every collocation point; a raw `jnp.sqrt`
    or `jnp.linalg.norm` that vanishes there gives NaN derivatives, so guard
    it with the double-`where` idiom.

    Args:
        viscosity: Kinematic viscosity $\nu \ge 0$.
        forcing: Optional vorticity forcing $F(s)$, a scalar function of
            ``s = (t, x, y)`` (for Kolmogorov flow,
            $F = -n f_0 \cos(n y)$ for the body force $f_0 \sin(n y)\,\hat x$).
            ``None`` means unforced.

    Returns:
        ``residual_fn(field, s) -> (1,)``, where ``s = (t, x, y)`` and the
        streamfunction is ``field(s)[0]``.

    Raises:
        ValueError: If ``viscosity`` is negative.

    Examples:
        The decaying Taylor–Green vortex $\psi = \sin x \sin y\,
        e^{-2\nu t}$ solves the unforced equation exactly:

        >>> import jax.numpy as jnp
        >>> from vardax.examples import navier_stokes_vorticity_2d
        >>> nu = 0.1
        >>> def tg(s):
        ...     return jnp.array(
        ...         [jnp.sin(s[1]) * jnp.sin(s[2]) * jnp.exp(-2 * nu * s[0])]
        ...     )
        >>> res = navier_stokes_vorticity_2d(nu)
        >>> r = res(tg, jnp.array([0.5, 0.3, 1.1]))
        >>> r.shape, bool(jnp.abs(r[0]) < 1e-5)
        ((1,), True)
    """
    if viscosity < 0:
        raise ConfigurationError(f"viscosity must be non-negative; got {viscosity}.")

    def residual(field: Field, point: Float[Array, " 3"]) -> Float[Array, " 1"]:
        def streamfunction(space_time_point: Float[Array, " 3"]) -> Float[Array, ""]:
            return field(space_time_point)[0]

        def vorticity(space_time_point: Float[Array, " 3"]) -> Float[Array, ""]:
            # ω = Δψ at fixed t: ()
            return diffops.laplacian(_at_time(streamfunction, space_time_point[0]))(
                space_time_point[1:]
            )

        # ∇ₛψ = (∂ₜψ, ∂ₓψ, ∂ᵧψ): (3,)
        streamfunction_gradient = diffops.grad(streamfunction)(point)
        # u = ∇^⊥ψ = (-∂ᵧψ, ∂ₓψ): (2,)
        velocity = jnp.stack([-streamfunction_gradient[2], streamfunction_gradient[1]])
        # ∇ₛω = (∂ₜω, ∂ₓω, ∂ᵧω): (3,)
        vorticity_gradient = diffops.grad(vorticity)(point)
        # Δω at fixed t: ()
        vorticity_laplacian = diffops.laplacian(_at_time(vorticity, point[0]))(
            point[1:]
        )
        # N[ψ] = ∂ₜω + u·∇ω - nu Δω - F: ()
        pde_residual = (
            vorticity_gradient[0]
            + jnp.dot(velocity, vorticity_gradient[1:])
            - viscosity * vorticity_laplacian
        )
        if forcing is not None:
            pde_residual = pde_residual - forcing(point)
        return einx.id(" -> 1", pde_residual)

    return residual

Priors

Implementations of the Prior Protocol. IdentityPrior gives plain Tikhonov regularisation; L63Prior / L96Prior encode Lorenz dynamics as a model-consistency penalty; the autoencoder priors (MLP, convolutional, and bilinear variants in 1D, 2D, and 2D-multivariate) are learned priors that penalise distance from a trained reconstruction manifold.

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

IdentityPrior

Bases: Module

Trivial identity prior: \(\varphi(x) = x\).

Zero parameters. Useful as a pure obs-driven baseline (the prior cost vanishes everywhere) and as a sanity-check building block in the linear-Gaussian agreement tests.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import IdentityPrior
>>> prior = IdentityPrior()
>>> x = jnp.arange(6.0).reshape(1, 2, 3)
>>> bool(jnp.all(prior(x) == x))
True

The prior cost \(\|x - \varphi(x)\|^2\) is therefore zero, for any shape:

>>> y = jnp.ones((4,))
>>> float(jnp.sum((y - prior(y)) ** 2))
0.0

It satisfies the Prior protocol:

>>> import vardax as vdx
>>> isinstance(prior, vdx.Prior)
True

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

class IdentityPrior(eqx.Module):
    r"""Trivial identity prior: $\varphi(x) = x$.

    Zero parameters. Useful as a pure obs-driven baseline (the prior
    cost vanishes everywhere) and as a sanity-check building block in
    the linear-Gaussian agreement tests.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import IdentityPrior
        >>> prior = IdentityPrior()
        >>> x = jnp.arange(6.0).reshape(1, 2, 3)
        >>> bool(jnp.all(prior(x) == x))
        True

        The prior cost $\|x - \varphi(x)\|^2$ is therefore zero, for any shape:

        >>> y = jnp.ones((4,))
        >>> float(jnp.sum((y - prior(y)) ** 2))
        0.0

        It satisfies the [`Prior`][vardax.Prior] protocol:

        >>> import vardax as vdx
        >>> isinstance(prior, vdx.Prior)
        True
    """

    def __call__(self, x: Float[Array, ...]) -> Float[Array, ...]:
        r"""Return the input unchanged.

        Args:
            x: State to project.

        Returns:
            The same array, $\varphi(x) = x$.
        """
        return x

L63Prior

Bases: Module

Learned prior for the Lorenz-63 system.

A simple MLP autoencoder designed for the 3-dimensional Lorenz-63 attractor, \(\varphi(x) = D(E(x))\). The state is treated as a flat vector of length 3.

Algorithm

1. Encode with a tanh hidden layer and a linear map to the latent code.
2. Decode with a tanh hidden layer and a linear map back to the state.

Parameters:

Name	Type	Description	Default
latent_dim	int	Dimensionality of the latent code.	3
hidden_dim	int	Hidden layer width.	32
state_dim	int	Dimensionality of the state vector.	3
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import L63Prior
>>> prior = L63Prior(key=jax.random.PRNGKey(0))
>>> prior(jnp.ones((5, 3))).shape
(5, 3)

The latent and hidden widths are free parameters:

>>> small = L63Prior(latent_dim=2, hidden_dim=4, key=jax.random.PRNGKey(0))
>>> small(jnp.zeros((1, 3))).shape
(1, 3)

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

class L63Prior(eqx.Module):
    r"""Learned prior for the Lorenz-63 system.

    A simple MLP autoencoder designed for the 3-dimensional Lorenz-63
    attractor, $\varphi(x) = D(E(x))$. The state is treated as a flat vector of
    length ``3``.

    Algorithm:
        1. Encode with a tanh hidden layer and a linear map to the latent code.
        2. Decode with a tanh hidden layer and a linear map back to the state.

    Args:
        latent_dim: Dimensionality of the latent code.
        hidden_dim: Hidden layer width.
        state_dim: Dimensionality of the state vector.
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import L63Prior
        >>> prior = L63Prior(key=jax.random.PRNGKey(0))
        >>> prior(jnp.ones((5, 3))).shape
        (5, 3)

        The latent and hidden widths are free parameters:

        >>> small = L63Prior(latent_dim=2, hidden_dim=4, key=jax.random.PRNGKey(0))
        >>> small(jnp.zeros((1, 3))).shape
        (1, 3)
    """

    encoder_hidden: eqx.nn.Linear
    encoder_out: eqx.nn.Linear
    decoder_hidden: eqx.nn.Linear
    decoder_out: eqx.nn.Linear

    def __init__(
        self,
        latent_dim: int = 3,
        hidden_dim: int = 32,
        state_dim: int = 3,
        *,
        key: PRNGKeyArray,
    ) -> None:
        key_enc1, key_enc2, key_dec1, key_dec2 = jax.random.split(key, 4)
        self.encoder_hidden = eqx.nn.Linear(state_dim, hidden_dim, key=key_enc1)
        self.encoder_out = eqx.nn.Linear(hidden_dim, latent_dim, key=key_enc2)
        self.decoder_hidden = eqx.nn.Linear(latent_dim, hidden_dim, key=key_dec1)
        self.decoder_out = eqx.nn.Linear(hidden_dim, state_dim, key=key_dec2)

    def __call__(self, x: Float[Array, "B N"]) -> Float[Array, "B N"]:
        latent = jnp.tanh(jax.vmap(self.encoder_hidden)(x))
        latent = jax.vmap(self.encoder_out)(latent)
        hidden = jnp.tanh(jax.vmap(self.decoder_hidden)(latent))
        return jax.vmap(self.decoder_out)(hidden)

L96Prior

Bases: Module

Learned prior for the Lorenz-96 system.

A simple MLP autoencoder designed for the N-dimensional Lorenz-96 attractor, \(\varphi(x) = D(E(x))\). The state is treated as a flat vector of length N.

Algorithm

1. Encode with a tanh hidden layer and a linear map to the latent code.
2. Decode with a tanh hidden layer and a linear map back to the state.

Parameters:

Name	Type	Description	Default
latent_dim	int	Dimensionality of the latent code.	16
hidden_dim	int	Hidden layer width.	64
state_dim	int	Dimensionality of the state vector.	40
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import L96Prior
>>> prior = L96Prior(state_dim=8, latent_dim=4, key=jax.random.PRNGKey(0))
>>> prior(jnp.ones((3, 8))).shape
(3, 8)

The defaults target the standard 40-variable system:

>>> L96Prior(key=jax.random.PRNGKey(0))(jnp.ones((1, 40))).shape
(1, 40)

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

class L96Prior(eqx.Module):
    r"""Learned prior for the Lorenz-96 system.

    A simple MLP autoencoder designed for the N-dimensional Lorenz-96
    attractor, $\varphi(x) = D(E(x))$. The state is treated as a flat vector of
    length ``N``.

    Algorithm:
        1. Encode with a tanh hidden layer and a linear map to the latent code.
        2. Decode with a tanh hidden layer and a linear map back to the state.

    Args:
        latent_dim: Dimensionality of the latent code.
        hidden_dim: Hidden layer width.
        state_dim: Dimensionality of the state vector.
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import L96Prior
        >>> prior = L96Prior(state_dim=8, latent_dim=4, key=jax.random.PRNGKey(0))
        >>> prior(jnp.ones((3, 8))).shape
        (3, 8)

        The defaults target the standard 40-variable system:

        >>> L96Prior(key=jax.random.PRNGKey(0))(jnp.ones((1, 40))).shape
        (1, 40)
    """

    encoder_hidden: eqx.nn.Linear
    encoder_out: eqx.nn.Linear
    decoder_hidden: eqx.nn.Linear
    decoder_out: eqx.nn.Linear

    def __init__(
        self,
        latent_dim: int = 16,
        hidden_dim: int = 64,
        state_dim: int = 40,
        *,
        key: PRNGKeyArray,
    ) -> None:
        key_enc1, key_enc2, key_dec1, key_dec2 = jax.random.split(key, 4)
        self.encoder_hidden = eqx.nn.Linear(state_dim, hidden_dim, key=key_enc1)
        self.encoder_out = eqx.nn.Linear(hidden_dim, latent_dim, key=key_enc2)
        self.decoder_hidden = eqx.nn.Linear(latent_dim, hidden_dim, key=key_dec1)
        self.decoder_out = eqx.nn.Linear(hidden_dim, state_dim, key=key_dec2)

    def __call__(self, x: Float[Array, "B N"]) -> Float[Array, "B N"]:
        latent = jnp.tanh(jax.vmap(self.encoder_hidden)(x))
        latent = jax.vmap(self.encoder_out)(latent)
        hidden = jnp.tanh(jax.vmap(self.decoder_hidden)(latent))
        return jax.vmap(self.decoder_out)(hidden)

MLPAEPrior1D

Bases: Module

MLP autoencoder prior for 1-D data.

A two-layer MLP encoder and decoder around a latent code, \(\varphi(x) = D(E(x))\), applied to the flattened (T, N) window.

Algorithm

1. Flatten each sample over time and space.
2. Encode with a ReLU hidden layer and a linear map to the latent code.
3. Decode with a ReLU hidden layer and a linear map back, then reshape to (T, N).

Parameters:

Name	Type	Description	Default
state_dim	int	Spatial size of the input (N).	required
latent_dim	int	Dimensionality of the latent code.	required
hidden_dim	int	Hidden layer width.	64
n_time	int	Number of time steps (T).	1
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import MLPAEPrior1D
>>> prior = MLPAEPrior1D(
...     state_dim=4,
...     latent_dim=2,
...     hidden_dim=8,
...     n_time=3,
...     key=jax.random.PRNGKey(0),
... )
>>> prior(jnp.ones((2, 3, 4))).shape
(2, 3, 4)

A single time step is the default:

>>> flat = MLPAEPrior1D(state_dim=4, latent_dim=2, key=jax.random.PRNGKey(0))
>>> flat(jnp.ones((5, 1, 4))).shape
(5, 1, 4)

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

class MLPAEPrior1D(eqx.Module):
    r"""MLP autoencoder prior for 1-D data.

    A two-layer MLP encoder and decoder around a latent code,
    $\varphi(x) = D(E(x))$, applied to the flattened ``(T, N)`` window.

    Algorithm:
        1. Flatten each sample over time and space.
        2. Encode with a ReLU hidden layer and a linear map to the latent code.
        3. Decode with a ReLU hidden layer and a linear map back, then reshape
           to ``(T, N)``.

    Args:
        state_dim: Spatial size of the input (``N``).
        latent_dim: Dimensionality of the latent code.
        hidden_dim: Hidden layer width.
        n_time: Number of time steps (``T``).
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import MLPAEPrior1D
        >>> prior = MLPAEPrior1D(
        ...     state_dim=4,
        ...     latent_dim=2,
        ...     hidden_dim=8,
        ...     n_time=3,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> prior(jnp.ones((2, 3, 4))).shape
        (2, 3, 4)

        A single time step is the default:

        >>> flat = MLPAEPrior1D(state_dim=4, latent_dim=2, key=jax.random.PRNGKey(0))
        >>> flat(jnp.ones((5, 1, 4))).shape
        (5, 1, 4)
    """

    n_time: int = eqx.field(static=True)
    encoder_hidden: eqx.nn.Linear
    encoder_out: eqx.nn.Linear
    decoder_hidden: eqx.nn.Linear
    decoder_out: eqx.nn.Linear

    def __init__(
        self,
        state_dim: int,
        latent_dim: int,
        hidden_dim: int = 64,
        n_time: int = 1,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.n_time = n_time
        in_features = n_time * state_dim
        key_enc1, key_enc2, key_dec1, key_dec2 = jax.random.split(key, 4)
        self.encoder_hidden = eqx.nn.Linear(in_features, hidden_dim, key=key_enc1)
        self.encoder_out = eqx.nn.Linear(hidden_dim, latent_dim, key=key_enc2)
        self.decoder_hidden = eqx.nn.Linear(latent_dim, hidden_dim, key=key_dec1)
        self.decoder_out = eqx.nn.Linear(hidden_dim, in_features, key=key_dec2)

    def __call__(self, x: Float[Array, "B T N"]) -> Float[Array, "B T N"]:
        state_flat = einx.id("b t n -> b (t n)", x)
        latent = jax.nn.relu(jax.vmap(self.encoder_hidden)(state_flat))
        latent = jax.vmap(self.encoder_out)(latent)
        hidden = jax.nn.relu(jax.vmap(self.decoder_hidden)(latent))
        decoded = jax.vmap(self.decoder_out)(hidden)
        return einx.id("b (t n) -> b t n", decoded, t=x.shape[1])

ConvAEPrior1D

Bases: Module

Convolutional autoencoder prior for 1-D spatially-structured data.

Uses circular (periodic) padding suitable for systems with periodic boundary conditions such as Lorenz-96, \(\varphi(x) = D(E(x))\) with convolutional \(E\) and \(D\).

Algorithm

1. Pad the spatial axis periodically and apply the encoder convolution.
2. Apply a ReLU.
3. Pad periodically again and apply the decoder convolution back to n_time channels.

Parameters:

Name	Type	Description	Default
latent_channels	int	Number of channels in the latent representation.	16
kernel_size	int	Convolution kernel size (must be a positive odd integer).	3
n_time	int	Number of time steps T; used as the input/output channels and validated against the runtime input shape.	1
key	PRNGKeyArray	PRNG key for the convolution weights.	required

Raises:

Type	Description
ConfigurationError	If kernel_size is not a positive odd integer.
ShapeError	On call, if the input's time dimension differs from n_time.

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import ConvAEPrior1D
>>> prior = ConvAEPrior1D(n_time=3, key=jax.random.PRNGKey(0))
>>> prior(jnp.ones((2, 3, 10))).shape
(2, 3, 10)

Periodic padding makes the prior shift-equivariant on the ring:

>>> x = jax.random.normal(jax.random.PRNGKey(1), (1, 3, 10))
>>> shifted = prior(jnp.roll(x, 2, axis=-1))
>>> bool(jnp.allclose(shifted, jnp.roll(prior(x), 2, axis=-1), atol=1e-5))
True

An even kernel is rejected:

>>> ConvAEPrior1D(
...     kernel_size=4, key=jax.random.PRNGKey(0)
... )
Traceback (most recent call last):
...
ConfigurationError: kernel_size must be a positive odd integer, got 4.

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

class ConvAEPrior1D(eqx.Module):
    r"""Convolutional autoencoder prior for 1-D spatially-structured data.

    Uses circular (periodic) padding suitable for systems with periodic
    boundary conditions such as Lorenz-96, $\varphi(x) = D(E(x))$ with
    convolutional $E$ and $D$.

    Algorithm:
        1. Pad the spatial axis periodically and apply the encoder convolution.
        2. Apply a ReLU.
        3. Pad periodically again and apply the decoder convolution back to
           ``n_time`` channels.

    Args:
        latent_channels: Number of channels in the latent representation.
        kernel_size: Convolution kernel size (must be a positive odd integer).
        n_time: Number of time steps ``T``; used as the input/output channels
            and validated against the runtime input shape.
        key: PRNG key for the convolution weights.

    Raises:
        ConfigurationError: If ``kernel_size`` is not a positive odd integer.
        ShapeError: On call, if the input's time dimension differs from
            ``n_time``.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import ConvAEPrior1D
        >>> prior = ConvAEPrior1D(n_time=3, key=jax.random.PRNGKey(0))
        >>> prior(jnp.ones((2, 3, 10))).shape
        (2, 3, 10)

        Periodic padding makes the prior shift-equivariant on the ring:

        >>> x = jax.random.normal(jax.random.PRNGKey(1), (1, 3, 10))
        >>> shifted = prior(jnp.roll(x, 2, axis=-1))
        >>> bool(jnp.allclose(shifted, jnp.roll(prior(x), 2, axis=-1), atol=1e-5))
        True

        An even kernel is rejected:

        >>> ConvAEPrior1D(
        ...     kernel_size=4, key=jax.random.PRNGKey(0)
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: kernel_size must be a positive odd integer, got 4.
    """

    kernel_size: int = eqx.field(static=True)
    n_time: int = eqx.field(static=True)
    _enc_conv: eqx.nn.Conv1d
    _dec_conv: eqx.nn.Conv1d

    def __init__(
        self,
        latent_channels: int = 16,
        kernel_size: int = 3,
        n_time: int = 1,
        *,
        key: PRNGKeyArray,
    ) -> None:
        if kernel_size <= 0 or kernel_size % 2 == 0:
            raise ConfigurationError(
                f"kernel_size must be a positive odd integer, got {kernel_size}."
            )
        self.kernel_size = kernel_size
        self.n_time = n_time
        key_enc, key_dec = jax.random.split(key)
        # eqx.nn.Conv1d uses channels-first: (in_channels, length)
        self._enc_conv = eqx.nn.Conv1d(
            in_channels=n_time,
            out_channels=latent_channels,
            kernel_size=kernel_size,
            padding=0,  # we apply circular padding manually
            key=key_enc,
        )
        self._dec_conv = eqx.nn.Conv1d(
            in_channels=latent_channels,
            out_channels=n_time,
            kernel_size=kernel_size,
            padding=0,
            key=key_dec,
        )

    def __call__(self, x: Float[Array, "B T N"]) -> Float[Array, "B T N"]:
        n_time = x.shape[1]
        if n_time != self.n_time:
            raise ShapeError(
                f"Input time dimension {n_time} does not match n_time={self.n_time}."
            )
        # eqx Conv1d expects channels-first: (in_channels=T, length=N)
        # x is already (B, T, N) so vmap over the batch dim works directly.
        pad = self.kernel_size // 2

        def _forward(sample: Float[Array, "T N"]) -> Float[Array, "T N"]:
            # circular padding on the spatial axis only: the grid is periodic
            periodic_pad = ((0, 0), (pad, pad))
            hidden = jnp.pad(sample, periodic_pad, mode="wrap")
            hidden = self._enc_conv(hidden)
            hidden = jax.nn.relu(hidden)
            hidden = jnp.pad(hidden, periodic_pad, mode="wrap")
            return self._dec_conv(hidden)

        return jax.vmap(_forward)(x)

BilinAEPrior1D

Bases: Module

Bilinear autoencoder prior for 1-D data.

The encoder maps the input to a low-dimensional latent code; the decoder reconstructs the original space, \(\varphi(x) = D(E(x))\). The prior cost is \(\|x - \varphi(x)\|^2\).

Algorithm

1. Flatten each sample over time and space.
2. Encode with the bilinear block \(\mathrm{relu}(Ax) \odot \tanh(Bx)\).
3. Decode with a dense layer and reshape to (T, N).

Parameters:

Name	Type	Description	Default
state_dim	int	Spatial size of the input (N).	required
latent_dim	int	Dimensionality of the latent code.	required
n_time	int	Number of time steps (T).	1
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import BilinAEPrior1D
>>> prior = BilinAEPrior1D(
...     state_dim=4, latent_dim=2, n_time=3, key=jax.random.PRNGKey(0)
... )
>>> prior(jnp.ones((2, 3, 4))).shape
(2, 3, 4)

The latent code is low-dimensional, and decoding it is the same as calling the prior:

>>> x = jnp.ones((2, 3, 4))
>>> prior.encode(x).shape
(2, 2)
>>> bool(jnp.allclose(prior.decode(prior.encode(x)), prior(x), atol=1e-6))
True

Used as prior_fn of a solver step, with a pass-through modulator:

>>> import vardax as vdx
>>> batch = vdx.Batch1D(input=x, mask=jnp.ones_like(x))
>>> state = vdx.init_solver_state_1d(batch, hidden_dim=5)
>>> passthrough = lambda grad, state, lstm: (grad, lstm)
>>> vdx.solver_step_1d(state, batch, prior, passthrough).x.shape
(2, 3, 4)

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

class BilinAEPrior1D(eqx.Module):
    r"""Bilinear autoencoder prior for 1-D data.

    The encoder maps the input to a low-dimensional latent code; the decoder
    reconstructs the original space, $\varphi(x) = D(E(x))$.  The prior cost is
    $\|x - \varphi(x)\|^2$.

    Algorithm:
        1. Flatten each sample over time and space.
        2. Encode with the bilinear block $\mathrm{relu}(Ax) \odot \tanh(Bx)$.
        3. Decode with a dense layer and reshape to ``(T, N)``.

    Args:
        state_dim: Spatial size of the input (``N``).
        latent_dim: Dimensionality of the latent code.
        n_time: Number of time steps (``T``).
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import BilinAEPrior1D
        >>> prior = BilinAEPrior1D(
        ...     state_dim=4, latent_dim=2, n_time=3, key=jax.random.PRNGKey(0)
        ... )
        >>> prior(jnp.ones((2, 3, 4))).shape
        (2, 3, 4)

        The latent code is low-dimensional, and decoding it is the same as
        calling the prior:

        >>> x = jnp.ones((2, 3, 4))
        >>> prior.encode(x).shape
        (2, 2)
        >>> bool(jnp.allclose(prior.decode(prior.encode(x)), prior(x), atol=1e-6))
        True

        Used as ``prior_fn`` of a solver step, with a pass-through modulator:

        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(input=x, mask=jnp.ones_like(x))
        >>> state = vdx.init_solver_state_1d(batch, hidden_dim=5)
        >>> passthrough = lambda grad, state, lstm: (grad, lstm)
        >>> vdx.solver_step_1d(state, batch, prior, passthrough).x.shape
        (2, 3, 4)
    """

    state_dim: int = eqx.field(static=True)
    latent_dim: int = eqx.field(static=True)
    n_time: int = eqx.field(static=True)
    _bilin: _BilinearBlock
    _decode_dense: eqx.nn.Linear

    def __init__(
        self,
        state_dim: int,
        latent_dim: int,
        n_time: int = 1,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.state_dim = state_dim
        self.latent_dim = latent_dim
        self.n_time = n_time
        in_features = n_time * state_dim
        key_bilin, key_dec = jax.random.split(key)
        self._bilin = _BilinearBlock(in_features, latent_dim, key=key_bilin)
        self._decode_dense = eqx.nn.Linear(latent_dim, n_time * state_dim, key=key_dec)

    def __call__(self, x: Float[Array, "B T N"]) -> Float[Array, "B T N"]:
        state_flat = einx.id("b t n -> b (t n)", x)
        latent = self._bilin(state_flat)
        decoded = jax.vmap(self._decode_dense)(latent)
        return einx.id("b (t n) -> b t n", decoded, t=x.shape[1])

    def encode(self, x: Float[Array, "B T N"]) -> Float[Array, "B Z"]:
        """Encode input to latent space.

        Args:
            x: State window to encode.

        Returns:
            Latent code $E(x)$ of dimension ``latent_dim`` per sample.
        """
        state_flat = einx.id("b t n -> b (t n)", x)
        return self._bilin(state_flat)

    def decode(self, z: Float[Array, "B Z"]) -> Float[Array, "B T N"]:
        """Decode latent code to state space.

        Args:
            z: Latent code $z$.

        Returns:
            Decoded state window $D(z)$.
        """
        decoded = jax.vmap(self._decode_dense)(z)
        return einx.id("b (t n) -> b t n", decoded, t=self.n_time, n=self.state_dim)

encode

encode(x: Float[Array, 'B T N']) -> Float[Array, 'B Z']

Encode input to latent space.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'B T N']	State window to encode.	required

Returns:

Type	Description
Float[Array, 'B Z']	Latent code \(E(x)\) of dimension latent_dim per sample.

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

def encode(self, x: Float[Array, "B T N"]) -> Float[Array, "B Z"]:
    """Encode input to latent space.

    Args:
        x: State window to encode.

    Returns:
        Latent code $E(x)$ of dimension ``latent_dim`` per sample.
    """
    state_flat = einx.id("b t n -> b (t n)", x)
    return self._bilin(state_flat)

decode

decode(z: Float[Array, 'B Z']) -> Float[Array, 'B T N']

Decode latent code to state space.

Parameters:

Name	Type	Description	Default
z	Float[Array, 'B Z']	Latent code \(z\).	required

Returns:

Type	Description
Float[Array, 'B T N']	Decoded state window \(D(z)\).

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

def decode(self, z: Float[Array, "B Z"]) -> Float[Array, "B T N"]:
    """Decode latent code to state space.

    Args:
        z: Latent code $z$.

    Returns:
        Decoded state window $D(z)$.
    """
    decoded = jax.vmap(self._decode_dense)(z)
    return einx.id("b (t n) -> b t n", decoded, t=self.n_time, n=self.state_dim)

BilinAEPrior2D

Bases: Module

Bilinear autoencoder prior for 2-D data.

The window is flattened, encoded with the bilinear block \(\mathrm{relu}(Ax) \odot \tanh(Bx)\) and decoded by a dense layer, \(\varphi(x) = D(E(x))\).

Algorithm

1. Flatten each sample over time and both spatial axes.
2. Encode with the bilinear block.
3. Decode with a dense layer and reshape to (T, H, W).

Parameters:

Name	Type	Description	Default
latent_dim	int	Dimensionality of the latent code.	required
n_time	int	Number of time steps (T).	required
height	int	Spatial height H.	required
width	int	Spatial width W.	required
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import BilinAEPrior2D
>>> prior = BilinAEPrior2D(
...     latent_dim=2, n_time=3, height=4, width=5, key=jax.random.PRNGKey(0)
... )
>>> prior(jnp.ones((1, 3, 4, 5))).shape
(1, 3, 4, 5)

It is the prior inside FourDVarNet2D:

>>> import vardax as vdx
>>> model = vdx.FourDVarNet2D(
...     n_time=3, height=4, width=5, key=jax.random.PRNGKey(0)
... )
>>> type(model.prior).__name__
'BilinAEPrior2D'

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

class BilinAEPrior2D(eqx.Module):
    r"""Bilinear autoencoder prior for 2-D data.

    The window is flattened, encoded with the bilinear block
    $\mathrm{relu}(Ax) \odot \tanh(Bx)$ and decoded by a dense layer,
    $\varphi(x) = D(E(x))$.

    Algorithm:
        1. Flatten each sample over time and both spatial axes.
        2. Encode with the bilinear block.
        3. Decode with a dense layer and reshape to ``(T, H, W)``.

    Args:
        latent_dim: Dimensionality of the latent code.
        n_time: Number of time steps (``T``).
        height: Spatial height ``H``.
        width: Spatial width ``W``.
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import BilinAEPrior2D
        >>> prior = BilinAEPrior2D(
        ...     latent_dim=2, n_time=3, height=4, width=5, key=jax.random.PRNGKey(0)
        ... )
        >>> prior(jnp.ones((1, 3, 4, 5))).shape
        (1, 3, 4, 5)

        It is the prior inside [`FourDVarNet2D`][vardax.FourDVarNet2D]:

        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet2D(
        ...     n_time=3, height=4, width=5, key=jax.random.PRNGKey(0)
        ... )
        >>> type(model.prior).__name__
        'BilinAEPrior2D'
    """

    n_time: int = eqx.field(static=True)
    height: int = eqx.field(static=True)
    width: int = eqx.field(static=True)
    _bilin: _BilinearBlock
    _decode_dense: eqx.nn.Linear

    def __init__(
        self,
        latent_dim: int,
        n_time: int,
        height: int,
        width: int,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.n_time = n_time
        self.height = height
        self.width = width
        in_features = n_time * height * width
        key_bilin, key_dec = jax.random.split(key)
        self._bilin = _BilinearBlock(in_features, latent_dim, key=key_bilin)
        self._decode_dense = eqx.nn.Linear(latent_dim, in_features, key=key_dec)

    def __call__(self, x: Float[Array, "B T H W"]) -> Float[Array, "B T H W"]:
        _, n_time, height, width = x.shape
        state_flat = einx.id("b t h w -> b (t h w)", x)
        latent = self._bilin(state_flat)
        decoded = jax.vmap(self._decode_dense)(latent)
        return einx.id("b (t h w) -> b t h w", decoded, t=n_time, h=height, w=width)

BilinAEPrior2DMultivar

Bases: Module

Bilinear autoencoder prior for 2-D multivariate data.

The window is flattened over time, channels and space, encoded with the bilinear block \(\mathrm{relu}(Ax) \odot \tanh(Bx)\) and decoded by a dense layer, \(\varphi(x) = D(E(x))\).

Algorithm

1. Flatten each sample over time, channel and both spatial axes.
2. Encode with the bilinear block.
3. Decode with a dense layer and reshape to (T, C, H, W).

Parameters:

Name	Type	Description	Default
latent_dim	int	Dimensionality of the latent code.	required
n_time	int	Number of time steps (T).	required
n_channels	int	Number of channels C.	required
height	int	Spatial height H.	required
width	int	Spatial width W.	required
key	PRNGKeyArray	PRNG key for the layer weights.	required

Examples:

>>> import jax, jax.numpy as jnp
>>> from vardax import BilinAEPrior2DMultivar
>>> prior = BilinAEPrior2DMultivar(
...     latent_dim=2,
...     n_time=3,
...     n_channels=2,
...     height=4,
...     width=5,
...     key=jax.random.PRNGKey(0),
... )
>>> prior(jnp.ones((1, 3, 2, 4, 5))).shape
(1, 3, 2, 4, 5)

The batch axis is free, so a larger batch works unchanged:

>>> prior(jnp.ones((4, 3, 2, 4, 5))).shape
(4, 3, 2, 4, 5)

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

class BilinAEPrior2DMultivar(eqx.Module):
    r"""Bilinear autoencoder prior for 2-D multivariate data.

    The window is flattened over time, channels and space, encoded with the
    bilinear block $\mathrm{relu}(Ax) \odot \tanh(Bx)$ and decoded by a dense
    layer, $\varphi(x) = D(E(x))$.

    Algorithm:
        1. Flatten each sample over time, channel and both spatial axes.
        2. Encode with the bilinear block.
        3. Decode with a dense layer and reshape to ``(T, C, H, W)``.

    Args:
        latent_dim: Dimensionality of the latent code.
        n_time: Number of time steps (``T``).
        n_channels: Number of channels ``C``.
        height: Spatial height ``H``.
        width: Spatial width ``W``.
        key: PRNG key for the layer weights.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> from vardax import BilinAEPrior2DMultivar
        >>> prior = BilinAEPrior2DMultivar(
        ...     latent_dim=2,
        ...     n_time=3,
        ...     n_channels=2,
        ...     height=4,
        ...     width=5,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> prior(jnp.ones((1, 3, 2, 4, 5))).shape
        (1, 3, 2, 4, 5)

        The batch axis is free, so a larger batch works unchanged:

        >>> prior(jnp.ones((4, 3, 2, 4, 5))).shape
        (4, 3, 2, 4, 5)
    """

    n_time: int = eqx.field(static=True)
    n_channels: int = eqx.field(static=True)
    height: int = eqx.field(static=True)
    width: int = eqx.field(static=True)
    _bilin: _BilinearBlock
    _decode_dense: eqx.nn.Linear

    def __init__(
        self,
        latent_dim: int,
        n_time: int,
        n_channels: int,
        height: int,
        width: int,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.n_time = n_time
        self.n_channels = n_channels
        self.height = height
        self.width = width
        in_features = n_time * n_channels * height * width
        key_bilin, key_dec = jax.random.split(key)
        self._bilin = _BilinearBlock(in_features, latent_dim, key=key_bilin)
        self._decode_dense = eqx.nn.Linear(latent_dim, in_features, key=key_dec)

    def __call__(self, x: Float[Array, "B T C H W"]) -> Float[Array, "B T C H W"]:
        _, n_time, n_channels, height, width = x.shape
        state_flat = einx.id("b t c h w -> b (t c h w)", x)
        latent = self._bilin(state_flat)
        decoded = jax.vmap(self._decode_dense)(latent)
        return einx.id(
            "b (t c h w) -> b t c h w",
            decoded,
            t=n_time,
            c=n_channels,
            h=height,
            w=width,
        )

Dynamical priors

ODE-based temporal priors ported from mfourdvar (Decision D18): DynIncrements scores one-step increments, DynTrajectory scores the full rollout from the initial state. Both satisfy the TemporalPrior protocol; bind(ts) adapts either to the static Prior seam, and as_forward_model(dt) adapts the wrapped ODE to pipekit_cycle.ForwardModel.

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

DynamicalPrior

Bases: Module

Base class for ODE-based dynamical priors.

Wraps a diffrax diffeqsolve around an ODE right-hand side \(\mathrm{d}x/\mathrm{d}t = f(t, x; \theta)\) with pluggable solver, step-size controller, and adjoint strategy. The adjoint choice trades memory against speed for reverse-mode differentiation through the solve (relevant for long windows).

Concrete subclasses implement __call__ (the integration) and loss (the dynamical residual). Instantiating DynamicalPrior directly and calling it raises NotImplementedError.

Parameters:

Name	Type	Description	Default
model	Callable	Diffrax-compatible ODE right-hand side f(t, y, args) -> dy.	required
params	PyTree | None	Optional default args threaded to model / diffeqsolve (e.g. learnable ODE parameters \(\theta\)). Overridable per call.	None
solver	Solver | None	Diffrax solver (default :class:diffrax.Tsit5).	None
stepsize	StepSizeController | None	Diffrax step-size controller. Defaults to an adaptive :class:diffrax.PIDController for solvers that provide error estimates (:class:diffrax.AbstractAdaptiveSolver subclasses), and to :class:diffrax.ConstantStepSize for fixed-step solvers such as :class:diffrax.Euler, which cannot drive an adaptive controller.	None
adjoint	Adjoint | None	Diffrax adjoint strategy (default :class:diffrax.RecursiveCheckpointAdjoint).	None

Examples:

The concrete subclasses share the constructor; the step-size controller follows the solver:

>>> import diffrax as dfx
>>> import vardax as vdx
>>> def decay(t, y, args):
...     return -y
>>> type(vdx.DynTrajectory(model=decay).stepsize).__name__
'PIDController'
>>> euler = vdx.DynTrajectory(model=decay, solver=dfx.Euler())
>>> type(euler.stepsize).__name__
'ConstantStepSize'

The base class itself has no integration to run:

>>> import jax.numpy as jnp
>>> vdx.DynamicalPrior(model=decay)(
...     jnp.ones(2), jnp.array([0.0, 1.0])
... )
Traceback (most recent call last):
...
NotImplementedError

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

class DynamicalPrior(eqx.Module):
    r"""Base class for ODE-based dynamical priors.

    Wraps a diffrax ``diffeqsolve`` around an ODE right-hand side
    $\mathrm{d}x/\mathrm{d}t = f(t, x; \theta)$ with pluggable solver, step-size
    controller, and adjoint strategy. The adjoint choice trades memory against
    speed for reverse-mode differentiation through the solve (relevant for long
    windows).

    Concrete subclasses implement ``__call__`` (the integration) and
    ``loss`` (the dynamical residual). Instantiating ``DynamicalPrior``
    directly and calling it raises ``NotImplementedError``.

    Args:
        model: Diffrax-compatible ODE right-hand side ``f(t, y, args) -> dy``.
        params: Optional default ``args`` threaded to ``model`` /
            ``diffeqsolve`` (e.g. learnable ODE parameters $\theta$).
            Overridable per call.
        solver: Diffrax solver (default :class:`diffrax.Tsit5`).
        stepsize: Diffrax step-size controller. Defaults to an adaptive
            :class:`diffrax.PIDController` for solvers that provide error
            estimates (:class:`diffrax.AbstractAdaptiveSolver` subclasses),
            and to :class:`diffrax.ConstantStepSize` for fixed-step solvers
            such as :class:`diffrax.Euler`, which cannot drive an adaptive
            controller.
        adjoint: Diffrax adjoint strategy (default
            :class:`diffrax.RecursiveCheckpointAdjoint`).

    Examples:
        The concrete subclasses share the constructor; the step-size controller
        follows the solver:

        >>> import diffrax as dfx
        >>> import vardax as vdx
        >>> def decay(t, y, args):
        ...     return -y
        >>> type(vdx.DynTrajectory(model=decay).stepsize).__name__
        'PIDController'
        >>> euler = vdx.DynTrajectory(model=decay, solver=dfx.Euler())
        >>> type(euler.stepsize).__name__
        'ConstantStepSize'

        The base class itself has no integration to run:

        >>> import jax.numpy as jnp
        >>> vdx.DynamicalPrior(model=decay)(
        ...     jnp.ones(2), jnp.array([0.0, 1.0])
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        NotImplementedError
    """

    model: tp.Callable
    params: PyTree | None
    solver: Solver
    stepsize: StepSizeController
    adjoint: Adjoint

    def __init__(
        self,
        model: tp.Callable,
        params: PyTree | None = None,
        solver: Solver | None = None,
        stepsize: StepSizeController | None = None,
        adjoint: Adjoint | None = None,
    ) -> None:
        self.model = model
        self.params = params
        self.solver = solver if solver is not None else dfx.Tsit5()
        if stepsize is not None:
            self.stepsize = stepsize
        elif isinstance(self.solver, dfx.AbstractAdaptiveSolver):
            self.stepsize = dfx.PIDController(rtol=1e-5, atol=1e-5)
        else:
            # Fixed-step solvers (e.g. Euler) provide no error estimate and
            # cannot drive an adaptive controller.
            self.stepsize = dfx.ConstantStepSize()
        self.adjoint = (
            adjoint if adjoint is not None else dfx.RecursiveCheckpointAdjoint()
        )

    def _integrate(
        self,
        y0: PyTree,
        ts: Array,
        saveat: dfx.SaveAt,
        dt: float | None,
        params: PyTree | None,
    ) -> PyTree:
        """Integrate ``model`` from ``ts[0]`` to ``ts[-1]`` starting at ``y0``."""
        start_time = ts[0]
        end_time = ts[-1]
        initial_step = ts[1] - ts[0] if dt is None else dt
        solution = dfx.diffeqsolve(
            terms=dfx.ODETerm(self.model),
            solver=self.solver,
            t0=start_time,
            t1=end_time,
            dt0=initial_step,
            y0=y0,
            saveat=saveat,
            adjoint=self.adjoint,
            stepsize_controller=self.stepsize,
            args=params if params is not None else self.params,
        )
        return solution.ys

    def __call__(
        self,
        x: Array,
        ts: Array,
        dt: float | None = None,
        params: PyTree | None = None,
    ) -> PyTree:
        raise NotImplementedError

    def loss(
        self,
        x: Array,
        ts: Array,
        x_gt: Array | None = None,
        params: PyTree | None = None,
    ) -> Array:
        raise NotImplementedError

    def reconstruct(
        self,
        x: Array,
        ts: Array,
        params: PyTree | None = None,
    ) -> Array:
        r"""Dynamics-consistent reconstruction of a state sequence.

        Returns an array of the *same shape* as ``x`` such that
        $\|x - \mathrm{reconstruct}(x, t_s)\|^2$ is the prior's dynamical residual,
        the temporal analogue of the autoencoder reconstruction used by the
        static [`Prior`][vardax.Prior] seam (Decision D18).

        Args:
            x: State sequence.
            ts: Time coordinates $t_s$ of the sequence.
            params: ODE ``args`` override.

        Returns:
            The dynamics-consistent sequence, with the shape of ``x``.
        """
        raise NotImplementedError

    def bind(
        self,
        ts: Array,
        params: PyTree | None = None,
    ) -> _BoundTemporalPrior:
        """Bind a time grid, yielding a one-argument `Prior` adapter.

        The returned callable satisfies the static
        [`Prior`][vardax.Prior] protocol (``__call__(x) -> x_prior``) by
        closing over ``ts``, so a dynamical prior can be used anywhere a
        static prior is expected, e.g. as the ``prior_fn`` of
        [`variational_cost`][vardax.variational_cost], turning the prior
        term into a weak-constraint dynamical residual.

        Args:
            ts: Time coordinates $t_s$ to close over.
            params: ODE ``args`` override to close over.

        Returns:
            A `Prior`-conforming module wrapping this temporal prior.

        Examples:
            >>> import jax.numpy as jnp
            >>> from vardax import DynTrajectory, IdentityPrior, Prior
            >>> def decay(t, y, args):
            ...     return -y
            >>> ts = jnp.linspace(0.0, 0.5, 6)
            >>> bound = DynTrajectory(model=decay).bind(ts)
            >>> isinstance(bound, Prior)
            True
            >>> traj = DynTrajectory(model=decay)(jnp.ones(3), ts)
            >>> bool(jnp.allclose(bound(traj), traj, atol=1e-4))
            True

            A sequence that violates the dynamics is pulled back onto the flow, so
            its reconstruction differs from it:

            >>> bool(jnp.allclose(bound(traj + 1.0), traj + 1.0, atol=1e-4))
            False
        """
        return _BoundTemporalPrior(prior=self, ts=ts, params=params)

    def as_forward_model(self, dt: float) -> _DynamicalForwardModel:
        r"""Adapt this prior's dynamics to ``pipekit_cycle.ForwardModel``.

        The returned adapter exposes ``step(state, dt)``, ``dt``, and
        ``state_signature``, so the wrapped ODE can drive anything that
        consumes the pipekit forward-model seam,
        [`StrongFourDVar`][vardax.StrongFourDVar], ``pipekit_cycle.DACycle``,
        etc. (Decisions D8/D18). ``step`` integrates the ODE over
        $[0, \Delta t]$; the dynamics are therefore treated as autonomous
        (time-shift invariant), which holds for the standard testbeds
        (Lorenz-63/96) and any RHS that ignores ``t``.

        Args:
            dt: Default integration step $\Delta t$ advertised as the adapter's ``dt``.

        Returns:
            A ``ForwardModel``-conforming adapter around this prior.

        Examples:
            One step of exponential decay matches the closed form:

            >>> import jax.numpy as jnp
            >>> import vardax as vdx
            >>> def decay(t, y, args):
            ...     return -y
            >>> fwd = vdx.DynTrajectory(model=decay).as_forward_model(0.1)
            >>> fwd.dt, fwd.state_signature
            (0.1, None)
            >>> x = jnp.array([1.0, 2.0])
            >>> bool(jnp.allclose(fwd.step(x, 0.1), x * jnp.exp(-0.1), atol=1e-4))
            True

            A zero step leaves the state alone:

            >>> bool(jnp.allclose(fwd.step(x, 0.0), x))
            True
        """
        return _DynamicalForwardModel(prior=self, dt=dt)

reconstruct

reconstruct(
    x: Array, ts: Array, params: PyTree | None = None
) -> Array

Dynamics-consistent reconstruction of a state sequence.

Returns an array of the same shape as x such that \(\|x - \mathrm{reconstruct}(x, t_s)\|^2\) is the prior's dynamical residual, the temporal analogue of the autoencoder reconstruction used by the static Prior seam (Decision D18).

Parameters:

Name	Type	Description	Default
x	Array	State sequence.	required
ts	Array	Time coordinates \(t_s\) of the sequence.	required
params	PyTree | None	ODE args override.	None

Returns:

Type	Description
Array	The dynamics-consistent sequence, with the shape of x.

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

def reconstruct(
    self,
    x: Array,
    ts: Array,
    params: PyTree | None = None,
) -> Array:
    r"""Dynamics-consistent reconstruction of a state sequence.

    Returns an array of the *same shape* as ``x`` such that
    $\|x - \mathrm{reconstruct}(x, t_s)\|^2$ is the prior's dynamical residual,
    the temporal analogue of the autoencoder reconstruction used by the
    static [`Prior`][vardax.Prior] seam (Decision D18).

    Args:
        x: State sequence.
        ts: Time coordinates $t_s$ of the sequence.
        params: ODE ``args`` override.

    Returns:
        The dynamics-consistent sequence, with the shape of ``x``.
    """
    raise NotImplementedError

bind

bind(
    ts: Array, params: PyTree | None = None
) -> _BoundTemporalPrior

Bind a time grid, yielding a one-argument Prior adapter.

The returned callable satisfies the static Prior protocol (__call__(x) -> x_prior) by closing over ts, so a dynamical prior can be used anywhere a static prior is expected, e.g. as the prior_fn of variational_cost, turning the prior term into a weak-constraint dynamical residual.

Parameters:

Name	Type	Description	Default
ts	Array	Time coordinates \(t_s\) to close over.	required
params	PyTree | None	ODE args override to close over.	None

Returns:

Type	Description
_BoundTemporalPrior	A Prior-conforming module wrapping this temporal prior.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import DynTrajectory, IdentityPrior, Prior
>>> def decay(t, y, args):
...     return -y
>>> ts = jnp.linspace(0.0, 0.5, 6)
>>> bound = DynTrajectory(model=decay).bind(ts)
>>> isinstance(bound, Prior)
True
>>> traj = DynTrajectory(model=decay)(jnp.ones(3), ts)
>>> bool(jnp.allclose(bound(traj), traj, atol=1e-4))
True

A sequence that violates the dynamics is pulled back onto the flow, so its reconstruction differs from it:

>>> bool(jnp.allclose(bound(traj + 1.0), traj + 1.0, atol=1e-4))
False

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

def bind(
    self,
    ts: Array,
    params: PyTree | None = None,
) -> _BoundTemporalPrior:
    """Bind a time grid, yielding a one-argument `Prior` adapter.

    The returned callable satisfies the static
    [`Prior`][vardax.Prior] protocol (``__call__(x) -> x_prior``) by
    closing over ``ts``, so a dynamical prior can be used anywhere a
    static prior is expected, e.g. as the ``prior_fn`` of
    [`variational_cost`][vardax.variational_cost], turning the prior
    term into a weak-constraint dynamical residual.

    Args:
        ts: Time coordinates $t_s$ to close over.
        params: ODE ``args`` override to close over.

    Returns:
        A `Prior`-conforming module wrapping this temporal prior.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import DynTrajectory, IdentityPrior, Prior
        >>> def decay(t, y, args):
        ...     return -y
        >>> ts = jnp.linspace(0.0, 0.5, 6)
        >>> bound = DynTrajectory(model=decay).bind(ts)
        >>> isinstance(bound, Prior)
        True
        >>> traj = DynTrajectory(model=decay)(jnp.ones(3), ts)
        >>> bool(jnp.allclose(bound(traj), traj, atol=1e-4))
        True

        A sequence that violates the dynamics is pulled back onto the flow, so
        its reconstruction differs from it:

        >>> bool(jnp.allclose(bound(traj + 1.0), traj + 1.0, atol=1e-4))
        False
    """
    return _BoundTemporalPrior(prior=self, ts=ts, params=params)

as_forward_model

as_forward_model(dt: float) -> _DynamicalForwardModel

Adapt this prior's dynamics to pipekit_cycle.ForwardModel.

The returned adapter exposes step(state, dt), dt, and state_signature, so the wrapped ODE can drive anything that consumes the pipekit forward-model seam, StrongFourDVar, pipekit_cycle.DACycle, etc. (Decisions D8/D18). step integrates the ODE over \([0, \Delta t]\); the dynamics are therefore treated as autonomous (time-shift invariant), which holds for the standard testbeds (Lorenz-63/96) and any RHS that ignores t.

Parameters:

Name	Type	Description	Default
dt	float	Default integration step \(\Delta t\) advertised as the adapter's dt.	required

Returns:

Type	Description
_DynamicalForwardModel	A ForwardModel-conforming adapter around this prior.

Examples:

One step of exponential decay matches the closed form:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> def decay(t, y, args):
...     return -y
>>> fwd = vdx.DynTrajectory(model=decay).as_forward_model(0.1)
>>> fwd.dt, fwd.state_signature
(0.1, None)
>>> x = jnp.array([1.0, 2.0])
>>> bool(jnp.allclose(fwd.step(x, 0.1), x * jnp.exp(-0.1), atol=1e-4))
True

A zero step leaves the state alone:

>>> bool(jnp.allclose(fwd.step(x, 0.0), x))
True

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

def as_forward_model(self, dt: float) -> _DynamicalForwardModel:
    r"""Adapt this prior's dynamics to ``pipekit_cycle.ForwardModel``.

    The returned adapter exposes ``step(state, dt)``, ``dt``, and
    ``state_signature``, so the wrapped ODE can drive anything that
    consumes the pipekit forward-model seam,
    [`StrongFourDVar`][vardax.StrongFourDVar], ``pipekit_cycle.DACycle``,
    etc. (Decisions D8/D18). ``step`` integrates the ODE over
    $[0, \Delta t]$; the dynamics are therefore treated as autonomous
    (time-shift invariant), which holds for the standard testbeds
    (Lorenz-63/96) and any RHS that ignores ``t``.

    Args:
        dt: Default integration step $\Delta t$ advertised as the adapter's ``dt``.

    Returns:
        A ``ForwardModel``-conforming adapter around this prior.

    Examples:
        One step of exponential decay matches the closed form:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> def decay(t, y, args):
        ...     return -y
        >>> fwd = vdx.DynTrajectory(model=decay).as_forward_model(0.1)
        >>> fwd.dt, fwd.state_signature
        (0.1, None)
        >>> x = jnp.array([1.0, 2.0])
        >>> bool(jnp.allclose(fwd.step(x, 0.1), x * jnp.exp(-0.1), atol=1e-4))
        True

        A zero step leaves the state alone:

        >>> bool(jnp.allclose(fwd.step(x, 0.0), x))
        True
    """
    return _DynamicalForwardModel(prior=self, dt=dt)

DynIncrements

Bases: DynamicalPrior

One-step-increment dynamical prior.

The loss integrates each state a single step forward and compares it to the next observed state, so the dynamics act locally in time:

\[ R(u; \theta) = \sum_t \| u_{t+1} - \varphi_{\Delta t}(u_t; \theta) \|^2 . \]

Algorithm

1. Split the time grid into consecutive pairs \((t_t, t_{t+1})\) with time_patches.
2. Propagate every state \(u_t\) one step with the ODE solver, in parallel over \(t\).
3. Sum the squared differences to the next states.

Parameters:

Name	Type	Description	Default
model	Callable	Diffrax-compatible ODE right-hand side f(t, y, args) -> dy.	required
params	PyTree | None	Optional default args threaded to model (e.g. \(\theta\)).	None
solver	Solver | None	Diffrax solver (default :class:diffrax.Tsit5).	None
stepsize	StepSizeController | None	Diffrax step-size controller (adaptive by default for adaptive solvers).	None
adjoint	Adjoint | None	Diffrax adjoint strategy.	None

Raises:

Type	Description
ShapeError	From loss, if the lengths of x, ts and x_gt are inconsistent.

Examples:

A consistent trajectory (produced by rolling out the same model) has (near-)zero increment loss.

>>> import jax.numpy as jnp
>>> from vardax import DynIncrements, DynTrajectory
>>> def decay(t, y, args):
...     return -y
>>> ts = jnp.linspace(0.0, 0.5, 6)
>>> x0 = jnp.array([1.0, 2.0, -1.0])
>>> traj = DynTrajectory(model=decay)(x0, ts)
>>> bool(DynIncrements(model=decay).loss(traj, ts) < 1e-3)
True

A perturbed trajectory is penalised:

>>> bool(DynIncrements(model=decay).loss(traj + 1.0, ts) > 1e-3)
True

A time grid that does not match the sequence is rejected:

>>> DynIncrements(model=decay).loss(
...     traj[:3], ts
... )
Traceback (most recent call last):
...
ShapeError: Size mismatch for DynIncrements.loss: expected ...

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

class DynIncrements(DynamicalPrior):
    r"""One-step-increment dynamical prior.

    The loss integrates each state a *single* step forward and compares it
    to the next observed state, so the dynamics act locally in time:

    $$
    R(u; \theta) = \sum_t \| u_{t+1} - \varphi_{\Delta t}(u_t; \theta) \|^2 .
    $$

    Algorithm:
        1. Split the time grid into consecutive pairs
           $(t_t, t_{t+1})$ with [`time_patches`][vardax.time_patches].
        2. Propagate every state $u_t$ one step with the ODE solver, in
           parallel over $t$.
        3. Sum the squared differences to the next states.

    Args:
        model: Diffrax-compatible ODE right-hand side ``f(t, y, args) -> dy``.
        params: Optional default ``args`` threaded to ``model`` (e.g. $\theta$).
        solver: Diffrax solver (default :class:`diffrax.Tsit5`).
        stepsize: Diffrax step-size controller (adaptive by default for
            adaptive solvers).
        adjoint: Diffrax adjoint strategy.

    Raises:
        ShapeError: From ``loss``, if the lengths of ``x``, ``ts`` and ``x_gt``
            are inconsistent.

    Examples:
        A consistent trajectory (produced by rolling out the same model)
        has (near-)zero increment loss.

        >>> import jax.numpy as jnp
        >>> from vardax import DynIncrements, DynTrajectory
        >>> def decay(t, y, args):
        ...     return -y
        >>> ts = jnp.linspace(0.0, 0.5, 6)
        >>> x0 = jnp.array([1.0, 2.0, -1.0])
        >>> traj = DynTrajectory(model=decay)(x0, ts)
        >>> bool(DynIncrements(model=decay).loss(traj, ts) < 1e-3)
        True

        A perturbed trajectory is penalised:

        >>> bool(DynIncrements(model=decay).loss(traj + 1.0, ts) > 1e-3)
        True

        A time grid that does not match the sequence is rejected:

        >>> DynIncrements(model=decay).loss(
        ...     traj[:3], ts
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ShapeError: Size mismatch for DynIncrements.loss: expected ...
    """

    def __call__(
        self,
        x: Array,
        ts: Array,
        dt: float | None = None,
        params: PyTree | None = None,
    ) -> Array:
        """Integrate ``x`` one step from ``ts[0]`` to ``ts[-1]``.

        Args:
            x: State at ``ts[0]``.
            ts: Integration endpoints (or more; only the first and last are used).
            dt: Initial step size. Defaults to ``ts[1] - ts[0]``.
            params: ODE ``args`` override (defaults to ``self.params``).

        Returns:
            State at the final time.
        """
        saveat = dfx.SaveAt(ts=jnp.asarray([ts[-1]]))
        saved_states = self._integrate(x, ts, saveat, dt, params)
        return saved_states[-1]

    def loss(
        self,
        x: Array,
        ts: Array,
        x_gt: Array | None = None,
        params: PyTree | None = None,
    ) -> Array:
        r"""One-step-increment dynamical loss.

        Args:
            x: State sequence whose entries are the sources for the one-step
                predictions.
            ts: Time coordinates of the sequence.
            x_gt: Target state sequence. Defaults to ``x`` (self-consistency).
            params: ODE ``args`` override.

        Returns:
            Scalar dynamical residual $\sum_t \|\varphi(u_t) - u_{t+1}\|^2$.

        Raises:
            ShapeError: If ``len(x) - 1``, the number of time pairs and
                ``len(x_gt) - 1`` differ.
        """
        if x_gt is None:
            x_gt = x
        ts_pairs = time_patches(ts)
        if not len(x) - 1 == len(ts_pairs) == len(x_gt) - 1:
            msg = (
                "Size mismatch for DynIncrements.loss: expected "
                f"len(x) - 1 == len(time_patches(ts)) == len(x_gt) - 1, got "
                f"{len(x)} | {len(ts_pairs)} | {len(x_gt)}"
            )
            raise ShapeError(msg)
        propagate = ft.partial(self, params=params)
        predicted_states = jax.vmap(propagate, in_axes=(0, 0))(x[:-1], ts_pairs)
        # R = Σ_t ‖φ(u_t) − u_{t+1}‖²
        return jnp.sum((predicted_states - x_gt[1:]) ** 2)

    def reconstruct(
        self,
        x: Array,
        ts: Array,
        params: PyTree | None = None,
    ) -> Array:
        r"""One-step-increment reconstruction.

        ``reconstruct(x, ts)[t+1]`` is the one-step propagation
        $\varphi_{\Delta t}(x_t)$ and ``reconstruct(x, ts)[0] = x[0]``, so
        $\|x - \text{reconstruct}(x, t_s)\|^2$ equals the increment
        residual of [`loss`][vardax.DynIncrements.loss] (the $t=0$ term vanishes).

        Args:
            x: State sequence.
            ts: Time coordinates $t_s$ of the sequence.
            params: ODE ``args`` override.

        Returns:
            The sequence $[x_0, \varphi(x_0), \dots, \varphi(x_{T-2})]$.

        Examples:
            The residual of the reconstruction is the increment loss:

            >>> import jax.numpy as jnp
            >>> from vardax import DynIncrements
            >>> def decay(t, y, args):
            ...     return -y
            >>> prior = DynIncrements(model=decay)
            >>> ts = jnp.linspace(0.0, 0.5, 6)
            >>> x = jnp.ones((6, 2)) * jnp.arange(6.0)[:, None]
            >>> rec = prior.reconstruct(x, ts)
            >>> rec.shape == x.shape, bool(jnp.all(rec[0] == x[0]))
            (True, True)
            >>> bool(
            ...     jnp.allclose(jnp.sum((x - rec) ** 2), prior.loss(x, ts), rtol=1e-4)
            ... )
            True
        """
        ts_pairs = time_patches(ts)
        propagate = ft.partial(self, params=params)
        predicted_states = jax.vmap(propagate, in_axes=(0, 0))(x[:-1], ts_pairs)
        # x̂ = [x₀; φ(x₀), …, φ(x_{T−1})]
        # jnp, not einx: a one-state trajectory has no predictions to size
        return jnp.concatenate([x[:1], predicted_states])

loss

loss(
    x: Array,
    ts: Array,
    x_gt: Array | None = None,
    params: PyTree | None = None,
) -> Array

One-step-increment dynamical loss.

Parameters:

Name	Type	Description	Default
x	Array	State sequence whose entries are the sources for the one-step predictions.	required
ts	Array	Time coordinates of the sequence.	required
x_gt	Array | None	Target state sequence. Defaults to x (self-consistency).	None
params	PyTree | None	ODE args override.	None

Returns:

Type	Description
Array	Scalar dynamical residual \(\sum_t \|\varphi(u_t) - u_{t+1}\|^2\).

Raises:

Type	Description
ShapeError	If len(x) - 1, the number of time pairs and len(x_gt) - 1 differ.

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

def loss(
    self,
    x: Array,
    ts: Array,
    x_gt: Array | None = None,
    params: PyTree | None = None,
) -> Array:
    r"""One-step-increment dynamical loss.

    Args:
        x: State sequence whose entries are the sources for the one-step
            predictions.
        ts: Time coordinates of the sequence.
        x_gt: Target state sequence. Defaults to ``x`` (self-consistency).
        params: ODE ``args`` override.

    Returns:
        Scalar dynamical residual $\sum_t \|\varphi(u_t) - u_{t+1}\|^2$.

    Raises:
        ShapeError: If ``len(x) - 1``, the number of time pairs and
            ``len(x_gt) - 1`` differ.
    """
    if x_gt is None:
        x_gt = x
    ts_pairs = time_patches(ts)
    if not len(x) - 1 == len(ts_pairs) == len(x_gt) - 1:
        msg = (
            "Size mismatch for DynIncrements.loss: expected "
            f"len(x) - 1 == len(time_patches(ts)) == len(x_gt) - 1, got "
            f"{len(x)} | {len(ts_pairs)} | {len(x_gt)}"
        )
        raise ShapeError(msg)
    propagate = ft.partial(self, params=params)
    predicted_states = jax.vmap(propagate, in_axes=(0, 0))(x[:-1], ts_pairs)
    # R = Σ_t ‖φ(u_t) − u_{t+1}‖²
    return jnp.sum((predicted_states - x_gt[1:]) ** 2)

reconstruct

reconstruct(
    x: Array, ts: Array, params: PyTree | None = None
) -> Array

One-step-increment reconstruction.

reconstruct(x, ts)[t+1] is the one-step propagation \(\varphi_{\Delta t}(x_t)\) and reconstruct(x, ts)[0] = x[0], so \(\|x - \text{reconstruct}(x, t_s)\|^2\) equals the increment residual of loss (the \(t=0\) term vanishes).

Parameters:

Name	Type	Description	Default
x	Array	State sequence.	required
ts	Array	Time coordinates \(t_s\) of the sequence.	required
params	PyTree | None	ODE args override.	None

Returns:

Type	Description
Array	The sequence \([x_0, \varphi(x_0), \dots, \varphi(x_{T-2})]\).

Examples:

The residual of the reconstruction is the increment loss:

>>> import jax.numpy as jnp
>>> from vardax import DynIncrements
>>> def decay(t, y, args):
...     return -y
>>> prior = DynIncrements(model=decay)
>>> ts = jnp.linspace(0.0, 0.5, 6)
>>> x = jnp.ones((6, 2)) * jnp.arange(6.0)[:, None]
>>> rec = prior.reconstruct(x, ts)
>>> rec.shape == x.shape, bool(jnp.all(rec[0] == x[0]))
(True, True)
>>> bool(
...     jnp.allclose(jnp.sum((x - rec) ** 2), prior.loss(x, ts), rtol=1e-4)
... )
True

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

def reconstruct(
    self,
    x: Array,
    ts: Array,
    params: PyTree | None = None,
) -> Array:
    r"""One-step-increment reconstruction.

    ``reconstruct(x, ts)[t+1]`` is the one-step propagation
    $\varphi_{\Delta t}(x_t)$ and ``reconstruct(x, ts)[0] = x[0]``, so
    $\|x - \text{reconstruct}(x, t_s)\|^2$ equals the increment
    residual of [`loss`][vardax.DynIncrements.loss] (the $t=0$ term vanishes).

    Args:
        x: State sequence.
        ts: Time coordinates $t_s$ of the sequence.
        params: ODE ``args`` override.

    Returns:
        The sequence $[x_0, \varphi(x_0), \dots, \varphi(x_{T-2})]$.

    Examples:
        The residual of the reconstruction is the increment loss:

        >>> import jax.numpy as jnp
        >>> from vardax import DynIncrements
        >>> def decay(t, y, args):
        ...     return -y
        >>> prior = DynIncrements(model=decay)
        >>> ts = jnp.linspace(0.0, 0.5, 6)
        >>> x = jnp.ones((6, 2)) * jnp.arange(6.0)[:, None]
        >>> rec = prior.reconstruct(x, ts)
        >>> rec.shape == x.shape, bool(jnp.all(rec[0] == x[0]))
        (True, True)
        >>> bool(
        ...     jnp.allclose(jnp.sum((x - rec) ** 2), prior.loss(x, ts), rtol=1e-4)
        ... )
        True
    """
    ts_pairs = time_patches(ts)
    propagate = ft.partial(self, params=params)
    predicted_states = jax.vmap(propagate, in_axes=(0, 0))(x[:-1], ts_pairs)
    # x̂ = [x₀; φ(x₀), …, φ(x_{T−1})]
    # jnp, not einx: a one-state trajectory has no predictions to size
    return jnp.concatenate([x[:1], predicted_states])

DynTrajectory

Bases: DynamicalPrior

Full-rollout dynamical prior.

The loss integrates the initial state across the whole window and compares the resulting trajectory to the target sequence, enforcing the dynamics as a global (strong) constraint:

\[ R(u; \theta) = \sum_t \| u_t - \varphi_t(u_0; \theta) \|^2 . \]

This is the propagation used by strong-constraint 4DVar, see strong_variational_cost.

Algorithm

1. Integrate the ODE from the initial state \(u_0\), saving the state at every time in ts.
2. Sum the squared differences between the saved trajectory and the target sequence.

Parameters:

Name	Type	Description	Default
model	Callable	Diffrax-compatible ODE right-hand side f(t, y, args) -> dy.	required
params	PyTree | None	Optional default args threaded to model (e.g. \(\theta\)).	None
solver	Solver | None	Diffrax solver (default :class:diffrax.Tsit5).	None
stepsize	StepSizeController | None	Diffrax step-size controller (adaptive by default for adaptive solvers).	None
adjoint	Adjoint | None	Diffrax adjoint strategy.	None

Examples:

Roll out exponential decay from an initial condition:

>>> import jax.numpy as jnp
>>> from vardax import DynTrajectory
>>> def decay(t, y, args):
...     return -y
>>> ts = jnp.linspace(0.0, 0.5, 6)
>>> prior = DynTrajectory(model=decay)
>>> traj = prior(jnp.array([1.0, 2.0]), ts)
>>> traj.shape
(6, 2)
>>> bool(
...     jnp.allclose(traj[-1], jnp.array([1.0, 2.0]) * jnp.exp(-0.5), atol=1e-3)
... )
True

A trajectory generated by the model has (near-)zero loss, and a shifted one does not:

>>> bool(prior.loss(traj, ts) < 1e-6)
True
>>> bool(prior.loss(traj + 1.0, ts, x_gt=traj) > 1.0)
True

Through bind it plugs into the static prior seam:

>>> bool(jnp.allclose(prior.bind(ts)(traj), traj, atol=1e-4))
True

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

class DynTrajectory(DynamicalPrior):
    r"""Full-rollout dynamical prior.

    The loss integrates the initial state across the whole window and
    compares the resulting trajectory to the target sequence, enforcing the
    dynamics as a global (strong) constraint:

    $$
    R(u; \theta) = \sum_t \| u_t - \varphi_t(u_0; \theta) \|^2 .
    $$

    This is the propagation used by strong-constraint 4DVar, see
    [`strong_variational_cost`][vardax.strong_variational_cost].

    Algorithm:
        1. Integrate the ODE from the initial state $u_0$, saving the state at
           every time in ``ts``.
        2. Sum the squared differences between the saved trajectory and the
           target sequence.

    Args:
        model: Diffrax-compatible ODE right-hand side ``f(t, y, args) -> dy``.
        params: Optional default ``args`` threaded to ``model`` (e.g. $\theta$).
        solver: Diffrax solver (default :class:`diffrax.Tsit5`).
        stepsize: Diffrax step-size controller (adaptive by default for
            adaptive solvers).
        adjoint: Diffrax adjoint strategy.

    Examples:
        Roll out exponential decay from an initial condition:

        >>> import jax.numpy as jnp
        >>> from vardax import DynTrajectory
        >>> def decay(t, y, args):
        ...     return -y
        >>> ts = jnp.linspace(0.0, 0.5, 6)
        >>> prior = DynTrajectory(model=decay)
        >>> traj = prior(jnp.array([1.0, 2.0]), ts)
        >>> traj.shape
        (6, 2)
        >>> bool(
        ...     jnp.allclose(traj[-1], jnp.array([1.0, 2.0]) * jnp.exp(-0.5), atol=1e-3)
        ... )
        True

        A trajectory generated by the model has (near-)zero loss, and a shifted
        one does not:

        >>> bool(prior.loss(traj, ts) < 1e-6)
        True
        >>> bool(prior.loss(traj + 1.0, ts, x_gt=traj) > 1.0)
        True

        Through [`bind`][vardax.DynamicalPrior.bind] it plugs into the static
        prior seam:

        >>> bool(jnp.allclose(prior.bind(ts)(traj), traj, atol=1e-4))
        True
    """

    def __call__(
        self,
        x: Array,
        ts: Array,
        dt: float | None = None,
        params: PyTree | None = None,
    ) -> Array:
        """Integrate the initial state ``x`` across all of ``ts``.

        Args:
            x: Initial condition (state at ``ts[0]``).
            ts: Time coordinates at which to save the state.
            dt: Initial step size. Defaults to ``ts[1] - ts[0]``.
            params: ODE ``args`` override.

        Returns:
            State trajectory with ``traj[0] == x``, one row per entry of ``ts``.
        """
        saveat = dfx.SaveAt(ts=ts)
        return self._integrate(x, ts, saveat, dt, params)

    def loss(
        self,
        x: Array,
        ts: Array,
        x_gt: Array | None = None,
        params: PyTree | None = None,
    ) -> Array:
        r"""Full-trajectory dynamical loss.

        Args:
            x: State sequence; the initial condition ``x[0]`` seeds the rollout.
            ts: Time coordinates of the sequence.
            x_gt: Target trajectory. Defaults to ``x``.
            params: ODE ``args`` override.

        Returns:
            Scalar dynamical residual $\sum_t \|\varphi_t(u_0) - u_t\|^2$.
        """
        if x_gt is None:
            x_gt = x
        predicted_states = self(x[0], ts, params=params)
        # R = Σ_t ‖φ_t(u_0) − u_t‖²
        return jnp.sum((predicted_states - x_gt) ** 2)

    def reconstruct(
        self,
        x: Array,
        ts: Array,
        params: PyTree | None = None,
    ) -> Array:
        r"""Full-rollout reconstruction: the trajectory from ``x[0]``.

        $\|x - \text{reconstruct}(x, t_s)\|^2$ equals the trajectory
        residual of [`loss`][vardax.DynTrajectory.loss].

        Args:
            x: State sequence; ``x[0]`` seeds the rollout.
            ts: Time coordinates $t_s$ of the sequence.
            params: ODE ``args`` override.

        Returns:
            The rolled-out trajectory, same shape as ``x``.
        """
        return self(x[0], ts, params=params)

loss

loss(
    x: Array,
    ts: Array,
    x_gt: Array | None = None,
    params: PyTree | None = None,
) -> Array

Full-trajectory dynamical loss.

Parameters:

Name	Type	Description	Default
x	Array	State sequence; the initial condition x[0] seeds the rollout.	required
ts	Array	Time coordinates of the sequence.	required
x_gt	Array | None	Target trajectory. Defaults to x.	None
params	PyTree | None	ODE args override.	None

Returns:

Type	Description
Array	Scalar dynamical residual \(\sum_t \|\varphi_t(u_0) - u_t\|^2\).

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

def loss(
    self,
    x: Array,
    ts: Array,
    x_gt: Array | None = None,
    params: PyTree | None = None,
) -> Array:
    r"""Full-trajectory dynamical loss.

    Args:
        x: State sequence; the initial condition ``x[0]`` seeds the rollout.
        ts: Time coordinates of the sequence.
        x_gt: Target trajectory. Defaults to ``x``.
        params: ODE ``args`` override.

    Returns:
        Scalar dynamical residual $\sum_t \|\varphi_t(u_0) - u_t\|^2$.
    """
    if x_gt is None:
        x_gt = x
    predicted_states = self(x[0], ts, params=params)
    # R = Σ_t ‖φ_t(u_0) − u_t‖²
    return jnp.sum((predicted_states - x_gt) ** 2)

reconstruct

reconstruct(
    x: Array, ts: Array, params: PyTree | None = None
) -> Array

Full-rollout reconstruction: the trajectory from x[0].

\(\|x - \text{reconstruct}(x, t_s)\|^2\) equals the trajectory residual of loss.

Parameters:

Name	Type	Description	Default
x	Array	State sequence; x[0] seeds the rollout.	required
ts	Array	Time coordinates \(t_s\) of the sequence.	required
params	PyTree | None	ODE args override.	None

Returns:

Type	Description
Array	The rolled-out trajectory, same shape as x.

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

def reconstruct(
    self,
    x: Array,
    ts: Array,
    params: PyTree | None = None,
) -> Array:
    r"""Full-rollout reconstruction: the trajectory from ``x[0]``.

    $\|x - \text{reconstruct}(x, t_s)\|^2$ equals the trajectory
    residual of [`loss`][vardax.DynTrajectory.loss].

    Args:
        x: State sequence; ``x[0]`` seeds the rollout.
        ts: Time coordinates $t_s$ of the sequence.
        params: ODE ``args`` override.

    Returns:
        The rolled-out trajectory, same shape as ``x``.
    """
    return self(x[0], ts, params=params)

Reduced bases

Control-variable transforms: a ReducedBasis parameterises the analysis increment in a lower-dimensional space, so the inner solve runs over basis coefficients rather than grid points. The constructors build the common families; CompositeBasis concatenates several.

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

LinearBasis

Bases: Module

Generic linear reduced basis with a diagonal-\(Q\) prior.

Build via :func:linear_basis (or a family constructor: :func:vardax.rbf_basis, :func:vardax.fourier_basis, :func:vardax.wavelet_basis, :func:vardax.eof_basis); direct instantiation expects the precomputed design_matrix (N, M) matrix, variance (M,), and the grid output_shape with prod(output_shape) == N.

Satisfies vardax.protocols.ReducedBasis (operg / prior_inv / nbasis). Stationary: operg ignores the t argument.

\[\delta x = \Phi X, \quad Q = \mathrm{diag}(\sigma^2)\]

Examples:

A hand-built design matrix with a non-unit prior variance:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> basis = vdx.linear_basis(
...     jnp.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]), variance=2.0
... )
>>> basis.operg(0.0, jnp.array([1.0, 2.0])).shape
(3,)
>>> bool(jnp.allclose(basis.prior_inv(jnp.ones(2)), 0.5))
True

An output_shape reshapes the increment onto a grid:

>>> gridded = vdx.linear_basis(jnp.eye(4), output_shape=(2, 2))
>>> gridded.operg(0.0, jnp.ones(4)).shape
(2, 2)

Source code in src/vardax/_src/basis/linear.py

class LinearBasis(eqx.Module):
    r"""Generic linear reduced basis with a diagonal-$Q$ prior.

    Build via :func:`linear_basis` (or a family constructor:
    :func:`vardax.rbf_basis`, :func:`vardax.fourier_basis`,
    :func:`vardax.wavelet_basis`, :func:`vardax.eof_basis`); direct
    instantiation expects the precomputed ``design_matrix`` ``(N, M)`` matrix,
    ``variance`` ``(M,)``, and the grid ``output_shape`` with
    ``prod(output_shape) == N``.

    Satisfies `vardax.protocols.ReducedBasis` (`operg` / `prior_inv` /
    `nbasis`). Stationary: ``operg`` ignores the ``t`` argument.

    $$\delta x = \Phi X, \quad Q = \mathrm{diag}(\sigma^2)$$

    Examples:
        A hand-built design matrix with a non-unit prior variance:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> basis = vdx.linear_basis(
        ...     jnp.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]), variance=2.0
        ... )
        >>> basis.operg(0.0, jnp.array([1.0, 2.0])).shape
        (3,)
        >>> bool(jnp.allclose(basis.prior_inv(jnp.ones(2)), 0.5))
        True

        An ``output_shape`` reshapes the increment onto a grid:

        >>> gridded = vdx.linear_basis(jnp.eye(4), output_shape=(2, 2))
        >>> gridded.operg(0.0, jnp.ones(4)).shape
        (2, 2)
    """

    design_matrix: Float[Array, "N M"]
    variance: Float[Array, " M"]
    output_shape: tuple[int, ...] = eqx.field(static=True)

    def operg(
        self,
        t: float,
        coefficients: Float[Array, " M"],
        state: Float[Array, ...] | None = None,
    ) -> Float[Array, ...]:
        r"""Apply $\Phi X$ and return the increment on ``output_shape``.

        Args:
            t: Time; ignored because the basis is stationary.
            coefficients: Coefficient vector.
            state: Current state; ignored by this basis.

        Returns:
            State increment $\Phi X$ reshaped onto ``output_shape``.
        """
        # δx = Φ X
        increment = einx.dot("n m, m -> n", self.design_matrix, coefficients)
        return einx.id("(s...) -> s...", increment, s=self.output_shape)

    def prior_inv(self, coefficients: Float[Array, " M"]) -> Float[Array, " M"]:
        r"""Apply $Q^{-1} X$ with diagonal $Q = \mathrm{diag}(\sigma^2)$.

        Args:
            coefficients: Coefficient vector.

        Returns:
            The vector $X / \sigma^2$.
        """
        return coefficients / self.variance

    @property
    def nbasis(self) -> int:
        return int(self.design_matrix.shape[1])

operg

operg(
    t: float,
    coefficients: Float[Array, " M"],
    state: Float[Array, ...] | None = None,
) -> Float[Array, ...]

Apply \(\Phi X\) and return the increment on output_shape.

Parameters:

Name	Type	Description	Default
t	float	Time; ignored because the basis is stationary.	required
coefficients	Float[Array, ' M']	Coefficient vector.	required
state	Float[Array, ...] | None	Current state; ignored by this basis.	None

Returns:

Type	Description
Float[Array, ...]	State increment \(\Phi X\) reshaped onto output_shape.

Source code in src/vardax/_src/basis/linear.py

def operg(
    self,
    t: float,
    coefficients: Float[Array, " M"],
    state: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    r"""Apply $\Phi X$ and return the increment on ``output_shape``.

    Args:
        t: Time; ignored because the basis is stationary.
        coefficients: Coefficient vector.
        state: Current state; ignored by this basis.

    Returns:
        State increment $\Phi X$ reshaped onto ``output_shape``.
    """
    # δx = Φ X
    increment = einx.dot("n m, m -> n", self.design_matrix, coefficients)
    return einx.id("(s...) -> s...", increment, s=self.output_shape)

prior_inv

prior_inv(
    coefficients: Float[Array, " M"],
) -> Float[Array, " M"]

Apply \(Q^{-1} X\) with diagonal \(Q = \mathrm{diag}(\sigma^2)\).

Parameters:

Name	Type	Description	Default
coefficients	Float[Array, ' M']	Coefficient vector.	required

Returns:

Type	Description
Float[Array, ' M']	The vector \(X / \sigma^2\).

Source code in src/vardax/_src/basis/linear.py

def prior_inv(self, coefficients: Float[Array, " M"]) -> Float[Array, " M"]:
    r"""Apply $Q^{-1} X$ with diagonal $Q = \mathrm{diag}(\sigma^2)$.

    Args:
        coefficients: Coefficient vector.

    Returns:
        The vector $X / \sigma^2$.
    """
    return coefficients / self.variance

CompositeBasis

Bases: Module

Concatenate several reduced bases into one control vector.

Satisfies vardax.protocols.ReducedBasis. nbasis is the sum of component sizes; operg slices coefficients, applies each component, sums the per-grid increments; prior_inv slices coefficients and concatenates the per-component blocks. Components may mix families (e.g. a broad rbf_basis layer plus a wavelet_basis detail layer), provided their operg outputs share one grid shape.

\[\delta x = \sum_c \Phi_c X_c\]

\[Q^{-1} = \mathrm{blockdiag}(Q_1^{-1}, \ldots, Q_C^{-1})\]

Algorithm

1. Split \(X\) at the cumulative component sizes (splits).
2. operg: apply each component to its slice and sum the grid increments.
3. prior_inv: apply each component's prior_inv to its slice and concatenate the blocks.

Examples:

Two linear bases on a 3-point grid, stacked into one control vector:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> b1 = vdx.linear_basis(jnp.eye(3)[:, :2], variance=2.0)
>>> b2 = vdx.linear_basis(jnp.ones((3, 1)), variance=0.5)
>>> basis = vdx.composite_basis(b1, b2)
>>> basis.nbasis
3
>>> basis.operg(0.0, jnp.array([1.0, 2.0, 3.0])).shape
(3,)

The prior is block-diagonal, so prior_inv divides each block by its own variance:

>>> bool(jnp.allclose(basis.prior_inv(jnp.ones(3)), jnp.array([0.5, 0.5, 2.0])))
True

Source code in src/vardax/_src/basis/composite.py

class CompositeBasis(eqx.Module):
    r"""Concatenate several reduced bases into one control vector.

    Satisfies `vardax.protocols.ReducedBasis`. ``nbasis`` is the sum of
    component sizes; ``operg`` slices ``coefficients``, applies each component,
    sums the per-grid increments; ``prior_inv`` slices ``coefficients`` and
    concatenates the per-component blocks. Components may mix families
    (e.g. a broad `rbf_basis` layer plus a `wavelet_basis` detail
    layer), provided their ``operg`` outputs share one grid shape.

    $$\delta x = \sum_c \Phi_c X_c$$

    $$Q^{-1} = \mathrm{blockdiag}(Q_1^{-1}, \ldots, Q_C^{-1})$$

    Algorithm:
        1. Split $X$ at the cumulative component sizes (``splits``).
        2. ``operg``: apply each component to its slice and sum the grid
           increments.
        3. ``prior_inv``: apply each component's ``prior_inv`` to its slice
           and concatenate the blocks.

    Examples:
        Two linear bases on a 3-point grid, stacked into one control vector:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> b1 = vdx.linear_basis(jnp.eye(3)[:, :2], variance=2.0)
        >>> b2 = vdx.linear_basis(jnp.ones((3, 1)), variance=0.5)
        >>> basis = vdx.composite_basis(b1, b2)
        >>> basis.nbasis
        3
        >>> basis.operg(0.0, jnp.array([1.0, 2.0, 3.0])).shape
        (3,)

        The prior is block-diagonal, so ``prior_inv`` divides each block by
        its own variance:

        >>> bool(jnp.allclose(basis.prior_inv(jnp.ones(3)), jnp.array([0.5, 0.5, 2.0])))
        True
    """

    components: tuple[Any, ...]
    splits: tuple[int, ...] = eqx.field(static=True)

    def operg(
        self,
        t: float,
        coefficients: Float[Array, " M"],
        state: Float[Array, ...] | None = None,
    ) -> Float[Array, ...]:
        parts = self._split(coefficients)
        out = self.components[0].operg(t, parts[0], state)
        for comp, part in zip(self.components[1:], parts[1:], strict=True):
            out = out + comp.operg(t, part, state)
        return out

    def prior_inv(self, coefficients: Float[Array, " M"]) -> Float[Array, " M"]:
        parts = self._split(coefficients)
        blocks = [c.prior_inv(p) for c, p in zip(self.components, parts, strict=True)]
        return jnp.concatenate(blocks)

    @property
    def nbasis(self) -> int:
        return int(sum(c.nbasis for c in self.components))

    def _split(self, coefficients: Float[Array, " M"]) -> list[Float[Array, ...]]:
        parts, start = [], 0
        for cut in self.splits:
            parts.append(coefficients[start:cut])
            start = cut
        parts.append(coefficients[start:])
        return parts

linear_basis

linear_basis(
    design_matrix: Float[ArrayLike, "N M"],
    *,
    variance: float | Float[ArrayLike, " M"] = 1.0,
    output_shape: tuple[int, ...] | None = None,
) -> LinearBasis

Build a LinearBasis from an arbitrary precomputed design matrix.

The escape hatch for basis families without a dedicated constructor: any design matrix (spherical harmonics, Slepian, Gabor frames, hand-rolled …) becomes a ReducedBasis.

Parameters:

Name	Type	Description	Default
design_matrix	Float[ArrayLike, 'N M']	Design matrix \(\Phi\) mapping coefficients to the flattened state.	required
variance	float | Float[ArrayLike, ' M']	Scalar or per-coefficient prior variance.	1.0
output_shape	tuple[int, ...] | None	Grid shape the increment is reshaped to; must satisfy prod(output_shape) == N. Defaults to (N,).	None

Returns:

Type	Description
LinearBasis	Configured LinearBasis.

Raises:

Type	Description
ShapeError	If design_matrix is not 2-D or output_shape does not multiply to the number of rows.
ConfigurationError	If variance is not finite and positive.

Examples:

Wrap a matrix; the default grid shape is (N,):

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> basis = vdx.linear_basis(jnp.ones((3, 2)))
>>> basis.nbasis, basis.output_shape
(2, (3,))

A mismatched grid shape is rejected:

>>> vdx.linear_basis(
...     jnp.ones((3, 2)), output_shape=(2, 2)
... )
Traceback (most recent call last):
...
ShapeError: prod(output_shape)=4 does not match design_matrix rows N=3.

Source code in src/vardax/_src/basis/linear.py

def linear_basis(
    design_matrix: Float[ArrayLike, "N M"],
    *,
    variance: float | Float[ArrayLike, " M"] = 1.0,
    output_shape: tuple[int, ...] | None = None,
) -> LinearBasis:
    r"""Build a `LinearBasis` from an arbitrary precomputed design matrix.

    The escape hatch for basis families without a dedicated constructor:
    any design matrix (spherical harmonics, Slepian, Gabor frames,
    hand-rolled …) becomes a `ReducedBasis`.

    Args:
        design_matrix: Design matrix $\Phi$ mapping coefficients to the flattened
            state.
        variance: Scalar or per-coefficient prior variance.
        output_shape: Grid shape the increment is reshaped to; must
            satisfy ``prod(output_shape) == N``. Defaults to ``(N,)``.

    Returns:
        Configured `LinearBasis`.

    Raises:
        ShapeError: If ``design_matrix`` is not 2-D or ``output_shape`` does not
            multiply to the number of rows.
        ConfigurationError: If ``variance`` is not finite and positive.

    Examples:
        Wrap a matrix; the default grid shape is ``(N,)``:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> basis = vdx.linear_basis(jnp.ones((3, 2)))
        >>> basis.nbasis, basis.output_shape
        (2, (3,))

        A mismatched grid shape is rejected:

        >>> vdx.linear_basis(
        ...     jnp.ones((3, 2)), output_shape=(2, 2)
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ShapeError: prod(output_shape)=4 does not match design_matrix rows N=3.
    """
    design_matrix = jnp.asarray(design_matrix)
    if design_matrix.ndim != 2:
        raise ShapeError(
            f"design_matrix must be 2-D (N, M); got shape {design_matrix.shape}."
        )
    state_dim, n_atoms = int(design_matrix.shape[0]), int(design_matrix.shape[1])
    shape = (
        (state_dim,)
        if output_shape is None
        else tuple(int(extent) for extent in output_shape)
    )
    if math.prod(shape) != state_dim:
        raise ShapeError(
            f"prod(output_shape)={math.prod(shape)} does not match design_matrix rows "
            f"N={state_dim}."
        )
    return LinearBasis(
        design_matrix=design_matrix,
        # the prior variance follows Φ's dtype, so x64 callers keep float64
        variance=jnp.asarray(
            _as_variance(variance, n_atoms),
            dtype=jnp.result_type(design_matrix, float),
        ),
        output_shape=shape,
    )

composite_basis

composite_basis(*components: Any) -> CompositeBasis

Build a CompositeBasis from a sequence of ReducedBasis components.

Parameters:

Name	Type	Description	Default
*components	Any	Reduced bases to concatenate, in control-vector order. Their operg outputs must share one grid shape.	()

Returns:

Type	Description
CompositeBasis	The composite basis, with nbasis the sum of component sizes.

Raises:

Type	Description
ConfigurationError	If no component is given.

Examples:

Composing two bases and checking the control-vector size:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> a = vdx.linear_basis(jnp.eye(2))
>>> b = vdx.linear_basis(jnp.ones((2, 1)))
>>> vdx.composite_basis(a, b).nbasis
3

An empty call is rejected:

>>> vdx.composite_basis()
Traceback (most recent call last):
...
ConfigurationError: composite_basis requires at least one component.

Source code in src/vardax/_src/basis/composite.py

def composite_basis(*components: Any) -> CompositeBasis:
    """Build a `CompositeBasis` from a sequence of `ReducedBasis` components.

    Args:
        *components: Reduced bases to concatenate, in control-vector order.
            Their ``operg`` outputs must share one grid shape.

    Returns:
        The composite basis, with ``nbasis`` the sum of component sizes.

    Raises:
        ConfigurationError: If no component is given.

    Examples:
        Composing two bases and checking the control-vector size:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> a = vdx.linear_basis(jnp.eye(2))
        >>> b = vdx.linear_basis(jnp.ones((2, 1)))
        >>> vdx.composite_basis(a, b).nbasis
        3

        An empty call is rejected:

        >>> vdx.composite_basis()  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: composite_basis requires at least one component.
    """
    if not components:
        raise ConfigurationError("composite_basis requires at least one component.")
    cuts, total = [], 0
    for c in components[:-1]:
        total += c.nbasis
        cuts.append(total)
    return CompositeBasis(components=tuple(components), splits=tuple(cuts))

eof_basis

eof_basis(
    data: Float[ArrayLike, "T N"],
    n_modes: int,
    *,
    center: bool = True,
    variance: float
    | Float[ArrayLike, " M"]
    | str = "explained",
    output_shape: tuple[int, ...] | None = None,
) -> LinearBasis

Data-driven EOF (PCA) reduced basis from a sample matrix.

The leading n_modes empirical orthogonal functions of data become the columns of \(\Phi\). With variance="explained" the prior variance of mode \(a\) is its sample variance \(\sigma_a^2 / (T - 1)\), so the background term matches the climatology of the training samples.

Parameters:

Name	Type	Description	Default
data	Float[ArrayLike, 'T N']	Sample matrix, T snapshots over the flattened grid.	required
n_modes	int	Number of leading EOFs to keep.	required
center	bool	Subtract the per-point sample mean before the SVD.	True
variance	float | Float[ArrayLike, ' M'] | str	"explained" (default), or a scalar or per-coefficient array overriding the prior variance.	'explained'
output_shape	tuple[int, ...] | None	Grid shape to reshape increments to; must satisfy prod(output_shape) == N. Defaults to (N,).	None

Algorithm

1. Optionally subtract the per-point sample mean of data.
2. Take the SVD and keep the leading n_modes right singular vectors as \(\Phi\) with singular values \(\sigma_a\).
3. Set the prior variance to \(\sigma_a^2 / (T - 1)\) when variance="explained", otherwise use variance.

Returns:

Type	Description
LinearBasis	LinearBasis with orthonormal columns and nbasis == n_modes.

Raises:

Type	Description
ConfigurationError	If variance is a string other than "explained".
DataError	If variance="explained" and data has fewer than two samples.

Examples:

Two leading EOFs of ten snapshots over four grid points:

>>> import jax
>>> import vardax as vdx
>>> data = jax.random.normal(jax.random.PRNGKey(0), (10, 4))
>>> basis = vdx.eof_basis(data, 2)
>>> basis.design_matrix.shape, basis.variance.shape
((4, 2), (2,))

A single snapshot cannot give an explained variance:

>>> vdx.eof_basis(data[:1], 1)
Traceback (most recent call last):
...
DataError: variance="explained" needs at least 2 samples (T >= 2); got T=1.

Source code in src/vardax/_src/basis/constructors.py

def eof_basis(
    data: Float[ArrayLike, "T N"],
    n_modes: int,
    *,
    center: bool = True,
    variance: float | Float[ArrayLike, " M"] | str = "explained",
    output_shape: tuple[int, ...] | None = None,
) -> LinearBasis:
    r"""Data-driven EOF (PCA) reduced basis from a sample matrix.

    The leading ``n_modes`` empirical orthogonal functions of ``data``
    become the columns of $\Phi$. With ``variance="explained"`` the
    prior variance of mode $a$ is its sample variance
    $\sigma_a^2 / (T - 1)$, so the background term matches the
    climatology of the training samples.

    Args:
        data: Sample matrix, ``T`` snapshots over the flattened grid.
        n_modes: Number of leading EOFs to keep.
        center: Subtract the per-point sample mean before the SVD.
        variance: ``"explained"`` (default), or a scalar or
            per-coefficient array overriding the prior variance.
        output_shape: Grid shape to reshape increments to; must satisfy
            ``prod(output_shape) == N``. Defaults to ``(N,)``.

    Algorithm:
        1. Optionally subtract the per-point sample mean of ``data``.
        2. Take the SVD and keep the leading ``n_modes`` right singular
           vectors as $\Phi$ with singular values $\sigma_a$.
        3. Set the prior variance to $\sigma_a^2 / (T - 1)$ when
           ``variance="explained"``, otherwise use ``variance``.

    Returns:
        `LinearBasis` with orthonormal columns and ``nbasis == n_modes``.

    Raises:
        ConfigurationError: If ``variance`` is a string other than
            ``"explained"``.
        DataError: If ``variance="explained"`` and ``data`` has fewer than
            two samples.

    Examples:
        Two leading EOFs of ten snapshots over four grid points:

        >>> import jax
        >>> import vardax as vdx
        >>> data = jax.random.normal(jax.random.PRNGKey(0), (10, 4))
        >>> basis = vdx.eof_basis(data, 2)
        >>> basis.design_matrix.shape, basis.variance.shape
        ((4, 2), (2,))

        A single snapshot cannot give an explained variance:

        >>> vdx.eof_basis(data[:1], 1)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        DataError: variance="explained" needs at least 2 samples (T >= 2); got T=1.
    """
    data_array = jnp.asarray(data)
    # Φ: leading EOFs of the data, singular values σ
    design_matrix, singular_values = _gnx.eof_basis(data_array, n_modes, center=center)
    if isinstance(variance, str):
        if variance != "explained":
            raise ConfigurationError(
                f'variance must be "explained", a scalar, or an array; '
                f"got {variance!r}."
            )
        n_samples = int(data_array.shape[0])
        if n_samples < 2:
            raise DataError(
                'variance="explained" needs at least 2 samples (T >= 2); '
                f"got T={n_samples}."
            )
        # σ_a² / (T − 1)
        prior_variance: float | Float[Array, " M"] = singular_values**2 / (
            n_samples - 1
        )
    else:
        prior_variance = variance
    return linear_basis(
        design_matrix, variance=prior_variance, output_shape=output_shape
    )

fourier_basis

fourier_basis(
    grid_coords: Sequence[Float[ArrayLike, " _"]],
    num_modes: int | tuple[int, ...],
    *,
    boundary_scale: float = 1.25,
    variance: float | Float[ArrayLike, " M"] = 1.0,
    spectral_density: Callable[
        [Float[Array, " M"]], Float[Array, " M"]
    ]
    | None = None,
) -> LinearBasis

Laplacian-eigenfunction (Dirichlet/HSGP) reduced basis on a box.

Tensor-product sine eigenfunctions of \(-\Delta\) on the (recentred, slightly enlarged) grid bounding box. Passing spectral_density evaluates it at the Laplacian eigenvalues \(\lambda_j\) and scales it by variance — supply a kernel's spectral density to obtain the Hilbert-space GP (HSGP) approximation of that kernel's prior with amplitude variance.

Parameters:

Name	Type	Description	Default
grid_coords	Sequence[Float[ArrayLike, ' _']]	Per-axis 1-D coordinate arrays; any dimension d >= 1. Each axis must span a non-degenerate interval.	required
num_modes	int | tuple[int, ...]	Number of 1-D modes per axis; an int is broadcast to all axes. Total nbasis is the product.	required
boundary_scale	float	Domain half-widths are the grid half-spans times this factor (>= 1). The Dirichlet eigenfunctions vanish on the box boundary, so a value > 1 keeps the basis expressive at the grid edges.	1.25
variance	float | Float[ArrayLike, ' M']	Scalar or per-coefficient prior variance; with spectral_density it acts as a scalar amplitude.	1.0
spectral_density	Callable[[Float[Array, ' M']], Float[Array, ' M']] | None	Optional callable mapping the eigenvalues \(\lambda_j\) to per-mode prior power (HSGP prior).	None

Algorithm

1. Flatten the grid, recentre it, and enlarge the half-spans by boundary_scale.
2. Build the tensor-product Dirichlet eigenfunctions \(\Phi\) and eigenvalues \(\lambda\) with num_modes modes per axis.
3. Set the prior variance to variance or, with a spectral density, to variance * spectral_density(lambda).

Returns:

Type	Description
LinearBasis	LinearBasis with output_shape the grid shape.

Raises:

Type	Description
ConfigurationError	If boundary_scale < 1.
GridError	If an axis is empty, not 1-D, or has zero span.

Examples:

Four sine modes on a 1-D grid:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> x = jnp.linspace(0.0, 1.0, 6)
>>> basis = vdx.fourier_basis([x], 4)
>>> basis.nbasis, basis.output_shape
(4, (6,))

A spectral density turns the prior into an HSGP prior with decaying variance:

>>> hsgp = vdx.fourier_basis(
...     [x], 4, spectral_density=lambda lam: 1.0 / (1.0 + lam)
... )
>>> bool(hsgp.variance[0] > hsgp.variance[-1])
True

A shrinking boundary scale is rejected:

>>> vdx.fourier_basis(
...     [x], 4, boundary_scale=0.5
... )
Traceback (most recent call last):
...
ConfigurationError: boundary_scale must be >= 1, got 0.5.

Source code in src/vardax/_src/basis/constructors.py

def fourier_basis(
    grid_coords: Sequence[Float[ArrayLike, " _"]],
    num_modes: int | tuple[int, ...],
    *,
    boundary_scale: float = 1.25,
    variance: float | Float[ArrayLike, " M"] = 1.0,
    spectral_density: Callable[[Float[Array, " M"]], Float[Array, " M"]] | None = None,
) -> LinearBasis:
    r"""Laplacian-eigenfunction (Dirichlet/HSGP) reduced basis on a box.

    Tensor-product sine eigenfunctions of $-\Delta$ on the (recentred,
    slightly enlarged) grid bounding box. Passing ``spectral_density``
    evaluates it at the Laplacian eigenvalues $\lambda_j$ and scales it
    by ``variance`` — supply a kernel's spectral density to obtain the
    Hilbert-space GP (HSGP) approximation of that kernel's prior with
    amplitude ``variance``.

    Args:
        grid_coords: Per-axis 1-D coordinate arrays; any dimension
            ``d >= 1``. Each axis must span a non-degenerate interval.
        num_modes: Number of 1-D modes per axis; an ``int`` is broadcast
            to all axes. Total ``nbasis`` is the product.
        boundary_scale: Domain half-widths are the grid half-spans
            times this factor (``>= 1``). The Dirichlet eigenfunctions
            vanish on the box boundary, so a value ``> 1`` keeps the
            basis expressive at the grid edges.
        variance: Scalar or per-coefficient prior variance;
            with ``spectral_density`` it acts as a scalar amplitude.
        spectral_density: Optional callable mapping the eigenvalues
            $\lambda_j$ to per-mode prior power (HSGP prior).

    Algorithm:
        1. Flatten the grid, recentre it, and enlarge the half-spans by
           ``boundary_scale``.
        2. Build the tensor-product Dirichlet eigenfunctions $\Phi$ and
           eigenvalues $\lambda$ with ``num_modes`` modes per axis.
        3. Set the prior variance to ``variance`` or, with a spectral
           density, to ``variance * spectral_density(lambda)``.

    Returns:
        `LinearBasis` with ``output_shape`` the grid shape.

    Raises:
        ConfigurationError: If ``boundary_scale < 1``.
        GridError: If an axis is empty, not 1-D, or has zero span.

    Examples:
        Four sine modes on a 1-D grid:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> x = jnp.linspace(0.0, 1.0, 6)
        >>> basis = vdx.fourier_basis([x], 4)
        >>> basis.nbasis, basis.output_shape
        (4, (6,))

        A spectral density turns the prior into an HSGP prior with decaying
        variance:

        >>> hsgp = vdx.fourier_basis(
        ...     [x], 4, spectral_density=lambda lam: 1.0 / (1.0 + lam)
        ... )
        >>> bool(hsgp.variance[0] > hsgp.variance[-1])
        True

        A shrinking boundary scale is rejected:

        >>> vdx.fourier_basis(
        ...     [x], 4, boundary_scale=0.5
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: boundary_scale must be >= 1, got 0.5.
    """
    if boundary_scale < 1.0:
        raise ConfigurationError(f"boundary_scale must be >= 1, got {boundary_scale}.")
    points, grid_shape = _grid_points(grid_coords)
    lower, upper = points.min(axis=0), points.max(axis=0)
    half_span = (upper - lower) / 2.0
    if _np.any(half_span <= 0):
        raise GridError(
            "each grid_coords axis must span a non-degenerate interval "
            f"(got half-spans {tuple(half_span.tolist())})."
        )
    centered = points - (lower + upper) / 2.0
    ndim = points.shape[1]
    half_widths = tuple(float(span) * boundary_scale for span in half_span)
    modes = (
        (int(num_modes),) * ndim
        if isinstance(num_modes, int)
        else tuple(int(axis_modes) for axis_modes in num_modes)
    )
    # Φ: Dirichlet eigenfunctions of −Δ on the box, eigenvalues λ
    design_matrix, eigvals = _gnx.fourier_basis(
        jnp.asarray(centered), modes, half_widths
    )
    # Q = variance · S(λ)
    prior_variance = (
        variance * spectral_density(eigvals)
        if spectral_density is not None
        else variance
    )
    return linear_basis(design_matrix, variance=prior_variance, output_shape=grid_shape)

rbf_basis

rbf_basis(
    grid_coords: Sequence[Float[ArrayLike, " _"]],
    centers: Float[ArrayLike, "M d"],
    *,
    widths: float | Float[ArrayLike, " M"],
    kernel: Literal[
        "gaussian", "wendland_c2", "wendland_c4"
    ] = "gaussian",
    variance: float | Float[ArrayLike, " M"] = 1.0,
) -> LinearBasis

Placeable radial-basis reduced basis (Gaussian or Wendland kernels).

Column \(a\) of \(\Phi\) is \(\varphi_a(g) = K(\lVert g - c_a \rVert / \ell_a)\) evaluated on the flattened grid. kernel="gaussian" reproduces the MASSH stationary Gaussian-RBF basis; the compactly supported "wendland_c2" / "wendland_c4" kernels give columns that are exactly zero past their width.

Parameters:

Name	Type	Description	Default
grid_coords	Sequence[Float[ArrayLike, ' _']]	Per-axis 1-D coordinate arrays, e.g. (y_axis, x_axis) for a 2-D grid. Any dimension d >= 1.	required
centers	Float[ArrayLike, 'M d']	Atom centre locations \(c_a\), one row per atom with \(d\) columns.	required
widths	float | Float[ArrayLike, ' M']	Scalar or per-atom width \(\ell_a\) (Gaussian length scale, or Wendland support radius); must be positive.	required
kernel	Literal['gaussian', 'wendland_c2', 'wendland_c4']	"gaussian", "wendland_c2", or "wendland_c4".	'gaussian'
variance	float | Float[ArrayLike, ' M']	Scalar or per-coefficient prior variance.	1.0

Algorithm

1. Flatten the grid axes into points \(g\).
2. Evaluate \(K(\lVert g - c_a \rVert / \ell_a)\) for every atom.
3. Wrap the design matrix with the diagonal prior variance.

Returns:

Type	Description
LinearBasis	LinearBasis with output_shape the grid shape and
LinearBasis	nbasis == M.

Raises:

Type	Description
GridError	If an axis of grid_coords is empty or not 1-D.
ShapeError	If centers does not match the grid dimension.
ConfigurationError	If a width is not strictly positive.

Examples:

Three Gaussian bumps on a 1-D grid of five points:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> x = jnp.linspace(0.0, 1.0, 5)
>>> basis = vdx.rbf_basis([x], jnp.array([[0.0], [0.5], [1.0]]), widths=0.3)
>>> basis.nbasis, basis.output_shape
(3, (5,))

A 2-D grid gives increments on the grid shape:

>>> y = jnp.linspace(0.0, 1.0, 3)
>>> b2 = vdx.rbf_basis([y, x], jnp.array([[0.5, 0.5]]), widths=0.4)
>>> b2.operg(0.0, jnp.ones(1)).shape
(3, 5)

A non-positive width is rejected:

>>> vdx.rbf_basis(
...     [x], jnp.zeros((1, 1)), widths=0.0
... )
Traceback (most recent call last):
...
ConfigurationError: widths must be strictly positive; got min=0.0.

Source code in src/vardax/_src/basis/constructors.py

def rbf_basis(
    grid_coords: Sequence[Float[ArrayLike, " _"]],
    centers: Float[ArrayLike, "M d"],
    *,
    widths: float | Float[ArrayLike, " M"],
    kernel: Literal["gaussian", "wendland_c2", "wendland_c4"] = "gaussian",
    variance: float | Float[ArrayLike, " M"] = 1.0,
) -> LinearBasis:
    r"""Placeable radial-basis reduced basis (Gaussian or Wendland kernels).

    Column $a$ of $\Phi$ is
    $\varphi_a(g) = K(\lVert g - c_a \rVert / \ell_a)$
    evaluated on the flattened grid. ``kernel="gaussian"`` reproduces the
    MASSH stationary Gaussian-RBF basis; the compactly supported
    ``"wendland_c2"`` / ``"wendland_c4"`` kernels give columns
    that are exactly zero past their width.

    Args:
        grid_coords: Per-axis 1-D coordinate arrays, e.g. ``(y_axis,
            x_axis)`` for a 2-D grid. Any dimension ``d >= 1``.
        centers: Atom centre locations $c_a$, one row per atom with $d$
            columns.
        widths: Scalar or per-atom width $\ell_a$ (Gaussian
            length scale, or Wendland support radius); must be positive.
        kernel: ``"gaussian"``, ``"wendland_c2"``, or ``"wendland_c4"``.
        variance: Scalar or per-coefficient prior variance.

    Algorithm:
        1. Flatten the grid axes into points $g$.
        2. Evaluate $K(\lVert g - c_a \rVert / \ell_a)$ for every atom.
        3. Wrap the design matrix with the diagonal prior ``variance``.

    Returns:
        `LinearBasis` with ``output_shape`` the grid shape and
        ``nbasis == M``.

    Raises:
        GridError: If an axis of ``grid_coords`` is empty or not 1-D.
        ShapeError: If ``centers`` does not match the grid dimension.
        ConfigurationError: If a width is not strictly positive.

    Examples:
        Three Gaussian bumps on a 1-D grid of five points:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> x = jnp.linspace(0.0, 1.0, 5)
        >>> basis = vdx.rbf_basis([x], jnp.array([[0.0], [0.5], [1.0]]), widths=0.3)
        >>> basis.nbasis, basis.output_shape
        (3, (5,))

        A 2-D grid gives increments on the grid shape:

        >>> y = jnp.linspace(0.0, 1.0, 3)
        >>> b2 = vdx.rbf_basis([y, x], jnp.array([[0.5, 0.5]]), widths=0.4)
        >>> b2.operg(0.0, jnp.ones(1)).shape
        (3, 5)

        A non-positive width is rejected:

        >>> vdx.rbf_basis(
        ...     [x], jnp.zeros((1, 1)), widths=0.0
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: widths must be strictly positive; got min=0.0.
    """
    points, grid_shape = _grid_points(grid_coords)
    ndim = points.shape[1]
    centers_array = _np.asarray(centers, dtype=float)
    if centers_array.ndim != 2 or centers_array.shape[1] != ndim:
        raise ShapeError(
            f"centers must be (M, {ndim}) to match {ndim}-D grid_coords; "
            f"got {centers_array.shape}."
        )
    n_atoms = centers_array.shape[0]
    widths_array = _np.broadcast_to(_np.asarray(widths, dtype=float), (n_atoms,))
    if _np.any(widths_array <= 0):
        raise ConfigurationError(
            f"widths must be strictly positive; got min={float(widths_array.min())}."
        )
    # Φ[:, a] = K(‖g − c_a‖ / ℓ_a)
    design_matrix = _gnx.rbf_basis(
        jnp.asarray(points),
        jnp.asarray(centers_array),
        jnp.asarray(widths_array),
        kernel=kernel,
    )
    return linear_basis(design_matrix, variance=variance, output_shape=grid_shape)

wavelet_basis

wavelet_basis(
    grid_shape: int | tuple[int, ...],
    *,
    wavelet: str = "haar",
    levels: int | None = None,
    variance: float | Float[ArrayLike, " M"] = 1.0,
) -> LinearBasis

Orthonormal wavelet (DWT) reduced basis on a periodic grid.

Haar / Daubechies synthesis basis: complete (\(M = N\)) and orthonormal (\(\Phi^\top \Phi = I\)), so it is a multiscale re-parameterisation of the grid — per-coefficient variance is the natural place to encode scale-dependent prior power.

Parameters:

Name	Type	Description	Default
grid_shape	int | tuple[int, ...]	n or (n,) for a 1-D grid, (ny, nx) for 2-D. Every extent must be a power of two >= 2.	required
wavelet	str	"haar", "db2", or "db4".	'haar'
levels	int | None	Decomposition levels; defaults to the full cascade.	None
variance	float | Float[ArrayLike, ' M']	Scalar or per-coefficient prior variance.	1.0

Algorithm

1. Build the 1-D or 2-D synthesis matrix \(\Phi\) for the chosen wavelet and number of levels.
2. Wrap it with the diagonal prior variance.

Returns:

Type	Description
LinearBasis	LinearBasis with output_shape == grid_shape and
LinearBasis	nbasis == prod(grid_shape).

Raises:

Type	Description
UnsupportedInputError	If grid_shape is neither 1-D nor 2-D.

Examples:

A Haar basis on 8 points is complete and orthonormal:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> basis = vdx.wavelet_basis(8)
>>> basis.nbasis
8
>>> bool(
...     jnp.allclose(
...         basis.design_matrix.T @ basis.design_matrix, jnp.eye(8), atol=1e-5
...     )
... )
True

A 2-D grid flattens to ny * nx coefficients:

>>> vdx.wavelet_basis((4, 4)).nbasis
16

A 3-D grid is rejected:

>>> vdx.wavelet_basis((2, 2, 2))
Traceback (most recent call last):
...
UnsupportedInputError: wavelet_basis supports 1-D or 2-D grids; ...

Source code in src/vardax/_src/basis/constructors.py

def wavelet_basis(
    grid_shape: int | tuple[int, ...],
    *,
    wavelet: str = "haar",
    levels: int | None = None,
    variance: float | Float[ArrayLike, " M"] = 1.0,
) -> LinearBasis:
    r"""Orthonormal wavelet (DWT) reduced basis on a periodic grid.

    Haar / Daubechies synthesis basis: complete ($M = N$) and
    orthonormal ($\Phi^\top \Phi = I$), so it is a multiscale
    *re-parameterisation* of the grid — per-coefficient ``variance``
    is the natural place to encode scale-dependent prior power.

    Args:
        grid_shape: ``n`` or ``(n,)`` for a 1-D grid, ``(ny, nx)`` for
            2-D. Every extent must be a power of two ``>= 2``.
        wavelet: ``"haar"``, ``"db2"``, or ``"db4"``.
        levels: Decomposition levels; defaults to the full cascade.
        variance: Scalar or per-coefficient prior variance.

    Algorithm:
        1. Build the 1-D or 2-D synthesis matrix $\Phi$ for the chosen
           wavelet and number of levels.
        2. Wrap it with the diagonal prior ``variance``.

    Returns:
        `LinearBasis` with ``output_shape == grid_shape`` and
        ``nbasis == prod(grid_shape)``.

    Raises:
        UnsupportedInputError: If ``grid_shape`` is neither 1-D nor 2-D.

    Examples:
        A Haar basis on 8 points is complete and orthonormal:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> basis = vdx.wavelet_basis(8)
        >>> basis.nbasis
        8
        >>> bool(
        ...     jnp.allclose(
        ...         basis.design_matrix.T @ basis.design_matrix, jnp.eye(8), atol=1e-5
        ...     )
        ... )
        True

        A 2-D grid flattens to ``ny * nx`` coefficients:

        >>> vdx.wavelet_basis((4, 4)).nbasis
        16

        A 3-D grid is rejected:

        >>> vdx.wavelet_basis((2, 2, 2))  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        UnsupportedInputError: wavelet_basis supports 1-D or 2-D grids; ...
    """
    shape = (grid_shape,) if isinstance(grid_shape, int) else tuple(grid_shape)
    if len(shape) == 1:
        design_matrix = _gnx.wavelet_basis_1d(
            int(shape[0]), wavelet=wavelet, levels=levels
        )
    elif len(shape) == 2:
        design_matrix = _gnx.wavelet_basis_2d(
            int(shape[0]), int(shape[1]), wavelet=wavelet, levels=levels
        )
    else:
        raise UnsupportedInputError(
            f"wavelet_basis supports 1-D or 2-D grids; got {len(shape)}-D "
            f"grid_shape {shape}."
        )
    return linear_basis(design_matrix, variance=variance, output_shape=shape)

State parameterizations

Implementations of the StateParameterization protocol, accepted as parameterization= by ThreeDVar, StrongFourDVar, WeakFourDVar and ReducedBasisFourDVar: the control becomes any pytree \(\theta\), the state \(\psi(\theta)\), and \(\mathrm{pen}(\theta)\) replaces the background term. LinearParameterization adapts any reduced basis (\(\psi(X) = x_b + \Phi X\)); NeuralFieldParameterization makes the state a coordinate network on the grid, with weight decay as its penalty.

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

LinearParameterization

Bases: Module

Adapter from a ReducedBasis to a parameterization.

\(\psi(X) = x_b + \Phi(t) X\) (with \(t = 0\) when not given) and \(\mathrm{pen}(X) = \tfrac12 X^\top Q^{-1} X\), so a model built with parameterization=LinearParameterization(basis, background) minimises the same cost as ReducedBasisFourDVar(basis=basis) with prior_mean=x_b.

Algorithm

1. decode applies the basis to the coefficients and adds \(x_b\).
2. penalty applies \(Q^{-1}\) through the basis.
3. init_control returns the zero coefficient vector.

Attributes:

Name	Type	Description
basis	Any	The ReducedBasis (\(\Phi\), \(Q^{-1}\), nbasis).
background	Float[Array, ...]	Background state the increment is added to.

Examples:

>>> import jax.numpy as jnp, vardax
>>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=2.0)
>>> param = vardax.LinearParameterization(basis, background=jnp.ones(4))
>>> X = param.init_control()
>>> X.shape
(2,)
>>> param.decode(jnp.array([1.0, -1.0]))
Array([2., 0., 1., 1.], dtype=float32)
>>> float(param.penalty(jnp.array([1.0, -1.0])))  # ½ (1 + 1) / 2
0.5

The initial control decodes to the background, whatever the basis.

>>> bool(jnp.allclose(param.decode(X), param.background))
True

It satisfies the StateParameterization protocol.

>>> isinstance(param, vardax.StateParameterization)
True

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

class LinearParameterization(eqx.Module):
    r"""Adapter from a [`ReducedBasis`][vardax.ReducedBasis] to a parameterization.

    $\psi(X) = x_b + \Phi(t) X$ (with $t = 0$ when not given) and
    $\mathrm{pen}(X) = \tfrac12 X^\top Q^{-1} X$, so a model built with
    ``parameterization=LinearParameterization(basis, background)`` minimises the
    same cost as ``ReducedBasisFourDVar(basis=basis)`` with ``prior_mean=x_b``.

    Algorithm:
        1. ``decode`` applies the basis to the coefficients and adds
           $x_b$.
        2. ``penalty`` applies $Q^{-1}$ through the basis.
        3. ``init_control`` returns the zero coefficient vector.

    Attributes:
        basis: The `ReducedBasis` ($\Phi$, $Q^{-1}$, ``nbasis``).
        background: Background state the increment is added to.

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> basis = vardax.linear_basis(jnp.eye(4)[:, :2], variance=2.0)
        >>> param = vardax.LinearParameterization(basis, background=jnp.ones(4))
        >>> X = param.init_control()
        >>> X.shape
        (2,)
        >>> param.decode(jnp.array([1.0, -1.0]))
        Array([2., 0., 1., 1.], dtype=float32)
        >>> float(param.penalty(jnp.array([1.0, -1.0])))  # ½ (1 + 1) / 2
        0.5

        The initial control decodes to the background, whatever the basis.

        >>> bool(jnp.allclose(param.decode(X), param.background))
        True

        It satisfies the [`StateParameterization`][vardax.StateParameterization]
        protocol.

        >>> isinstance(param, vardax.StateParameterization)
        True
    """

    basis: Any
    background: Float[Array, ...]

    def decode(self, theta: Float[Array, " M"], t: float | None = None) -> Array:
        """Return the state $x_b + \\Phi(t) X$, shaped like ``background``.

        Args:
            theta: Coefficient vector $X$.
            t: Time at which the basis is evaluated; ``None`` means ``0``.

        Returns:
            The decoded state.
        """
        increment = self.basis.operg(0.0 if t is None else t, theta)
        # x = x_b + Φ(t) X
        return self.background + as_shape(increment, self.background.shape)

    def penalty(self, theta: Float[Array, " M"]) -> Float[Array, ""]:
        """Return $\\tfrac12 X^\\top Q^{-1} X$.

        Args:
            theta: Coefficient vector $X$.

        Returns:
            The scalar penalty.
        """
        return 0.5 * jnp.sum(theta * self.basis.prior_inv(theta))

    def init_control(self, *, key: PRNGKeyArray | None = None) -> Array:
        """Return the zero coefficient vector (the background).

        Args:
            key: Unused; the start is deterministic.

        Returns:
            The zero coefficient vector $X = 0$.
        """
        del key
        return jnp.zeros((self.basis.nbasis,), dtype=self.background.dtype)

decode

decode(
    theta: Float[Array, " M"], t: float | None = None
) -> Array

Return the state \(x_b + \Phi(t) X\), shaped like background.

Parameters:

Name	Type	Description	Default
theta	Float[Array, ' M']	Coefficient vector \(X\).	required
t	float | None	Time at which the basis is evaluated; None means 0.	None

Returns:

Type	Description
Array	The decoded state.

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

def decode(self, theta: Float[Array, " M"], t: float | None = None) -> Array:
    """Return the state $x_b + \\Phi(t) X$, shaped like ``background``.

    Args:
        theta: Coefficient vector $X$.
        t: Time at which the basis is evaluated; ``None`` means ``0``.

    Returns:
        The decoded state.
    """
    increment = self.basis.operg(0.0 if t is None else t, theta)
    # x = x_b + Φ(t) X
    return self.background + as_shape(increment, self.background.shape)

penalty

penalty(theta: Float[Array, ' M']) -> Float[Array, '']

Return \(\tfrac12 X^\top Q^{-1} X\).

Parameters:

Name	Type	Description	Default
theta	Float[Array, ' M']	Coefficient vector \(X\).	required

Returns:

Type	Description
Float[Array, '']	The scalar penalty.

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

def penalty(self, theta: Float[Array, " M"]) -> Float[Array, ""]:
    """Return $\\tfrac12 X^\\top Q^{-1} X$.

    Args:
        theta: Coefficient vector $X$.

    Returns:
        The scalar penalty.
    """
    return 0.5 * jnp.sum(theta * self.basis.prior_inv(theta))

init_control

init_control(*, key: PRNGKeyArray | None = None) -> Array

Return the zero coefficient vector (the background).

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray | None	Unused; the start is deterministic.	None

Returns:

Type	Description
Array	The zero coefficient vector \(X = 0\).

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

def init_control(self, *, key: PRNGKeyArray | None = None) -> Array:
    """Return the zero coefficient vector (the background).

    Args:
        key: Unused; the start is deterministic.

    Returns:
        The zero coefficient vector $X = 0$.
    """
    del key
    return jnp.zeros((self.basis.nbasis,), dtype=self.background.dtype)

NeuralFieldParameterization

Bases: Module

A coordinate network evaluated on the grid as the state.

The control is the array half of the field, $\theta = $ eqx.filter(field, eqx.is_inexact_array); the rest of the module (activations, sizes) is fixed. The state is the field at the grid coordinates, \(\psi(\theta) = [f_\theta(s_j)]_{j=1}^{N}\), of shape (N,) for a scalar field and (N, C) for a C-channel one. With a time t the field is evaluated at \((s_j, t)\), so a space-time field takes D + 1 inputs. The penalty is weight decay, \(\mathrm{pen}(\theta) = \tfrac{\alpha}{2} \sum_i \theta_i^2\).

Algorithm

1. init_control extracts the array leaves of field selected by filter_spec as \(\theta\).
2. decode recombines \(\theta\) with the fixed part of field and evaluates it at every grid coordinate.
3. penalty returns the weight decay of \(\theta\).

Raises:

Type	Description
ShapeError	If coords is not two-dimensional.
ConfigurationError	If weight_decay is negative.

Attributes:

Name	Type	Description
field	Module	Single-example network field(s) -> (C,), for example a geonnax.SIREN. Its current parameters are the initial control.
coords	Float[Array, 'N D']	Grid coordinates \(s_j\), shape (N, D).
weight_decay	float	\(\alpha \ge 0\).
filter_spec	Any	Which leaves of field are the control: an equinox filter (a callable on leaves, or a pytree of bools shaped like field), used for both halves of the partition. Default eqx.is_inexact_array: every floating array.

Frozen arrays: geonnax (and equinox) keep fixed arrays — random Fourier frequencies, precomputed tables — as ordinary leaves read through jax.lax.stop_gradient. The default filter would make them controls, and the weight decay would shrink them to zero. Exclude them with filter_spec, e.g. eqx.tree_at(lambda f: f.B, jax.tree.map(eqx.is_inexact_array, field), False); they then stay fixed and outside the penalty.

Examples:

>>> import jax, jax.numpy as jnp, geonnax, vardax
>>> angle = 2 * jnp.pi * jnp.arange(40) / 40
>>> coords = jnp.stack([jnp.cos(angle), jnp.sin(angle)], axis=-1)
>>> siren = geonnax.SIREN.init(2, 16, 1, depth=3, key=jax.random.PRNGKey(0))
>>> param = vardax.NeuralFieldParameterization(siren, coords, 1e-4)
>>> theta = param.init_control()
>>> param.decode(theta).shape
(40,)
>>> bool(param.penalty(theta) > 0)
True

With a time, the field takes D + 1 inputs.

>>> st = geonnax.SIREN.init(3, 16, 1, depth=3, key=jax.random.PRNGKey(1))
>>> st_param = vardax.NeuralFieldParameterization(st, coords, 0.0)
>>> st_param.decode(st_param.init_control(), t=0.5).shape
(40,)

Negative weight decay is rejected.

>>> vardax.NeuralFieldParameterization(
...     siren, coords, -1.0
... )
Traceback (most recent call last):
...
ConfigurationError: weight_decay must be >= 0, got -1.0.

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

class NeuralFieldParameterization(eqx.Module):
    r"""A coordinate network evaluated on the grid as the state.

    The control is the array half of the field,
    $\theta = $ ``eqx.filter(field, eqx.is_inexact_array)``; the rest of
    the module (activations, sizes) is fixed. The state is the field at
    the grid coordinates,
    $\psi(\theta) = [f_\theta(s_j)]_{j=1}^{N}$, of shape ``(N,)`` for a
    scalar field and ``(N, C)`` for a ``C``-channel one. With a time ``t``
    the field is evaluated at $(s_j, t)$, so a space-time field takes
    ``D + 1`` inputs. The penalty is weight decay,
    $\mathrm{pen}(\theta) = \tfrac{\alpha}{2} \sum_i \theta_i^2$.

    Algorithm:
        1. ``init_control`` extracts the array leaves of ``field``
           selected by ``filter_spec`` as $\theta$.
        2. ``decode`` recombines $\theta$ with the fixed part of ``field``
           and evaluates it at every grid coordinate.
        3. ``penalty`` returns the weight decay of $\theta$.

    Raises:
        ShapeError: If ``coords`` is not two-dimensional.
        ConfigurationError: If ``weight_decay`` is negative.

    Attributes:
        field: Single-example network ``field(s) -> (C,)``, for example a
            `geonnax.SIREN`. Its current parameters are the initial control.
        coords: Grid coordinates $s_j$, shape ``(N, D)``.
        weight_decay: $\alpha \ge 0$.
        filter_spec: Which leaves of ``field`` are the control: an
            ``equinox`` filter (a callable on leaves, or a pytree of bools
            shaped like ``field``), used for both halves of the partition.
            Default ``eqx.is_inexact_array``: every floating array.

    Frozen arrays: geonnax (and equinox) keep fixed arrays — random Fourier
    frequencies, precomputed tables — as ordinary leaves read through
    ``jax.lax.stop_gradient``. The default filter would make them controls,
    and the weight decay would shrink them to zero. Exclude them with
    ``filter_spec``, e.g.
    ``eqx.tree_at(lambda f: f.B, jax.tree.map(eqx.is_inexact_array, field), False)``;
    they then stay fixed and outside the penalty.

    Examples:
        >>> import jax, jax.numpy as jnp, geonnax, vardax
        >>> angle = 2 * jnp.pi * jnp.arange(40) / 40
        >>> coords = jnp.stack([jnp.cos(angle), jnp.sin(angle)], axis=-1)
        >>> siren = geonnax.SIREN.init(2, 16, 1, depth=3, key=jax.random.PRNGKey(0))
        >>> param = vardax.NeuralFieldParameterization(siren, coords, 1e-4)
        >>> theta = param.init_control()
        >>> param.decode(theta).shape
        (40,)
        >>> bool(param.penalty(theta) > 0)
        True

        With a time, the field takes ``D + 1`` inputs.

        >>> st = geonnax.SIREN.init(3, 16, 1, depth=3, key=jax.random.PRNGKey(1))
        >>> st_param = vardax.NeuralFieldParameterization(st, coords, 0.0)
        >>> st_param.decode(st_param.init_control(), t=0.5).shape
        (40,)

        Negative weight decay is rejected.

        >>> vardax.NeuralFieldParameterization(  # doctest: +IGNORE_EXCEPTION_DETAIL
        ...     siren, coords, -1.0
        ... )
        Traceback (most recent call last):
        ...
        ConfigurationError: weight_decay must be >= 0, got -1.0.
    """

    field: eqx.Module
    coords: Float[Array, "N D"]
    weight_decay: float = eqx.field(static=True)
    filter_spec: Any = eqx.field(static=True)

    def __init__(
        self,
        field: eqx.Module,
        coords: Float[Array, "N D"],
        weight_decay: float,
        *,
        filter_spec: Any = eqx.is_inexact_array,
    ) -> None:
        coords = jnp.asarray(coords)
        if coords.ndim != 2:
            raise ShapeError(f"coords must be 2-D (N, D); got shape {coords.shape}.")
        if weight_decay < 0:
            raise ConfigurationError(f"weight_decay must be >= 0, got {weight_decay}.")
        self.field = field
        self.coords = coords
        self.weight_decay = weight_decay
        self.filter_spec = filter_spec

    def as_field(self, theta: PyTree[Array]) -> Callable[[Array], Array]:
        """The field with parameters ``theta``, as a callable ``s -> (C,)``.

        Args:
            theta: Control, the array half of the field.

        Returns:
            The field with ``theta`` combined into the fixed part.
        """
        static = eqx.filter(self.field, self.filter_spec, inverse=True)
        return eqx.combine(theta, static)

    def decode(self, theta: PyTree[Array], t: float | None = None) -> Array:
        """Evaluate the field with parameters ``theta`` on the grid.

        Args:
            theta: Control, the array half of the field.
            t: Time appended to every coordinate; ``None`` for a
                time-independent field.

        Returns:
            The field values at the grid coordinates.
        """
        field = self.as_field(theta)
        coords = self.coords
        if t is not None:
            time = jnp.full((coords.shape[0],), t, dtype=coords.dtype)
            coords = einx.id("n d, n -> n (d + 1)", coords, time)
        values = jax.vmap(field)(coords)
        if values.ndim == 2 and values.shape[1] == 1:
            return values[:, 0]
        return values

    def penalty(self, theta: PyTree[Array]) -> Float[Array, ""]:
        """Return $\\tfrac{\\alpha}{2} \\lVert \\theta \\rVert^2$.

        Args:
            theta: Control, the array half of the field.

        Returns:
            The scalar weight-decay penalty.
        """
        leaves = jtu.tree_leaves(theta)
        sq = sum((jnp.sum(jnp.square(leaf)) for leaf in leaves), jnp.zeros(()))
        return 0.5 * self.weight_decay * sq

    def init_control(self, *, key: PRNGKeyArray | None = None) -> PyTree[Array]:
        """Return the field's current parameters.

        Args:
            key: Unused; the start is the field's own parameters.

        Returns:
            The array half of the field selected by ``filter_spec``.
        """
        del key
        return eqx.filter(self.field, self.filter_spec)

as_field

as_field(theta: PyTree[Array]) -> Callable[[Array], Array]

The field with parameters theta, as a callable s -> (C,).

Parameters:

Name	Type	Description	Default
theta	PyTree[Array]	Control, the array half of the field.	required

Returns:

Type	Description
Callable[[Array], Array]	The field with theta combined into the fixed part.

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

def as_field(self, theta: PyTree[Array]) -> Callable[[Array], Array]:
    """The field with parameters ``theta``, as a callable ``s -> (C,)``.

    Args:
        theta: Control, the array half of the field.

    Returns:
        The field with ``theta`` combined into the fixed part.
    """
    static = eqx.filter(self.field, self.filter_spec, inverse=True)
    return eqx.combine(theta, static)

decode

decode(
    theta: PyTree[Array], t: float | None = None
) -> Array

Evaluate the field with parameters theta on the grid.

Parameters:

Name	Type	Description	Default
theta	PyTree[Array]	Control, the array half of the field.	required
t	float | None	Time appended to every coordinate; None for a time-independent field.	None

Returns:

Type	Description
Array	The field values at the grid coordinates.

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

def decode(self, theta: PyTree[Array], t: float | None = None) -> Array:
    """Evaluate the field with parameters ``theta`` on the grid.

    Args:
        theta: Control, the array half of the field.
        t: Time appended to every coordinate; ``None`` for a
            time-independent field.

    Returns:
        The field values at the grid coordinates.
    """
    field = self.as_field(theta)
    coords = self.coords
    if t is not None:
        time = jnp.full((coords.shape[0],), t, dtype=coords.dtype)
        coords = einx.id("n d, n -> n (d + 1)", coords, time)
    values = jax.vmap(field)(coords)
    if values.ndim == 2 and values.shape[1] == 1:
        return values[:, 0]
    return values

penalty

penalty(theta: PyTree[Array]) -> Float[Array, '']

Return \(\tfrac{\alpha}{2} \lVert \theta \rVert^2\).

Parameters:

Name	Type	Description	Default
theta	PyTree[Array]	Control, the array half of the field.	required

Returns:

Type	Description
Float[Array, '']	The scalar weight-decay penalty.

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

def penalty(self, theta: PyTree[Array]) -> Float[Array, ""]:
    """Return $\\tfrac{\\alpha}{2} \\lVert \\theta \\rVert^2$.

    Args:
        theta: Control, the array half of the field.

    Returns:
        The scalar weight-decay penalty.
    """
    leaves = jtu.tree_leaves(theta)
    sq = sum((jnp.sum(jnp.square(leaf)) for leaf in leaves), jnp.zeros(()))
    return 0.5 * self.weight_decay * sq

init_control

init_control(
    *, key: PRNGKeyArray | None = None
) -> PyTree[Array]

Return the field's current parameters.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray | None	Unused; the start is the field's own parameters.	None

Returns:

Type	Description
PyTree[Array]	The array half of the field selected by filter_spec.

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

def init_control(self, *, key: PRNGKeyArray | None = None) -> PyTree[Array]:
    """Return the field's current parameters.

    Args:
        key: Unused; the start is the field's own parameters.

    Returns:
        The array half of the field selected by ``filter_spec``.
    """
    del key
    return eqx.filter(self.field, self.filter_spec)

Minimisers

Implementations of the Minimiser protocol, accepted as minimiser= by ThreeDVar, StrongFourDVar, WeakFourDVar and ReducedBasisFourDVar. OptimistixMinimiser wraps an optimistix solver (the models' default, BFGS); OptaxMinimiser runs a first-order optax optimiser with early stopping for controls too large for BFGS, and reports the final gradient norm the implicit adjoint relies on. Both return a MinimiseResult from run.

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

OptaxMinimiser

Bases: Module

First-order minimisation with an optax optimiser.

\[ \theta_{k+1} = \theta_k + u_k, \qquad g_k = \nabla J(\theta_k) \]

with \(u_k\) the optimizer update from \(g_k\). The loop is a jax.lax.while_loop (so the whole minimisation is jit-able) and stops at the first of:

• max_steps steps;
• a small gradient, \(\lVert g_k \rVert_2 \le \mathrm{gtol}\);
• step_patience consecutive small steps, \(\lVert u_k \rVert_2 \le \mathrm{atol} + \mathrm{rtol}\,\lVert \theta_k \rVert_2\) (one small step is not enough: a learning-rate schedule starting at zero gives a zero first update);
• patience consecutive steps whose cost exceeds the lowest cost seen by more than float round-off.

There is deliberately no relative cost-change test: with a large irreducible misfit, \(J_k - J_{k-1}\) rounds to zero in float32 long before \(\theta\) converges. For the same reason a step counts as an improvement when \(J_k \le J_\text{best} + 8\,\varepsilon\,|J_\text{best}|\) (\(\varepsilon\) the machine epsilon of \(J\)), so a cost that is flat to round-off does not exhaust the patience. At a small gradient the loop stops before applying the update, so stored momentum cannot carry a stationary iterate away. It returns the best iterate seen: the lowest cost, with ties within round-off broken by the smaller gradient norm.

With polish, a second loop of at most polish_steps steps of that optimiser (for example optax.lbfgs()) starts from there, with its own, tight, tolerances (polish_gtol, polish_rtol, polish_atol, polish_patience; by default it stops only on polish_steps or step_patience exactly zero steps). The optimisers are called with the extra arguments value, grad and value_fn, so line-search optimisers such as optax.lbfgs work unchanged.

The implicit adjoint used to differentiate through the minimum assumes \(\nabla J(\theta^\star) = 0\), which a first-order optimiser stopped early does not guarantee: check MinimiseResult.grad_norm from run.

Algorithm

1. Evaluate the cost and gradient at the iterate \(\theta_k\).
2. Stop if the gradient norm is below gtol; otherwise apply the optimizer update.
3. Track the best iterate, and count consecutive small steps and steps without improvement.
4. Repeat until a stopping rule above fires; keep the best iterate.
5. With polish, repeat steps 1-4 from the best iterate with the polish optimiser and tolerances.

Raises:

Type	Description
ConfigurationError	If a step cap or step_patience is below 1, a tolerance is negative, or a patience is neither None nor at least 1.

Attributes:

Name	Type	Description
optimizer	GradientTransformation	The optax.GradientTransformation of the main loop.
max_steps	int	Step cap of the main loop.
gtol	float	Gradient-norm tolerance of the main loop; 0 disables it.
rtol	float	Relative step tolerance of the main loop.
atol	float	Absolute step tolerance of the main loop; with rtol = atol = 0 only exactly zero steps stop the loop.
step_patience	int	Consecutive small steps needed to stop (both loops).
patience	int | None	Steps without improvement before stopping; None disables the test.
polish	GradientTransformation | None	Optional second optimiser, run after the main loop.
polish_steps	int	Step cap of the polish loop.
polish_gtol	float	As gtol, for the polish loop.
polish_rtol	float	As rtol, for the polish loop.
polish_atol	float	As atol, for the polish loop.
polish_patience	int | None	As patience, for the polish loop.

Examples:

Adam on \(J(x) = \tfrac12 \lVert x - m \rVert^2\), polished by L-BFGS, reaches the closed-form minimum \(x^\star = m\).

>>> import jax.numpy as jnp, optax, vardax
>>> m = jnp.array([1.0, -2.0, 0.5])
>>> minimiser = vardax.OptaxMinimiser(
...     optimizer=optax.adam(1e-1), max_steps=500, polish=optax.lbfgs()
... )
>>> result = minimiser.run(
...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(3), m
... )
>>> bool(jnp.allclose(result.value, m, atol=1e-5))
True
>>> bool(result.grad_norm < 1e-5)
True

Calling the minimiser returns only the minimiser, and any pytree is a valid control.

>>> adam = vardax.OptaxMinimiser(optax.adam(1e-1), max_steps=300)
>>> x = adam(
...     lambda p, m: 0.5 * jnp.sum((p["a"] - m) ** 2), {"a": jnp.zeros(3)}, m
... )
>>> bool(jnp.allclose(x["a"], m, atol=1e-2))
True

An invalid tolerance is rejected.

>>> vardax.OptaxMinimiser(gtol=-1.0)
Traceback (most recent call last):
...
ConfigurationError: gtol must be >= 0, got -1.0.

References

Kingma, D. P. and Ba, J. (2015). Adam: a method for stochastic optimization. ICLR.

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

class OptaxMinimiser(eqx.Module):
    r"""First-order minimisation with an ``optax`` optimiser.

    $$ \theta_{k+1} = \theta_k + u_k, \qquad g_k = \nabla J(\theta_k) $$

    with $u_k$ the ``optimizer`` update from $g_k$. The loop is a
    ``jax.lax.while_loop`` (so the whole minimisation is ``jit``-able) and
    stops at the first of:

    - ``max_steps`` steps;
    - a small gradient, $\lVert g_k \rVert_2 \le \mathrm{gtol}$;
    - ``step_patience`` consecutive small steps,
      $\lVert u_k \rVert_2 \le \mathrm{atol} + \mathrm{rtol}\,\lVert \theta_k \rVert_2$
      (one small step is not enough: a learning-rate schedule starting at
      zero gives a zero first update);
    - ``patience`` consecutive steps whose cost exceeds the lowest cost
      seen by more than float round-off.

    There is deliberately no relative cost-change test: with a large
    irreducible misfit, $J_k - J_{k-1}$ rounds to zero in float32 long
    before $\theta$ converges. For the same reason a step counts as an
    improvement when $J_k \le J_\text{best} + 8\,\varepsilon\,|J_\text{best}|$
    ($\varepsilon$ the machine epsilon of $J$), so a cost that is flat to
    round-off does not exhaust the patience. At a small gradient the loop
    stops *before* applying the update, so stored momentum cannot carry a
    stationary iterate away. It returns the best iterate seen: the lowest
    cost, with ties within round-off broken by the smaller gradient norm.

    With ``polish``, a second loop of at most ``polish_steps`` steps of that
    optimiser (for example ``optax.lbfgs()``) starts from there, with its
    own, tight, tolerances (``polish_gtol``, ``polish_rtol``,
    ``polish_atol``, ``polish_patience``; by default it stops only on
    ``polish_steps`` or ``step_patience`` exactly zero steps). The optimisers are called
    with the extra arguments ``value``, ``grad`` and ``value_fn``, so
    line-search optimisers such as ``optax.lbfgs`` work unchanged.

    The implicit adjoint used to differentiate through the minimum assumes
    $\nabla J(\theta^\star) = 0$, which a first-order optimiser stopped
    early does not guarantee: check
    [`MinimiseResult.grad_norm`][vardax.MinimiseResult] from `run`.

    Algorithm:
        1. Evaluate the cost and gradient at the iterate $\theta_k$.
        2. Stop if the gradient norm is below ``gtol``; otherwise apply
           the ``optimizer`` update.
        3. Track the best iterate, and count consecutive small steps and
           steps without improvement.
        4. Repeat until a stopping rule above fires; keep the best
           iterate.
        5. With ``polish``, repeat steps 1-4 from the best iterate with
           the polish optimiser and tolerances.

    Raises:
        ConfigurationError: If a step cap or ``step_patience`` is below 1,
            a tolerance is negative, or a patience is neither ``None``
            nor at least 1.

    Attributes:
        optimizer: The ``optax.GradientTransformation`` of the main loop.
        max_steps: Step cap of the main loop.
        gtol: Gradient-norm tolerance of the main loop; ``0`` disables it.
        rtol: Relative step tolerance of the main loop.
        atol: Absolute step tolerance of the main loop; with ``rtol = atol
            = 0`` only exactly zero steps stop the loop.
        step_patience: Consecutive small steps needed to stop (both loops).
        patience: Steps without improvement before stopping; ``None``
            disables the test.
        polish: Optional second optimiser, run after the main loop.
        polish_steps: Step cap of the polish loop.
        polish_gtol: As ``gtol``, for the polish loop.
        polish_rtol: As ``rtol``, for the polish loop.
        polish_atol: As ``atol``, for the polish loop.
        polish_patience: As ``patience``, for the polish loop.

    Examples:
        Adam on $J(x) = \tfrac12 \lVert x - m \rVert^2$, polished by
        L-BFGS, reaches the closed-form minimum $x^\star = m$.

        >>> import jax.numpy as jnp, optax, vardax
        >>> m = jnp.array([1.0, -2.0, 0.5])
        >>> minimiser = vardax.OptaxMinimiser(
        ...     optimizer=optax.adam(1e-1), max_steps=500, polish=optax.lbfgs()
        ... )
        >>> result = minimiser.run(
        ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(3), m
        ... )
        >>> bool(jnp.allclose(result.value, m, atol=1e-5))
        True
        >>> bool(result.grad_norm < 1e-5)
        True

        Calling the minimiser returns only the minimiser, and any pytree
        is a valid control.

        >>> adam = vardax.OptaxMinimiser(optax.adam(1e-1), max_steps=300)
        >>> x = adam(
        ...     lambda p, m: 0.5 * jnp.sum((p["a"] - m) ** 2), {"a": jnp.zeros(3)}, m
        ... )
        >>> bool(jnp.allclose(x["a"], m, atol=1e-2))
        True

        An invalid tolerance is rejected.

        >>> vardax.OptaxMinimiser(gtol=-1.0)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: gtol must be >= 0, got -1.0.

    References:
        Kingma, D. P. and Ba, J. (2015). Adam: a method for stochastic
        optimization. ICLR.
    """

    optimizer: optax.GradientTransformation = eqx.field(static=True)
    max_steps: int = eqx.field(static=True)
    gtol: float = eqx.field(static=True)
    rtol: float = eqx.field(static=True)
    atol: float = eqx.field(static=True)
    step_patience: int = eqx.field(static=True)
    patience: int | None = eqx.field(static=True)
    polish: optax.GradientTransformation | None = eqx.field(static=True)
    polish_steps: int = eqx.field(static=True)
    polish_gtol: float = eqx.field(static=True)
    polish_rtol: float = eqx.field(static=True)
    polish_atol: float = eqx.field(static=True)
    polish_patience: int | None = eqx.field(static=True)

    def __init__(
        self,
        optimizer: optax.GradientTransformation | None = None,
        *,
        max_steps: int = 5000,
        gtol: float = 1e-6,
        rtol: float = 0.0,
        atol: float = 0.0,
        step_patience: int = 3,
        patience: int | None = 200,
        polish: optax.GradientTransformation | None = None,
        polish_steps: int = 200,
        polish_gtol: float = 0.0,
        polish_rtol: float = 0.0,
        polish_atol: float = 0.0,
        polish_patience: int | None = None,
    ) -> None:
        for name, steps in (
            ("max_steps", max_steps),
            ("polish_steps", polish_steps),
            ("step_patience", step_patience),
        ):
            if steps < 1:
                raise ConfigurationError(f"{name} must be >= 1, got {steps}.")
        for name, value in (
            ("gtol", gtol),
            ("rtol", rtol),
            ("atol", atol),
            ("polish_gtol", polish_gtol),
            ("polish_rtol", polish_rtol),
            ("polish_atol", polish_atol),
        ):
            if value < 0:
                raise ConfigurationError(f"{name} must be >= 0, got {value}.")
        for name, value in (
            ("patience", patience),
            ("polish_patience", polish_patience),
        ):
            if value is not None and value < 1:
                raise ConfigurationError(f"{name} must be >= 1 or None, got {value}.")
        self.optimizer = optimizer if optimizer is not None else optax.adam(1e-3)
        self.max_steps = max_steps
        self.gtol = gtol
        self.rtol = rtol
        self.atol = atol
        self.step_patience = step_patience
        self.patience = patience
        self.polish = polish
        self.polish_steps = polish_steps
        self.polish_gtol = polish_gtol
        self.polish_rtol = polish_rtol
        self.polish_atol = polish_atol
        self.polish_patience = polish_patience

    def run(
        self,
        cost_fn: Callable[[Any, Any], Float[Array, ""]],
        x0: PyTree[Array],
        batch: Any = None,
        *,
        adjoint: optx.AbstractAdjoint | None = None,
    ) -> MinimiseResult:
        """Minimise ``cost_fn(x, batch)`` from ``x0``, with diagnostics.

        ``result.cost`` is differentiable: by the envelope theorem
        $\\mathrm{d}J^\\star/\\mathrm{d}\\phi = \\partial J/\\partial\\phi$ at
        $\\theta^\\star$ for any parameter $\\phi$ of the cost.

        Args:
            cost_fn: Scalar cost ``cost_fn(x, batch)``.
            x0: Initial guess, any pytree of arrays (for a network, the
                array half of ``eqx.partition(model, eqx.is_inexact_array)``).
            batch: Passed as the second argument of ``cost_fn``.
            adjoint: How to differentiate through the minimum. Default
                ``optimistix.ImplicitAdjoint()``; the loop is then a
                ``jax.lax.while_loop``. ``optimistix.RecursiveCheckpointAdjoint``
                backpropagates through the iterates instead; it cannot
                differentiate an optimiser with an inner ``while_loop``
                (the line search of ``optax.lbfgs``).

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

        Examples:
            >>> import jax.numpy as jnp, optax, vardax
            >>> minimiser = vardax.OptaxMinimiser(optax.adam(1e-1), max_steps=200)
            >>> result = minimiser.run(
            ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
            ... )
            >>> result.num_steps > 0
            Array(True, dtype=bool)
        """
        x0 = jtu.tree_map(jnp.asarray, x0)
        closed_cost_fn = _close_over(cost_fn, x0, batch)
        inputs = (closed_cost_fn, self, x0, batch)
        value, (num_steps, converged) = (adjoint or optx.ImplicitAdjoint()).apply(
            _optax_primal, _optax_rewrite, inputs, frozenset()
        )
        cost, grad_norm = _cost_and_grad_norm(cost_fn, value, batch)
        return MinimiseResult(
            value=value,
            cost=cost,
            grad_norm=grad_norm,
            num_steps=num_steps,
            converged=converged,
        )

    def __call__(
        self,
        cost_fn: Callable[[Any, Any], Float[Array, ""]],
        x0: PyTree[Array],
        batch: Any = None,
    ) -> PyTree[Array]:
        """Return ``argmin_x cost_fn(x, batch)`` starting from ``x0``.

        Args:
            cost_fn: Scalar cost ``cost_fn(x, batch)``.
            x0: Initial guess, any pytree of arrays.
            batch: Passed as the second argument of ``cost_fn``.

        Returns:
            The minimiser, with the structure of ``x0``.
        """
        return self.run(cost_fn, x0, batch).value

run

run(
    cost_fn: Callable[[Any, Any], Float[Array, ""]],
    x0: PyTree[Array],
    batch: Any = None,
    *,
    adjoint: AbstractAdjoint | None = None,
) -> MinimiseResult

Minimise cost_fn(x, batch) from x0, with diagnostics.

result.cost is differentiable: by the envelope theorem \(\mathrm{d}J^\star/\mathrm{d}\phi = \partial J/\partial\phi\) at \(\theta^\star\) for any parameter \(\phi\) of the cost.

Parameters:

Name	Type	Description	Default
cost_fn	Callable[[Any, Any], Float[Array, '']]	Scalar cost cost_fn(x, batch).	required
x0	PyTree[Array]	Initial guess, any pytree of arrays (for a network, the array half of eqx.partition(model, eqx.is_inexact_array)).	required
batch	Any	Passed as the second argument of cost_fn.	None
adjoint	AbstractAdjoint | None	How to differentiate through the minimum. Default optimistix.ImplicitAdjoint(); the loop is then a jax.lax.while_loop. optimistix.RecursiveCheckpointAdjoint backpropagates through the iterates instead; it cannot differentiate an optimiser with an inner while_loop (the line search of optax.lbfgs).	None

Returns:

Type	Description
MinimiseResult	A MinimiseResult.

Examples:

>>> import jax.numpy as jnp, optax, vardax
>>> minimiser = vardax.OptaxMinimiser(optax.adam(1e-1), max_steps=200)
>>> result = minimiser.run(
...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
... )
>>> result.num_steps > 0
Array(True, dtype=bool)

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

def run(
    self,
    cost_fn: Callable[[Any, Any], Float[Array, ""]],
    x0: PyTree[Array],
    batch: Any = None,
    *,
    adjoint: optx.AbstractAdjoint | None = None,
) -> MinimiseResult:
    """Minimise ``cost_fn(x, batch)`` from ``x0``, with diagnostics.

    ``result.cost`` is differentiable: by the envelope theorem
    $\\mathrm{d}J^\\star/\\mathrm{d}\\phi = \\partial J/\\partial\\phi$ at
    $\\theta^\\star$ for any parameter $\\phi$ of the cost.

    Args:
        cost_fn: Scalar cost ``cost_fn(x, batch)``.
        x0: Initial guess, any pytree of arrays (for a network, the
            array half of ``eqx.partition(model, eqx.is_inexact_array)``).
        batch: Passed as the second argument of ``cost_fn``.
        adjoint: How to differentiate through the minimum. Default
            ``optimistix.ImplicitAdjoint()``; the loop is then a
            ``jax.lax.while_loop``. ``optimistix.RecursiveCheckpointAdjoint``
            backpropagates through the iterates instead; it cannot
            differentiate an optimiser with an inner ``while_loop``
            (the line search of ``optax.lbfgs``).

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

    Examples:
        >>> import jax.numpy as jnp, optax, vardax
        >>> minimiser = vardax.OptaxMinimiser(optax.adam(1e-1), max_steps=200)
        >>> result = minimiser.run(
        ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
        ... )
        >>> result.num_steps > 0
        Array(True, dtype=bool)
    """
    x0 = jtu.tree_map(jnp.asarray, x0)
    closed_cost_fn = _close_over(cost_fn, x0, batch)
    inputs = (closed_cost_fn, self, x0, batch)
    value, (num_steps, converged) = (adjoint or optx.ImplicitAdjoint()).apply(
        _optax_primal, _optax_rewrite, inputs, frozenset()
    )
    cost, grad_norm = _cost_and_grad_norm(cost_fn, value, batch)
    return MinimiseResult(
        value=value,
        cost=cost,
        grad_norm=grad_norm,
        num_steps=num_steps,
        converged=converged,
    )

OptimistixMinimiser

Bases: Module

An optimistix minimiser behind the Minimiser protocol.

OptimistixMinimiser() is the models' default made explicit: optx.BFGS(rtol=1e-6, atol=1e-6) with max_steps=200 and throw=False, so a model built with it returns exactly what the default model returns.

\[ \theta^\star = \arg\min_\theta J(\theta;\, \text{batch}) \]

Algorithm

1. Hand cost_fn(x, batch), the solver and x0 to optimistix.minimise, capped at max_steps iterations.
2. Return the minimiser; run also reports the cost, the gradient norm and the step count at the solution.

Raises:

Type	Description
ConfigurationError	If max_steps is below 1.

Attributes:

Name	Type	Description
solver	AbstractMinimiser	The optimistix.AbstractMinimiser.
max_steps	int	Iteration cap.
throw	bool	Passed to optimistix.minimise: raise on failure (non-convergence within max_steps) instead of returning the last iterate.

Examples:

A quadratic \(J(x) = \tfrac12 \lVert x - m \rVert^2\) has its minimum at \(m\).

>>> import jax.numpy as jnp, vardax
>>> m = jnp.array([1.0, -2.0, 0.5])
>>> minimiser = vardax.OptimistixMinimiser()
>>> x = minimiser(lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(3), m)
>>> bool(jnp.allclose(x, m, atol=1e-5))
True

The minimiser can return diagnostics, and any pytree is a valid control.

>>> result = minimiser.run(
...     lambda p, m: 0.5 * jnp.sum((p["a"] - m) ** 2) + 0.5 * p["b"] ** 2,
...     {"a": jnp.zeros(3), "b": jnp.ones(())},
...     m,
... )
>>> sorted(result.value)
['a', 'b']

An invalid iteration cap is rejected.

>>> vardax.OptimistixMinimiser(max_steps=0)
Traceback (most recent call last):
...
ConfigurationError: max_steps must be >= 1, got 0.

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

class OptimistixMinimiser(eqx.Module):
    r"""An ``optimistix`` minimiser behind the `Minimiser` protocol.

    ``OptimistixMinimiser()`` is the models' default made explicit:
    ``optx.BFGS(rtol=1e-6, atol=1e-6)`` with ``max_steps=200`` and
    ``throw=False``, so a model built with it returns exactly what the
    default model returns.

    $$ \theta^\star = \arg\min_\theta J(\theta;\, \text{batch}) $$

    Algorithm:
        1. Hand ``cost_fn(x, batch)``, the solver and ``x0`` to
           ``optimistix.minimise``, capped at ``max_steps`` iterations.
        2. Return the minimiser; `run` also reports the cost, the
           gradient norm and the step count at the solution.

    Raises:
        ConfigurationError: If ``max_steps`` is below 1.

    Attributes:
        solver: The ``optimistix.AbstractMinimiser``.
        max_steps: Iteration cap.
        throw: Passed to ``optimistix.minimise``: raise on failure
            (non-convergence within ``max_steps``) instead of returning
            the last iterate.

    Examples:
        A quadratic $J(x) = \tfrac12 \lVert x - m \rVert^2$ has its
        minimum at $m$.

        >>> import jax.numpy as jnp, vardax
        >>> m = jnp.array([1.0, -2.0, 0.5])
        >>> minimiser = vardax.OptimistixMinimiser()
        >>> x = minimiser(lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(3), m)
        >>> bool(jnp.allclose(x, m, atol=1e-5))
        True

        The minimiser can return diagnostics, and any pytree is a valid
        control.

        >>> result = minimiser.run(
        ...     lambda p, m: 0.5 * jnp.sum((p["a"] - m) ** 2) + 0.5 * p["b"] ** 2,
        ...     {"a": jnp.zeros(3), "b": jnp.ones(())},
        ...     m,
        ... )
        >>> sorted(result.value)
        ['a', 'b']

        An invalid iteration cap is rejected.

        >>> vardax.OptimistixMinimiser(max_steps=0)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: max_steps must be >= 1, got 0.
    """

    solver: optx.AbstractMinimiser = eqx.field(static=True)
    max_steps: int = eqx.field(static=True)
    throw: bool = eqx.field(static=True)

    def __init__(
        self,
        solver: optx.AbstractMinimiser | None = None,
        *,
        max_steps: int = 200,
        throw: bool = False,
    ) -> None:
        if max_steps < 1:
            raise ConfigurationError(f"max_steps must be >= 1, got {max_steps}.")
        self.solver = solver or optx.BFGS(rtol=1e-6, atol=1e-6)
        self.max_steps = max_steps
        self.throw = throw

    def run(
        self,
        cost_fn: Callable[[Any, Any], Float[Array, ""]],
        x0: PyTree[Array],
        batch: Any = None,
        *,
        adjoint: optx.AbstractAdjoint | None = None,
    ) -> MinimiseResult:
        """Minimise ``cost_fn(x, batch)`` from ``x0``, with diagnostics.

        Args:
            cost_fn: Scalar cost ``cost_fn(x, batch)``.
            x0: Initial guess, any pytree of arrays.
            batch: Passed as the second argument of ``cost_fn``.
            adjoint: How to differentiate through the minimum. Default
                ``optimistix.ImplicitAdjoint()``.

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

        Examples:
            >>> import jax.numpy as jnp, vardax
            >>> result = vardax.OptimistixMinimiser().run(
            ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
            ... )
            >>> bool(jnp.allclose(result.value, 1.0, atol=1e-5))
            True
        """
        solution = optx.minimise(
            fn=cost_fn,
            solver=self.solver,
            y0=x0,
            args=batch,
            max_steps=self.max_steps,
            adjoint=adjoint or optx.ImplicitAdjoint(),
            throw=self.throw,
        )
        value = solution.value
        cost, grad_norm = _cost_and_grad_norm(cost_fn, value, batch)
        return MinimiseResult(
            value=value,
            cost=cost,
            grad_norm=grad_norm,
            num_steps=solution.stats["num_steps"],
            converged=solution.result == optx.RESULTS.successful,
        )

    def __call__(
        self,
        cost_fn: Callable[[Any, Any], Float[Array, ""]],
        x0: PyTree[Array],
        batch: Any = None,
    ) -> PyTree[Array]:
        """Return ``argmin_x cost_fn(x, batch)`` starting from ``x0``.

        Args:
            cost_fn: Scalar cost ``cost_fn(x, batch)``.
            x0: Initial guess, any pytree of arrays.
            batch: Passed as the second argument of ``cost_fn``.

        Returns:
            The minimiser, with the structure of ``x0``.
        """
        return optx.minimise(
            fn=cost_fn,
            solver=self.solver,
            y0=x0,
            args=batch,
            max_steps=self.max_steps,
            throw=self.throw,
        ).value

run

run(
    cost_fn: Callable[[Any, Any], Float[Array, ""]],
    x0: PyTree[Array],
    batch: Any = None,
    *,
    adjoint: AbstractAdjoint | None = None,
) -> MinimiseResult

Minimise cost_fn(x, batch) from x0, with diagnostics.

Parameters:

Name	Type	Description	Default
cost_fn	Callable[[Any, Any], Float[Array, '']]	Scalar cost cost_fn(x, batch).	required
x0	PyTree[Array]	Initial guess, any pytree of arrays.	required
batch	Any	Passed as the second argument of cost_fn.	None
adjoint	AbstractAdjoint | None	How to differentiate through the minimum. Default optimistix.ImplicitAdjoint().	None

Returns:

Type	Description
MinimiseResult	A MinimiseResult.

Examples:

>>> import jax.numpy as jnp, vardax
>>> result = vardax.OptimistixMinimiser().run(
...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
... )
>>> bool(jnp.allclose(result.value, 1.0, atol=1e-5))
True

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

def run(
    self,
    cost_fn: Callable[[Any, Any], Float[Array, ""]],
    x0: PyTree[Array],
    batch: Any = None,
    *,
    adjoint: optx.AbstractAdjoint | None = None,
) -> MinimiseResult:
    """Minimise ``cost_fn(x, batch)`` from ``x0``, with diagnostics.

    Args:
        cost_fn: Scalar cost ``cost_fn(x, batch)``.
        x0: Initial guess, any pytree of arrays.
        batch: Passed as the second argument of ``cost_fn``.
        adjoint: How to differentiate through the minimum. Default
            ``optimistix.ImplicitAdjoint()``.

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

    Examples:
        >>> import jax.numpy as jnp, vardax
        >>> result = vardax.OptimistixMinimiser().run(
        ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), jnp.ones(2)
        ... )
        >>> bool(jnp.allclose(result.value, 1.0, atol=1e-5))
        True
    """
    solution = optx.minimise(
        fn=cost_fn,
        solver=self.solver,
        y0=x0,
        args=batch,
        max_steps=self.max_steps,
        adjoint=adjoint or optx.ImplicitAdjoint(),
        throw=self.throw,
    )
    value = solution.value
    cost, grad_norm = _cost_and_grad_norm(cost_fn, value, batch)
    return MinimiseResult(
        value=value,
        cost=cost,
        grad_norm=grad_norm,
        num_steps=solution.stats["num_steps"],
        converged=solution.result == optx.RESULTS.successful,
    )

MinimiseResult

Bases: Module

Outcome of a minimisation.

Attributes:

Name	Type	Description
value	PyTree[Array]	The minimiser \(\theta^\star\), with the structure of x0.
cost	Float[Array, '']	The cost \(J(\theta^\star)\).
grad_norm	Float[Array, '']	\(\lVert \nabla J(\theta^\star) \rVert_2\) over all leaves of \(\theta\). The implicit adjoint is exact only when this is (close to) zero.
num_steps	Int[Array, '']	Number of optimiser steps taken (both phases, for OptaxMinimiser with a polish).
converged	Bool[Array, ''] | None	Whether the minimiser met its own convergence test: optimistix's result == successful, or for OptaxMinimiser a small gradient or step_patience small steps (in the polish phase, when there is one); stopping on patience or max_steps is False. None for a custom Minimiser, which reports no status.

Examples:

The result carries the minimum and its diagnostics.

>>> import jax.numpy as jnp, vardax
>>> m = jnp.array([1.0, -2.0])
>>> result = vardax.OptimistixMinimiser().run(
...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), m
... )
>>> type(result).__name__
'MinimiseResult'
>>> bool(jnp.allclose(result.value, m, atol=1e-5))
True

At a minimum the gradient norm is near zero and the run converged.

>>> bool(result.grad_norm < 1e-4), bool(result.converged)
(True, True)

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

class MinimiseResult(eqx.Module):
    r"""Outcome of a minimisation.

    Attributes:
        value: The minimiser $\theta^\star$, with the structure of ``x0``.
        cost: The cost $J(\theta^\star)$.
        grad_norm: $\lVert \nabla J(\theta^\star) \rVert_2$ over all
            leaves of $\theta$. The implicit adjoint is exact only when
            this is (close to) zero.
        num_steps: Number of optimiser steps taken (both phases, for
            `OptaxMinimiser` with a polish).
        converged: Whether the minimiser met its own convergence test:
            optimistix's ``result == successful``, or for
            `OptaxMinimiser` a small gradient or ``step_patience`` small
            steps (in the polish phase, when there is one); stopping on
            ``patience`` or ``max_steps`` is ``False``. ``None`` for a custom
            [`Minimiser`][vardax.Minimiser], which reports no status.

    Examples:
        The result carries the minimum and its diagnostics.

        >>> import jax.numpy as jnp, vardax
        >>> m = jnp.array([1.0, -2.0])
        >>> result = vardax.OptimistixMinimiser().run(
        ...     lambda x, m: 0.5 * jnp.sum((x - m) ** 2), jnp.zeros(2), m
        ... )
        >>> type(result).__name__
        'MinimiseResult'
        >>> bool(jnp.allclose(result.value, m, atol=1e-5))
        True

        At a minimum the gradient norm is near zero and the run converged.

        >>> bool(result.grad_norm < 1e-4), bool(result.converged)
        (True, True)
    """

    value: PyTree[Array]
    cost: Float[Array, ""]
    grad_norm: Float[Array, ""]
    num_steps: Int[Array, ""]
    converged: Bool[Array, ""] | None = None

4DVarNet inner-loop solvers

The unrolled (and fixed-point) inner loop of 4DVarNet, exposed as pure functions over an explicit SolverState: initialise with init_solver_state_*, advance one modulated-gradient step with solver_step_* (or fp_solver_step_1d for the fixed-point formulation), or run the whole loop with solve_4dvarnet_*. The one_step_* variants pair with OneStepAdjoint for memory-frugal training.

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

SolverState1D

Bases: Module

Carry of the 1-D solver loop: iterate, ConvLSTM memory and step counter.

The state is the triple \((u_k, (h_k, c_k), k)\) threaded through the unrolled iteration \(u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k, c_k)\).

Attributes:

Name	Type	Description
x	Float[Array, 'B T N']	Current state estimate \(u_k\).
lstm	LSTMState1D	Current LSTM hidden/cell state \((h_k, c_k)\) of the gradient modulator.
step	int	Current iteration index \(k\).

Examples:

A state built by hand:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> state = vdx.SolverState1D(
...     x=jnp.zeros((2, 3, 4)), lstm=vdx.LSTMState1D.zeros(2, 5, 4), step=0
... )
>>> state.x.shape, state.lstm.hidden.shape, state.step
((2, 3, 4), (2, 5, 4), 0)

It is an Equinox module, so it is a pytree (x, hidden, cell and the step counter):

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

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

class SolverState1D(eqx.Module):
    r"""Carry of the 1-D solver loop: iterate, ConvLSTM memory and step counter.

    The state is the triple $(u_k, (h_k, c_k), k)$ threaded through the unrolled
    iteration $u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k, c_k)$.

    Attributes:
        x: Current state estimate $u_k$.
        lstm: Current LSTM hidden/cell state $(h_k, c_k)$ of the gradient modulator.
        step: Current iteration index $k$.

    Examples:
        A state built by hand:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> state = vdx.SolverState1D(
        ...     x=jnp.zeros((2, 3, 4)), lstm=vdx.LSTMState1D.zeros(2, 5, 4), step=0
        ... )
        >>> state.x.shape, state.lstm.hidden.shape, state.step
        ((2, 3, 4), (2, 5, 4), 0)

        It is an Equinox module, so it is a pytree (x, hidden, cell and the step counter):

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

    x: Float[Array, "B T N"]
    lstm: LSTMState1D
    step: int

SolverState2D

Bases: Module

Carry of the 2-D solver loop: iterate, ConvLSTM memory and step counter.

The state is the triple \((u_k, (h_k, c_k), k)\) threaded through the unrolled iteration \(u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k, c_k)\).

Attributes:

Name	Type	Description
x	Float[Array, 'B T H W']	Current state estimate \(u_k\).
lstm	LSTMState2D	Current LSTM hidden/cell state \((h_k, c_k)\) of the gradient modulator.
step	int	Current iteration index \(k\).

Examples:

A state built by hand:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> state = vdx.SolverState2D(
...     x=jnp.zeros((1, 3, 4, 5)),
...     lstm=vdx.LSTMState2D.zeros(1, 2, 4, 5),
...     step=0,
... )
>>> state.x.shape, state.lstm.cell.shape
((1, 3, 4, 5), (1, 2, 4, 5))

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

class SolverState2D(eqx.Module):
    r"""Carry of the 2-D solver loop: iterate, ConvLSTM memory and step counter.

    The state is the triple $(u_k, (h_k, c_k), k)$ threaded through the unrolled
    iteration $u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k, c_k)$.

    Attributes:
        x: Current state estimate $u_k$.
        lstm: Current LSTM hidden/cell state $(h_k, c_k)$ of the gradient modulator.
        step: Current iteration index $k$.

    Examples:
        A state built by hand:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> state = vdx.SolverState2D(
        ...     x=jnp.zeros((1, 3, 4, 5)),
        ...     lstm=vdx.LSTMState2D.zeros(1, 2, 4, 5),
        ...     step=0,
        ... )
        >>> state.x.shape, state.lstm.cell.shape
        ((1, 3, 4, 5), (1, 2, 4, 5))
    """

    x: Float[Array, "B T H W"]
    lstm: LSTMState2D
    step: int

init_solver_state_1d

init_solver_state_1d(
    batch: Batch1D, hidden_dim: int
) -> SolverState1D

Initialise a 1-D solver state from a batch.

The iterate starts at the masked observations, \(u_0 = m \odot y\), with a zero ConvLSTM memory and step counter zero.

Parameters:

Name	Type	Description	Default
batch	Batch1D	Input batch \((y, m)\). The initial state is set to the masked input (zeros where unobserved).	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required

Returns:

Type	Description
SolverState1D	Initial SolverState1D holding \(u_0\), a zero LSTM
SolverState1D	state and step=0.

Examples:

A typical call, with everything observed:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
>>> state = vdx.init_solver_state_1d(batch, hidden_dim=5)
>>> state.x.shape == (2, 3, 4), state.step
(True, 0)

Unobserved entries start at zero:

>>> empty = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.zeros((2, 3, 4)))
>>> bool(jnp.all(vdx.init_solver_state_1d(empty, hidden_dim=5).x == 0))
True

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

def init_solver_state_1d(
    batch: Batch1D,
    hidden_dim: int,
) -> SolverState1D:
    r"""Initialise a 1-D solver state from a batch.

    The iterate starts at the masked observations, $u_0 = m \odot y$, with a zero
    ConvLSTM memory and step counter zero.

    Args:
        batch: Input batch $(y, m)$.  The initial state is set to the masked
            input (zeros where unobserved).
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.

    Returns:
        Initial [`SolverState1D`][vardax.SolverState1D] holding $u_0$, a zero LSTM
        state and ``step=0``.

    Examples:
        A typical call, with everything observed:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
        >>> state = vdx.init_solver_state_1d(batch, hidden_dim=5)
        >>> state.x.shape == (2, 3, 4), state.step
        (True, 0)

        Unobserved entries start at zero:

        >>> empty = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.zeros((2, 3, 4)))
        >>> bool(jnp.all(vdx.init_solver_state_1d(empty, hidden_dim=5).x == 0))
        True
    """
    batch_size, _, n_space = batch.input.shape
    # u₀ = m ⊙ y
    initial_state = batch.input * batch.mask
    lstm = LSTMState1D.zeros(batch_size, hidden_dim, n_space)
    return SolverState1D(x=initial_state, lstm=lstm, step=0)

init_solver_state_2d

init_solver_state_2d(
    batch: Batch2D, hidden_dim: int
) -> SolverState2D

Initialise a 2-D solver state from a batch.

The iterate starts at the masked observations, \(u_0 = m \odot y\), with a zero ConvLSTM memory and step counter zero.

Parameters:

Name	Type	Description	Default
batch	Batch2D	Input batch \((y, m)\). The initial state is set to the masked input (zeros where unobserved).	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required

Returns:

Type	Description
SolverState2D	Initial SolverState2D holding \(u_0\), a zero LSTM
SolverState2D	state and step=0.

Examples:

A typical call, with everything observed:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
... )
>>> state = vdx.init_solver_state_2d(batch, hidden_dim=5)
>>> state.x.shape == (1, 3, 4, 5), state.step
(True, 0)

Unobserved entries start at zero:

>>> empty = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.zeros((1, 3, 4, 5))
... )
>>> bool(jnp.all(vdx.init_solver_state_2d(empty, hidden_dim=5).x == 0))
True

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

def init_solver_state_2d(
    batch: Batch2D,
    hidden_dim: int,
) -> SolverState2D:
    r"""Initialise a 2-D solver state from a batch.

    The iterate starts at the masked observations, $u_0 = m \odot y$, with a zero
    ConvLSTM memory and step counter zero.

    Args:
        batch: Input batch $(y, m)$.  The initial state is set to the masked
            input (zeros where unobserved).
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.

    Returns:
        Initial [`SolverState2D`][vardax.SolverState2D] holding $u_0$, a zero LSTM
        state and ``step=0``.

    Examples:
        A typical call, with everything observed:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
        ... )
        >>> state = vdx.init_solver_state_2d(batch, hidden_dim=5)
        >>> state.x.shape == (1, 3, 4, 5), state.step
        (True, 0)

        Unobserved entries start at zero:

        >>> empty = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.zeros((1, 3, 4, 5))
        ... )
        >>> bool(jnp.all(vdx.init_solver_state_2d(empty, hidden_dim=5).x == 0))
        True
    """
    batch_size, _, height, width = batch.input.shape
    # u₀ = m ⊙ y
    initial_state = batch.input * batch.mask
    lstm = LSTMState2D.zeros(batch_size, hidden_dim, height, width)
    return SolverState2D(x=initial_state, lstm=lstm, step=0)

solver_step_1d

solver_step_1d(
    solver_state: SolverState1D,
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
) -> SolverState1D

Perform a single 1-D solver iteration.

With the variational cost

\[ J(u) = \|m \odot (u - y)\|^2 + \lambda \|u - \varphi(u)\|^2 \]

one iteration moves the state along the learned, modulated gradient, \(u \leftarrow u - \alpha\, \mathcal{G}(\nabla J(u), u, h)\).

Algorithm

1. Evaluate the observation term \(\|m \odot (u - y)\|^2\) and the prior term \(\lambda \|u - \varphi(u)\|^2\) at the current state.
2. Differentiate the sum with jax.grad to obtain \(\nabla J(u)\).
3. Pass \((\nabla J, u, h)\) through grad_mod_fn to get the update and the new LSTM state.
4. Set \(u \leftarrow u - \alpha \cdot \text{update}\) and increment step.

Parameters:

Name	Type	Description	Default
solver_state	SolverState1D	Current solver state \((u, h, c, k)\).	required
batch	Batch1D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\) (prior autoencoder forward pass).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm), the gradient modulator \(\mathcal{G}\).	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0
prior_weight	float	Weighting factor \(\lambda\) for the prior cost term.	1.0

Returns:

Type	Description
SolverState1D	Updated SolverState1D: new iterate, new LSTM state
SolverState1D	and step incremented by one.

Examples:

With an identity prior and a pass-through modulator, one step with alpha=0.5 moves a zero state halfway to the observations (the observation gradient is \(-2\)):

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
>>> state = vdx.SolverState1D(
...     x=jnp.zeros((2, 3, 4)), lstm=vdx.LSTMState1D.zeros(2, 5, 4), step=0
... )
>>> passthrough = lambda grad, x, lstm: (grad, lstm)
>>> new = vdx.solver_step_1d(
...     state, batch, vdx.IdentityPrior(), passthrough, alpha=0.5
... )
>>> bool(jnp.allclose(new.x, 1.0)), new.step
(True, 1)

Where nothing is observed and the prior is the identity there is no gradient, so the state does not move:

>>> blind = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.zeros((2, 3, 4)))
>>> new = vdx.solver_step_1d(state, blind, vdx.IdentityPrior(), passthrough)
>>> bool(jnp.allclose(new.x, 0.0))
True

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

def solver_step_1d(
    solver_state: SolverState1D,
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
) -> SolverState1D:
    r"""Perform a single 1-D solver iteration.

    With the variational cost

    $$
    J(u) = \|m \odot (u - y)\|^2 + \lambda \|u - \varphi(u)\|^2
    $$

    one iteration moves the state along the learned, modulated gradient,
    $u \leftarrow u - \alpha\, \mathcal{G}(\nabla J(u), u, h)$.

    Algorithm:
        1. Evaluate the observation term $\|m \odot (u - y)\|^2$ and the prior
           term $\lambda \|u - \varphi(u)\|^2$ at the current state.
        2. Differentiate the sum with ``jax.grad`` to obtain $\nabla J(u)$.
        3. Pass $(\nabla J, u, h)$ through ``grad_mod_fn`` to get the update and
           the new LSTM state.
        4. Set $u \leftarrow u - \alpha \cdot \text{update}$ and increment ``step``.

    Args:
        solver_state: Current solver state $(u, h, c, k)$.
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$ (prior
            autoencoder forward pass).
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``, the
            gradient modulator $\mathcal{G}$.
        alpha: Step-size scaling factor $\alpha$.
        prior_weight: Weighting factor $\lambda$ for the prior cost term.

    Returns:
        Updated [`SolverState1D`][vardax.SolverState1D]: new iterate, new LSTM state
        and ``step`` incremented by one.

    Examples:
        With an identity prior and a pass-through modulator, one step with
        ``alpha=0.5`` moves a zero state halfway to the observations (the
        observation gradient is $-2$):

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
        >>> state = vdx.SolverState1D(
        ...     x=jnp.zeros((2, 3, 4)), lstm=vdx.LSTMState1D.zeros(2, 5, 4), step=0
        ... )
        >>> passthrough = lambda grad, x, lstm: (grad, lstm)
        >>> new = vdx.solver_step_1d(
        ...     state, batch, vdx.IdentityPrior(), passthrough, alpha=0.5
        ... )
        >>> bool(jnp.allclose(new.x, 1.0)), new.step
        (True, 1)

        Where nothing is observed and the prior is the identity there is no
        gradient, so the state does not move:

        >>> blind = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.zeros((2, 3, 4)))
        >>> new = vdx.solver_step_1d(state, blind, vdx.IdentityPrior(), passthrough)
        >>> bool(jnp.allclose(new.x, 0.0))
        True
    """
    state = solver_state.x

    def cost_fn(current_state):
        prior_state = prior_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 = prior_weight * jnp.sum((current_state - prior_state) ** 2)
        return obs_cost + prior_cost

    grad = jax.grad(cost_fn)(state)
    update, new_lstm = grad_mod_fn(grad, state, solver_state.lstm)
    # u ← u − α · update
    new_state = state - alpha * update

    return SolverState1D(x=new_state, lstm=new_lstm, step=solver_state.step + 1)

solver_step_2d

solver_step_2d(
    solver_state: SolverState2D,
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
) -> SolverState2D

Perform a single 2-D solver iteration.

With the variational cost

\[ J(u) = \|m \odot (u - y)\|^2 + \lambda \|u - \varphi(u)\|^2 \]

one iteration moves the state along the learned, modulated gradient, \(u \leftarrow u - \alpha\, \mathcal{G}(\nabla J(u), u, h)\).

Algorithm

1. Evaluate the observation term \(\|m \odot (u - y)\|^2\) and the prior term \(\lambda \|u - \varphi(u)\|^2\) at the current state.
2. Differentiate the sum with jax.grad to obtain \(\nabla J(u)\).
3. Pass \((\nabla J, u, h)\) through grad_mod_fn to get the update and the new LSTM state.
4. Set \(u \leftarrow u - \alpha \cdot \text{update}\) and increment step.

Parameters:

Name	Type	Description	Default
solver_state	SolverState2D	Current solver state \((u, h, c, k)\).	required
batch	Batch2D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\) (prior autoencoder forward pass).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm), the gradient modulator \(\mathcal{G}\).	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0
prior_weight	float	Weighting factor \(\lambda\) for the prior cost term.	1.0

Returns:

Type	Description
SolverState2D	Updated SolverState2D: new iterate, new LSTM state
SolverState2D	and step incremented by one.

Examples:

With an identity prior and a pass-through modulator, one step with alpha=0.5 moves a zero state halfway to the observations (the observation gradient is \(-2\)):

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
... )
>>> state = vdx.SolverState2D(
...     x=jnp.zeros((1, 3, 4, 5)),
...     lstm=vdx.LSTMState2D.zeros(1, 2, 4, 5),
...     step=0,
... )
>>> passthrough = lambda grad, x, lstm: (grad, lstm)
>>> new = vdx.solver_step_2d(
...     state, batch, vdx.IdentityPrior(), passthrough, alpha=0.5
... )
>>> bool(jnp.allclose(new.x, 1.0)), new.step
(True, 1)

Where nothing is observed and the prior is the identity there is no gradient, so the state does not move:

>>> blind = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.zeros((1, 3, 4, 5))
... )
>>> new = vdx.solver_step_2d(state, blind, vdx.IdentityPrior(), passthrough)
>>> bool(jnp.allclose(new.x, 0.0))
True

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

def solver_step_2d(
    solver_state: SolverState2D,
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
) -> SolverState2D:
    r"""Perform a single 2-D solver iteration.

    With the variational cost

    $$
    J(u) = \|m \odot (u - y)\|^2 + \lambda \|u - \varphi(u)\|^2
    $$

    one iteration moves the state along the learned, modulated gradient,
    $u \leftarrow u - \alpha\, \mathcal{G}(\nabla J(u), u, h)$.

    Algorithm:
        1. Evaluate the observation term $\|m \odot (u - y)\|^2$ and the prior
           term $\lambda \|u - \varphi(u)\|^2$ at the current state.
        2. Differentiate the sum with ``jax.grad`` to obtain $\nabla J(u)$.
        3. Pass $(\nabla J, u, h)$ through ``grad_mod_fn`` to get the update and
           the new LSTM state.
        4. Set $u \leftarrow u - \alpha \cdot \text{update}$ and increment ``step``.

    Args:
        solver_state: Current solver state $(u, h, c, k)$.
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$ (prior
            autoencoder forward pass).
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``, the
            gradient modulator $\mathcal{G}$.
        alpha: Step-size scaling factor $\alpha$.
        prior_weight: Weighting factor $\lambda$ for the prior cost term.

    Returns:
        Updated [`SolverState2D`][vardax.SolverState2D]: new iterate, new LSTM state
        and ``step`` incremented by one.

    Examples:
        With an identity prior and a pass-through modulator, one step with
        ``alpha=0.5`` moves a zero state halfway to the observations (the
        observation gradient is $-2$):

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
        ... )
        >>> state = vdx.SolverState2D(
        ...     x=jnp.zeros((1, 3, 4, 5)),
        ...     lstm=vdx.LSTMState2D.zeros(1, 2, 4, 5),
        ...     step=0,
        ... )
        >>> passthrough = lambda grad, x, lstm: (grad, lstm)
        >>> new = vdx.solver_step_2d(
        ...     state, batch, vdx.IdentityPrior(), passthrough, alpha=0.5
        ... )
        >>> bool(jnp.allclose(new.x, 1.0)), new.step
        (True, 1)

        Where nothing is observed and the prior is the identity there is no
        gradient, so the state does not move:

        >>> blind = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.zeros((1, 3, 4, 5))
        ... )
        >>> new = vdx.solver_step_2d(state, blind, vdx.IdentityPrior(), passthrough)
        >>> bool(jnp.allclose(new.x, 0.0))
        True
    """
    state = solver_state.x

    def cost_fn(current_state):
        prior_state = prior_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 = prior_weight * jnp.sum((current_state - prior_state) ** 2)
        return obs_cost + prior_cost

    grad = jax.grad(cost_fn)(state)
    update, new_lstm = grad_mod_fn(grad, state, solver_state.lstm)
    # u ← u − α · update
    new_state = state - alpha * update

    return SolverState2D(x=new_state, lstm=new_lstm, step=solver_state.step + 1)

fp_solver_step_1d

fp_solver_step_1d(
    x: Float[Array, "B T N"], batch: Batch1D, prior_fn: Any
) -> Float[Array, "B T N"]

Perform a single 1-D fixed-point projection step.

Applies the prior projection then re-inserts observations at observed locations:

\[ x \leftarrow \varphi(x), \quad x \leftarrow m \odot y + (1 - m) \odot x \]

Algorithm

1. Project the state with the prior, \(\varphi(x)\).
2. Overwrite observed entries with the data and keep the projection elsewhere, \(m \odot y + (1 - m) \odot \varphi(x)\).

Parameters:

Name	Type	Description	Default
x	Float[Array, 'B T N']	Current state estimate \(x\).	required
batch	Batch1D	Observed data batch containing input (observations \(y\)) and mask (\(m\)).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(x)\) (prior autoencoder forward pass).	required

Returns:

Type	Description
Float[Array, 'B T N']	Updated state estimate \(m \odot y + (1 - m) \odot \varphi(x)\).

Examples:

Observed entries are reinserted, unobserved ones come from the prior:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch1D(
...     input=jnp.ones((1, 2, 3)),
...     mask=jnp.array([[[1.0, 0.0, 1.0], [0.0, 0.0, 1.0]]]),
... )
>>> prior = lambda x: 5.0 * jnp.ones_like(x)
>>> out = vdx.fp_solver_step_1d(jnp.zeros((1, 2, 3)), batch, prior)
>>> out[0].tolist()
[[1.0, 5.0, 1.0], [5.0, 5.0, 1.0]]

A fully observed batch returns the data whatever the prior says:

>>> full = vdx.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
>>> bool(jnp.all(vdx.fp_solver_step_1d(jnp.zeros((1, 2, 3)), full, prior) == 1))
True

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

def fp_solver_step_1d(
    x: Float[Array, "B T N"],
    batch: Batch1D,
    prior_fn: Any,
) -> Float[Array, "B T N"]:
    r"""Perform a single 1-D fixed-point projection step.

    Applies the prior projection then re-inserts observations at observed
    locations:

    $$
    x \leftarrow \varphi(x), \quad
    x \leftarrow m \odot y + (1 - m) \odot x
    $$

    Algorithm:
        1. Project the state with the prior, $\varphi(x)$.
        2. Overwrite observed entries with the data and keep the projection
           elsewhere, $m \odot y + (1 - m) \odot \varphi(x)$.

    Args:
        x: Current state estimate $x$.
        batch: Observed data batch containing ``input`` (observations $y$)
            and ``mask`` ($m$).
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(x)$ (prior
            autoencoder forward pass).

    Returns:
        Updated state estimate $m \odot y + (1 - m) \odot \varphi(x)$.

    Examples:
        Observed entries are reinserted, unobserved ones come from the prior:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(
        ...     input=jnp.ones((1, 2, 3)),
        ...     mask=jnp.array([[[1.0, 0.0, 1.0], [0.0, 0.0, 1.0]]]),
        ... )
        >>> prior = lambda x: 5.0 * jnp.ones_like(x)
        >>> out = vdx.fp_solver_step_1d(jnp.zeros((1, 2, 3)), batch, prior)
        >>> out[0].tolist()
        [[1.0, 5.0, 1.0], [5.0, 5.0, 1.0]]

        A fully observed batch returns the data whatever the prior says:

        >>> full = vdx.Batch1D(input=jnp.ones((1, 2, 3)), mask=jnp.ones((1, 2, 3)))
        >>> bool(jnp.all(vdx.fp_solver_step_1d(jnp.zeros((1, 2, 3)), full, prior) == 1))
        True
    """
    prior_state = prior_fn(x)
    # u ← m ⊙ y + (1 − m) ⊙ φ(u)
    return batch.mask * batch.input + (1 - batch.mask) * prior_state

solve_4dvarnet_1d

solve_4dvarnet_1d(
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
) -> Float[Array, "B T N"]

Run the full 1-D 4DVarNet solver for n_steps iterations.

Unrolls \(u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k)\) from \(u_0 = m \odot y\), with \(J\) the variational cost of solver_step_1d.

Algorithm

1. Initialise the state with init_solver_state_1d.
2. Repeat n_steps times: apply solver_step_1d to the state.
3. Return the final iterate.

Parameters:

Name	Type	Description	Default
batch	Batch1D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm).	required
n_steps	int	Number of gradient-descent steps to unroll.	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0

Returns:

Type	Description
Float[Array, 'B T N']	Final state estimate \(u_{n}\) after n_steps iterations.

Examples:

Fully observed data is already a minimiser for an identity prior, so the solver returns it unchanged:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
>>> passthrough = lambda grad, x, lstm: (grad, lstm)
>>> out = vdx.solve_4dvarnet_1d(
...     batch,
...     vdx.IdentityPrior(),
...     passthrough,
...     n_steps=3,
...     hidden_dim=5,
...     alpha=0.5,
... )
>>> out.shape == (2, 3, 4), bool(jnp.allclose(out, 1.0))
(True, True)

With n_steps=0 the result is the masked input:

>>> out = vdx.solve_4dvarnet_1d(
...     batch, vdx.IdentityPrior(), passthrough, n_steps=0, hidden_dim=5
... )
>>> bool(jnp.allclose(out, batch.input * batch.mask))
True

The learned pieces of a FourDVarNet1D plug in directly and reproduce its unrolled forward pass:

>>> import jax
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     alpha=0.2,
...     key=jax.random.PRNGKey(0),
... )
>>> out = vdx.solve_4dvarnet_1d(
...     batch,
...     model.prior,
...     model.grad_mod,
...     n_steps=2,
...     hidden_dim=3,
...     alpha=0.2,
... )
>>> bool(jnp.allclose(out, model(batch), atol=1e-5))
True

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

def solve_4dvarnet_1d(
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
) -> Float[Array, "B T N"]:
    r"""Run the full 1-D 4DVarNet solver for ``n_steps`` iterations.

    Unrolls $u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k)$ from
    $u_0 = m \odot y$, with $J$ the variational cost of
    [`solver_step_1d`][vardax.solver_step_1d].

    Algorithm:
        1. Initialise the state with
           [`init_solver_state_1d`][vardax.init_solver_state_1d].
        2. Repeat ``n_steps`` times: apply
           [`solver_step_1d`][vardax.solver_step_1d] to the state.
        3. Return the final iterate.

    Args:
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$.
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``.
        n_steps: Number of gradient-descent steps to unroll.
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.
        alpha: Step-size scaling factor $\alpha$.

    Returns:
        Final state estimate $u_{n}$ after ``n_steps`` iterations.

    Examples:
        Fully observed data is already a minimiser for an identity prior, so
        the solver returns it unchanged:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
        >>> passthrough = lambda grad, x, lstm: (grad, lstm)
        >>> out = vdx.solve_4dvarnet_1d(
        ...     batch,
        ...     vdx.IdentityPrior(),
        ...     passthrough,
        ...     n_steps=3,
        ...     hidden_dim=5,
        ...     alpha=0.5,
        ... )
        >>> out.shape == (2, 3, 4), bool(jnp.allclose(out, 1.0))
        (True, True)

        With ``n_steps=0`` the result is the masked input:

        >>> out = vdx.solve_4dvarnet_1d(
        ...     batch, vdx.IdentityPrior(), passthrough, n_steps=0, hidden_dim=5
        ... )
        >>> bool(jnp.allclose(out, batch.input * batch.mask))
        True

        The learned pieces of a [`FourDVarNet1D`][vardax.FourDVarNet1D] plug in
        directly and reproduce its unrolled forward pass:

        >>> import jax
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     alpha=0.2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> out = vdx.solve_4dvarnet_1d(
        ...     batch,
        ...     model.prior,
        ...     model.grad_mod,
        ...     n_steps=2,
        ...     hidden_dim=3,
        ...     alpha=0.2,
        ... )
        >>> bool(jnp.allclose(out, model(batch), atol=1e-5))
        True
    """
    state = init_solver_state_1d(batch, hidden_dim)
    for _ in range(n_steps):
        state = solver_step_1d(state, batch, prior_fn, grad_mod_fn, alpha)
    return state.x

solve_4dvarnet_2d

solve_4dvarnet_2d(
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
) -> Float[Array, "B T H W"]

Run the full 2-D 4DVarNet solver for n_steps iterations.

Unrolls \(u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k)\) from \(u_0 = m \odot y\), with \(J\) the variational cost of solver_step_2d.

Algorithm

1. Initialise the state with init_solver_state_2d.
2. Repeat n_steps times: apply solver_step_2d to the state.
3. Return the final iterate.

Parameters:

Name	Type	Description	Default
batch	Batch2D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm).	required
n_steps	int	Number of gradient-descent steps to unroll.	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0

Returns:

Type	Description
Float[Array, 'B T H W']	Final state estimate \(u_{n}\) after n_steps iterations.

Examples:

Fully observed data is already a minimiser for an identity prior, so the solver returns it unchanged:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
... )
>>> passthrough = lambda grad, x, lstm: (grad, lstm)
>>> out = vdx.solve_4dvarnet_2d(
...     batch,
...     vdx.IdentityPrior(),
...     passthrough,
...     n_steps=3,
...     hidden_dim=5,
...     alpha=0.5,
... )
>>> out.shape == (1, 3, 4, 5), bool(jnp.allclose(out, 1.0))
(True, True)

With n_steps=0 the result is the masked input:

>>> out = vdx.solve_4dvarnet_2d(
...     batch, vdx.IdentityPrior(), passthrough, n_steps=0, hidden_dim=5
... )
>>> bool(jnp.allclose(out, batch.input * batch.mask))
True

The learned pieces of a FourDVarNet2D plug in directly and reproduce its unrolled forward pass:

>>> import jax
>>> model = vdx.FourDVarNet2D(
...     n_time=3,
...     height=4,
...     width=5,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=2,
...     alpha=0.2,
...     key=jax.random.PRNGKey(0),
... )
>>> out = vdx.solve_4dvarnet_2d(
...     batch,
...     model.prior,
...     model.grad_mod,
...     n_steps=2,
...     hidden_dim=3,
...     alpha=0.2,
... )
>>> bool(jnp.allclose(out, model(batch), atol=1e-5))
True

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

def solve_4dvarnet_2d(
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
) -> Float[Array, "B T H W"]:
    r"""Run the full 2-D 4DVarNet solver for ``n_steps`` iterations.

    Unrolls $u_{k+1} = u_k - \alpha\, \mathcal{G}(\nabla J(u_k), u_k, h_k)$ from
    $u_0 = m \odot y$, with $J$ the variational cost of
    [`solver_step_2d`][vardax.solver_step_2d].

    Algorithm:
        1. Initialise the state with
           [`init_solver_state_2d`][vardax.init_solver_state_2d].
        2. Repeat ``n_steps`` times: apply
           [`solver_step_2d`][vardax.solver_step_2d] to the state.
        3. Return the final iterate.

    Args:
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$.
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``.
        n_steps: Number of gradient-descent steps to unroll.
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.
        alpha: Step-size scaling factor $\alpha$.

    Returns:
        Final state estimate $u_{n}$ after ``n_steps`` iterations.

    Examples:
        Fully observed data is already a minimiser for an identity prior, so
        the solver returns it unchanged:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
        ... )
        >>> passthrough = lambda grad, x, lstm: (grad, lstm)
        >>> out = vdx.solve_4dvarnet_2d(
        ...     batch,
        ...     vdx.IdentityPrior(),
        ...     passthrough,
        ...     n_steps=3,
        ...     hidden_dim=5,
        ...     alpha=0.5,
        ... )
        >>> out.shape == (1, 3, 4, 5), bool(jnp.allclose(out, 1.0))
        (True, True)

        With ``n_steps=0`` the result is the masked input:

        >>> out = vdx.solve_4dvarnet_2d(
        ...     batch, vdx.IdentityPrior(), passthrough, n_steps=0, hidden_dim=5
        ... )
        >>> bool(jnp.allclose(out, batch.input * batch.mask))
        True

        The learned pieces of a [`FourDVarNet2D`][vardax.FourDVarNet2D] plug in
        directly and reproduce its unrolled forward pass:

        >>> import jax
        >>> model = vdx.FourDVarNet2D(
        ...     n_time=3,
        ...     height=4,
        ...     width=5,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=2,
        ...     alpha=0.2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> out = vdx.solve_4dvarnet_2d(
        ...     batch,
        ...     model.prior,
        ...     model.grad_mod,
        ...     n_steps=2,
        ...     hidden_dim=3,
        ...     alpha=0.2,
        ... )
        >>> bool(jnp.allclose(out, model(batch), atol=1e-5))
        True
    """
    state = init_solver_state_2d(batch, hidden_dim)
    for _ in range(n_steps):
        state = solver_step_2d(state, batch, prior_fn, grad_mod_fn, alpha)
    return state.x

solve_4dvarnet_1d_fixedpoint

solve_4dvarnet_1d_fixedpoint(
    batch: Batch1D, prior_fn: Any, n_fp_steps: int
) -> Float[Array, "B T N"]

Run n_fp_steps fixed-point projection steps using jax.lax.scan.

Iterates \(x_{k+1} = m \odot y + (1 - m) \odot \varphi(x_k)\) from the masked observations, the projection used by the implicit-adjoint path.

Algorithm

1. Initialise \(x_0 = m \odot y\).
2. Scan fp_solver_step_1d for n_fp_steps steps.
3. Return the final iterate.

Parameters:

Name	Type	Description	Default
batch	Batch1D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(x)\).	required
n_fp_steps	int	Number of fixed-point iterations.	required

Returns:

Type	Description
Float[Array, 'B T N']	Final state estimate \(x_{n}\) after n_fp_steps projections.

Examples:

With an identity prior the unobserved entries stay at zero:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> batch = vdx.Batch1D(
...     input=jnp.ones((1, 2, 3)),
...     mask=jnp.array([[[1.0, 0.0, 1.0], [0.0, 0.0, 1.0]]]),
... )
>>> out = vdx.solve_4dvarnet_1d_fixedpoint(batch, vdx.IdentityPrior(), 4)
>>> bool(jnp.allclose(out, batch.input * batch.mask))
True

Zero steps return the masked input, and a constant prior fills the gaps after one step:

>>> zero = vdx.solve_4dvarnet_1d_fixedpoint(batch, vdx.IdentityPrior(), 0)
>>> bool(jnp.allclose(zero, batch.input * batch.mask))
True
>>> out = vdx.solve_4dvarnet_1d_fixedpoint(batch, lambda x: 2 + 0 * x, 1)
>>> out[0, 0].tolist()
[1.0, 2.0, 1.0]

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

def solve_4dvarnet_1d_fixedpoint(
    batch: Batch1D,
    prior_fn: Any,
    n_fp_steps: int,
) -> Float[Array, "B T N"]:
    r"""Run ``n_fp_steps`` fixed-point projection steps using ``jax.lax.scan``.

    Iterates $x_{k+1} = m \odot y + (1 - m) \odot \varphi(x_k)$ from the masked
    observations, the projection used by the implicit-adjoint path.

    Algorithm:
        1. Initialise $x_0 = m \odot y$.
        2. Scan [`fp_solver_step_1d`][vardax.fp_solver_step_1d] for
           ``n_fp_steps`` steps.
        3. Return the final iterate.

    Args:
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(x)$.
        n_fp_steps: Number of fixed-point iterations.

    Returns:
        Final state estimate $x_{n}$ after ``n_fp_steps`` projections.

    Examples:
        With an identity prior the unobserved entries stay at zero:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> batch = vdx.Batch1D(
        ...     input=jnp.ones((1, 2, 3)),
        ...     mask=jnp.array([[[1.0, 0.0, 1.0], [0.0, 0.0, 1.0]]]),
        ... )
        >>> out = vdx.solve_4dvarnet_1d_fixedpoint(batch, vdx.IdentityPrior(), 4)
        >>> bool(jnp.allclose(out, batch.input * batch.mask))
        True

        Zero steps return the masked input, and a constant prior fills the
        gaps after one step:

        >>> zero = vdx.solve_4dvarnet_1d_fixedpoint(batch, vdx.IdentityPrior(), 0)
        >>> bool(jnp.allclose(zero, batch.input * batch.mask))
        True
        >>> out = vdx.solve_4dvarnet_1d_fixedpoint(batch, lambda x: 2 + 0 * x, 1)
        >>> out[0, 0].tolist()
        [1.0, 2.0, 1.0]
    """
    # u₀ = m ⊙ y
    initial_state = batch.input * batch.mask

    def scan_fn(carry: Float[Array, "B T N"], _: None) -> tuple:
        new_state = fp_solver_step_1d(carry, batch, prior_fn)
        return new_state, None

    final_state, _ = jax.lax.scan(scan_fn, initial_state, None, length=n_fp_steps)
    return final_state

one_step_solve_4dvarnet_1d

one_step_solve_4dvarnet_1d(
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
    k: int = 1,
) -> Float[Array, "B T N"]

Solve 4DVarNet-1D using k-step differentiation (Bolte et al., 2023).

Runs the same iteration as solve_4dvarnet_1d, \(u_{j+1} = u_j - \alpha\, \mathcal{G}(\nabla J(u_j), u_j, h_j)\), but only the last k steps carry gradients. This gives O(k) memory cost (k=1 matches implicit differentiation) while being as simple to implement as unrolled backprop.

Algorithm

1. Initialise the state from the masked observations.
2. Run n_steps - k warm-up steps and apply jax.lax.stop_gradient to the iterate and LSTM state (skipped when there is no warm-up).
3. Run the remaining k steps through which gradients flow.
4. Return the final iterate.

References

Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation of iterative algorithms." https://arxiv.org/abs/2305.13768

Parameters:

Name	Type	Description	Default
batch	Batch1D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm).	required
n_steps	int	Total number of solver iterations (warmup = n_steps - k, then k differentiable steps).	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0
prior_weight	float	Weighting factor \(\lambda\) for the prior cost term.	1.0
k	int	Number of trailing differentiable steps (clipped to n_steps).	1

Returns:

Type	Description
Float[Array, 'B T N']	Final state estimate after n_steps iterations.

Examples:

The forward value equals the fully unrolled solve; only the gradient path differs:

>>> 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=3,
...     alpha=0.2,
...     key=jax.random.PRNGKey(0),
... )
>>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
>>> kwargs = dict(n_steps=3, hidden_dim=3, alpha=0.2)
>>> full = vdx.solve_4dvarnet_1d(batch, model.prior, model.grad_mod, **kwargs)
>>> one = vdx.one_step_solve_4dvarnet_1d(
...     batch, model.prior, model.grad_mod, k=1, **kwargs
... )
>>> bool(jnp.allclose(full, one, atol=1e-5))
True

k >= n_steps disables the warm-up, so the solve is the same value again:

>>> big = vdx.one_step_solve_4dvarnet_1d(
...     batch, model.prior, model.grad_mod, k=10, **kwargs
... )
>>> bool(jnp.allclose(full, big, atol=1e-5))
True

It is what FourDVarNet1D runs under a OneStepAdjoint:

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

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

def one_step_solve_4dvarnet_1d(
    batch: Batch1D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
    k: int = 1,
) -> Float[Array, "B T N"]:
    r"""Solve 4DVarNet-1D using k-step differentiation (Bolte et al., 2023).

    Runs the same iteration as [`solve_4dvarnet_1d`][vardax.solve_4dvarnet_1d],
    $u_{j+1} = u_j - \alpha\, \mathcal{G}(\nabla J(u_j), u_j, h_j)$, but
    only the last ``k`` steps carry gradients.  This gives O(k) memory cost
    (k=1 matches implicit differentiation) while being as simple to
    implement as unrolled backprop.

    Algorithm:
        1. Initialise the state from the masked observations.
        2. Run ``n_steps - k`` warm-up steps and apply ``jax.lax.stop_gradient``
           to the iterate and LSTM state (skipped when there is no warm-up).
        3. Run the remaining ``k`` steps through which gradients flow.
        4. Return the final iterate.

    References:
        Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation of
        iterative algorithms." https://arxiv.org/abs/2305.13768

    Args:
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$.
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``.
        n_steps: Total number of solver iterations (warmup = n_steps - k,
            then k differentiable steps).
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.
        alpha: Step-size scaling factor $\alpha$.
        prior_weight: Weighting factor $\lambda$ for the prior cost term.
        k: Number of trailing differentiable steps (clipped to n_steps).

    Returns:
        Final state estimate after ``n_steps`` iterations.

    Examples:
        The forward value equals the fully unrolled solve; only the gradient
        path differs:

        >>> 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=3,
        ...     alpha=0.2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> batch = vdx.Batch1D(input=jnp.ones((2, 3, 4)), mask=jnp.ones((2, 3, 4)))
        >>> kwargs = dict(n_steps=3, hidden_dim=3, alpha=0.2)
        >>> full = vdx.solve_4dvarnet_1d(batch, model.prior, model.grad_mod, **kwargs)
        >>> one = vdx.one_step_solve_4dvarnet_1d(
        ...     batch, model.prior, model.grad_mod, k=1, **kwargs
        ... )
        >>> bool(jnp.allclose(full, one, atol=1e-5))
        True

        ``k >= n_steps`` disables the warm-up, so the solve is the same
        value again:

        >>> big = vdx.one_step_solve_4dvarnet_1d(
        ...     batch, model.prior, model.grad_mod, k=10, **kwargs
        ... )
        >>> bool(jnp.allclose(full, big, atol=1e-5))
        True

        It is what [`FourDVarNet1D`][vardax.FourDVarNet1D] runs under a
        [`OneStepAdjoint`][vardax.OneStepAdjoint]:

        >>> fast = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=3,
        ...     alpha=0.2,
        ...     solver_adjoint=vdx.OneStepAdjoint(),
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> bool(jnp.allclose(fast(batch), one, atol=1e-5))
        True
    """
    # --- warmup: run n_steps-k steps without tracking gradients ---
    state = init_solver_state_1d(batch, hidden_dim)
    live_steps = min(max(k, 1), n_steps) if n_steps >= 1 else 0
    warmup_steps = max(n_steps - live_steps, 0)
    for _ in range(warmup_steps):
        state = solver_step_1d(state, batch, prior_fn, grad_mod_fn, alpha, prior_weight)

    # detach the iterate so earlier steps don't contribute to the
    # gradient — but only when there was a warmup to detach: with
    # k >= n_steps the solve is fully differentiable, including
    # gradients with respect to the batch through the initial state.
    if warmup_steps > 0:
        state = SolverState1D(
            x=jax.lax.stop_gradient(state.x),
            lstm=jax.lax.stop_gradient(state.lstm),
            step=state.step,
        )

    # --- k differentiable steps ---
    for _ in range(live_steps):
        state = solver_step_1d(state, batch, prior_fn, grad_mod_fn, alpha, prior_weight)

    return state.x

one_step_solve_4dvarnet_2d

one_step_solve_4dvarnet_2d(
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
    k: int = 1,
) -> Float[Array, "B T H W"]

Solve 4DVarNet-2D using k-step differentiation (Bolte et al., 2023).

Runs the same iteration as solve_4dvarnet_2d, \(u_{j+1} = u_j - \alpha\, \mathcal{G}(\nabla J(u_j), u_j, h_j)\), but only the last k steps carry gradients. This gives O(k) memory cost (k=1 matches implicit differentiation) while being as simple to implement as unrolled backprop.

Algorithm

1. Initialise the state from the masked observations.
2. Run n_steps - k warm-up steps and apply jax.lax.stop_gradient to the iterate and LSTM state (skipped when there is no warm-up).
3. Run the remaining k steps through which gradients flow.
4. Return the final iterate.

References

Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation of iterative algorithms." https://arxiv.org/abs/2305.13768

Parameters:

Name	Type	Description	Default
batch	Batch2D	Observed data batch \((y, m)\).	required
prior_fn	Any	Callable x -> x_prior giving \(\varphi(u)\).	required
grad_mod_fn	Any	Callable (grad, x, lstm) -> (update, new_lstm).	required
n_steps	int	Total number of solver iterations (warmup = n_steps - k, then k differentiable steps).	required
hidden_dim	int	Hidden dimension of the ConvLSTM gradient modulator.	required
alpha	float	Step-size scaling factor \(\alpha\).	1.0
prior_weight	float	Weighting factor \(\lambda\) for the prior cost term.	1.0
k	int	Number of trailing differentiable steps (clipped to n_steps).	1

Returns:

Type	Description
Float[Array, 'B T H W']	Final state estimate after n_steps iterations.

Examples:

The forward value equals the fully unrolled solve; only the gradient path differs:

>>> 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=3,
...     alpha=0.2,
...     key=jax.random.PRNGKey(0),
... )
>>> batch = vdx.Batch2D(
...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
... )
>>> kwargs = dict(n_steps=3, hidden_dim=3, alpha=0.2)
>>> full = vdx.solve_4dvarnet_2d(batch, model.prior, model.grad_mod, **kwargs)
>>> one = vdx.one_step_solve_4dvarnet_2d(
...     batch, model.prior, model.grad_mod, k=1, **kwargs
... )
>>> bool(jnp.allclose(full, one, atol=1e-5))
True

k >= n_steps disables the warm-up, so the solve is the same value again:

>>> big = vdx.one_step_solve_4dvarnet_2d(
...     batch, model.prior, model.grad_mod, k=10, **kwargs
... )
>>> bool(jnp.allclose(full, big, atol=1e-5))
True

It is what FourDVarNet2D runs under a OneStepAdjoint:

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

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

def one_step_solve_4dvarnet_2d(
    batch: Batch2D,
    prior_fn: Any,
    grad_mod_fn: Any,
    n_steps: int,
    hidden_dim: int,
    alpha: float = 1.0,
    prior_weight: float = 1.0,
    k: int = 1,
) -> Float[Array, "B T H W"]:
    r"""Solve 4DVarNet-2D using k-step differentiation (Bolte et al., 2023).

    Runs the same iteration as [`solve_4dvarnet_2d`][vardax.solve_4dvarnet_2d],
    $u_{j+1} = u_j - \alpha\, \mathcal{G}(\nabla J(u_j), u_j, h_j)$, but
    only the last ``k`` steps carry gradients.  This gives O(k) memory cost
    (k=1 matches implicit differentiation) while being as simple to
    implement as unrolled backprop.

    Algorithm:
        1. Initialise the state from the masked observations.
        2. Run ``n_steps - k`` warm-up steps and apply ``jax.lax.stop_gradient``
           to the iterate and LSTM state (skipped when there is no warm-up).
        3. Run the remaining ``k`` steps through which gradients flow.
        4. Return the final iterate.

    References:
        Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation of
        iterative algorithms." https://arxiv.org/abs/2305.13768

    Args:
        batch: Observed data batch $(y, m)$.
        prior_fn: Callable ``x -> x_prior`` giving $\varphi(u)$.
        grad_mod_fn: Callable ``(grad, x, lstm) -> (update, new_lstm)``.
        n_steps: Total number of solver iterations (warmup = n_steps - k,
            then k differentiable steps).
        hidden_dim: Hidden dimension of the ConvLSTM gradient modulator.
        alpha: Step-size scaling factor $\alpha$.
        prior_weight: Weighting factor $\lambda$ for the prior cost term.
        k: Number of trailing differentiable steps (clipped to n_steps).

    Returns:
        Final state estimate after ``n_steps`` iterations.

    Examples:
        The forward value equals the fully unrolled solve; only the gradient
        path differs:

        >>> 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=3,
        ...     alpha=0.2,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> batch = vdx.Batch2D(
        ...     input=jnp.ones((1, 3, 4, 5)), mask=jnp.ones((1, 3, 4, 5))
        ... )
        >>> kwargs = dict(n_steps=3, hidden_dim=3, alpha=0.2)
        >>> full = vdx.solve_4dvarnet_2d(batch, model.prior, model.grad_mod, **kwargs)
        >>> one = vdx.one_step_solve_4dvarnet_2d(
        ...     batch, model.prior, model.grad_mod, k=1, **kwargs
        ... )
        >>> bool(jnp.allclose(full, one, atol=1e-5))
        True

        ``k >= n_steps`` disables the warm-up, so the solve is the same
        value again:

        >>> big = vdx.one_step_solve_4dvarnet_2d(
        ...     batch, model.prior, model.grad_mod, k=10, **kwargs
        ... )
        >>> bool(jnp.allclose(full, big, atol=1e-5))
        True

        It is what [`FourDVarNet2D`][vardax.FourDVarNet2D] runs under a
        [`OneStepAdjoint`][vardax.OneStepAdjoint]:

        >>> fast = vdx.FourDVarNet2D(
        ...     n_time=3,
        ...     height=4,
        ...     width=5,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=3,
        ...     alpha=0.2,
        ...     solver_adjoint=vdx.OneStepAdjoint(),
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> bool(jnp.allclose(fast(batch), one, atol=1e-5))
        True
    """
    # --- warmup: run n_steps-k steps without tracking gradients ---
    state = init_solver_state_2d(batch, hidden_dim)
    live_steps = min(max(k, 1), n_steps) if n_steps >= 1 else 0
    warmup_steps = max(n_steps - live_steps, 0)
    for _ in range(warmup_steps):
        state = solver_step_2d(state, batch, prior_fn, grad_mod_fn, alpha, prior_weight)

    # detach the iterate so earlier steps don't contribute to the
    # gradient — but only when there was a warmup to detach: with
    # k >= n_steps the solve is fully differentiable, including
    # gradients with respect to the batch through the initial state.
    if warmup_steps > 0:
        state = SolverState2D(
            x=jax.lax.stop_gradient(state.x),
            lstm=jax.lax.stop_gradient(state.lstm),
            step=state.step,
        )

    # --- k differentiable steps ---
    for _ in range(live_steps):
        state = solver_step_2d(state, batch, prior_fn, grad_mod_fn, alpha, prior_weight)

    return state.x