Utilities & Diagnostics¶

Supporting cast: a reference dynamical system for experiments, the statistical gates that decide whether an inference setup can be trusted, and the plotting helpers used throughout the end-to-end examples.

Dynamical systems¶

Lorenz-63 and Lorenz-96 are the standard chaotic testbeds for assimilation experiments; the simulate_* helpers generate trajectories for the Lorenz examples and the test suite, and time_patches produces the consecutive time pairs consumed by the one-step increment losses.

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

Lorenz63 ¶

Bases: Module

Lorenz-63 vector field as an Equinox module.

Implements the system

\[ \dot x = \sigma (y - x), \quad \dot y = x (\rho - z) - y, \quad \dot z = x y - \beta z . \]

The module is a diffrax-compatible right-hand side f(t, y, args).

Attributes:

Name	Type	Description
`sigma`	`float`	Prandtl number \(\sigma\) (classically 10).
`rho`	`float`	Rayleigh number \(\rho\) (classically 28).
`beta`	`float`	Geometric factor \(\beta\) (classically 8/3).

Examples:

The vector field at \((1, 1, 1)\) for the classical parameters:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> f = vdx.Lorenz63(sigma=10.0, rho=28.0, beta=8.0 / 3.0)
>>> dydt = f(0.0, jnp.ones(3), None)
>>> (
...     dydt.shape,
...     bool(jnp.allclose(dydt, jnp.array([0.0, 26.0, 1.0 - 8.0 / 3.0]))),
... )
((3,), True)

The origin is a fixed point:

>>> bool(jnp.all(f(0.0, jnp.zeros(3), None) == 0))
True

It drives a DynTrajectory prior:

>>> ts = jnp.linspace(0.0, 0.05, 6)
>>> vdx.DynTrajectory(model=f)(jnp.ones(3), ts).shape
(6, 3)

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

class Lorenz63(eqx.Module):
    r"""Lorenz-63 vector field as an Equinox module.

    Implements the system

    $$
    \dot x = \sigma (y - x), \quad
    \dot y = x (\rho - z) - y, \quad
    \dot z = x y - \beta z .
    $$

    The module is a diffrax-compatible right-hand side ``f(t, y, args)``.

    Attributes:
        sigma: Prandtl number $\sigma$ (classically 10).
        rho: Rayleigh number $\rho$ (classically 28).
        beta: Geometric factor $\beta$ (classically 8/3).

    Examples:
        The vector field at $(1, 1, 1)$ for the classical parameters:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> f = vdx.Lorenz63(sigma=10.0, rho=28.0, beta=8.0 / 3.0)
        >>> dydt = f(0.0, jnp.ones(3), None)
        >>> (
        ...     dydt.shape,
        ...     bool(jnp.allclose(dydt, jnp.array([0.0, 26.0, 1.0 - 8.0 / 3.0]))),
        ... )
        ((3,), True)

        The origin is a fixed point:

        >>> bool(jnp.all(f(0.0, jnp.zeros(3), None) == 0))
        True

        It drives a [`DynTrajectory`][vardax.DynTrajectory] prior:

        >>> ts = jnp.linspace(0.0, 0.05, 6)
        >>> vdx.DynTrajectory(model=f)(jnp.ones(3), ts).shape
        (6, 3)
    """

    sigma: float
    rho: float
    beta: float

    def __call__(
        self,
        t: float | Float[Array, ""],
        y: Float[Array, " 3"],
        args: None,
    ) -> Float[Array, " 3"]:
        r"""Evaluate the vector field at state ``y`` and time ``t``.

        Args:
            t: Time (unused; the system is autonomous).
            y: State $(x, y, z)$.
            args: Unused, present for the diffrax signature.

        Returns:
            The time derivative $(\dot x, \dot y, \dot z)$.
        """
        state_x, state_y, state_z = y[0], y[1], y[2]
        # dx/dt = σ(y − x)
        dx_dt = self.sigma * (state_y - state_x)
        # dy/dt = x(ρ − z) − y
        dy_dt = state_x * (self.rho - state_z) - state_y
        # dz/dt = xy − βz
        dz_dt = state_x * state_y - self.beta * state_z
        return jnp.array([dx_dt, dy_dt, dz_dt])

    def get_config(self) -> dict[str, Any]:
        return {
            # float(): the parameters are leaves and may be 0-d arrays
            "sigma": float(self.sigma),
            "rho": float(self.rho),
            "beta": float(self.beta),
        }

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

Lorenz96 ¶

Bases: Module

Lorenz-96 vector field as an Equinox module.

Implements the system (with periodic boundary conditions)

\[ \dot x_k = (x_{k+1} - x_{k-2})\, x_{k-1} - x_k + F . \]

The module is a diffrax-compatible right-hand side f(t, y, args).

Attributes:

Name	Type	Description
`forcing`	`float`	Forcing constant \(F\) (classically 8).

Examples:

The uniform state \(x_k = F\) is a fixed point, and at the origin the field equals the forcing:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> f = vdx.Lorenz96(forcing=8.0)
>>> bool(jnp.allclose(f(0.0, jnp.full(6, 8.0), None), 0.0))
True
>>> bool(jnp.allclose(f(0.0, jnp.zeros(6), None), 8.0))
True

It maps an \(N\)-vector to an \(N\)-vector, for any \(N\):

>>> f(0.0, jnp.arange(5.0), None).shape
(5,)

It drives a DynTrajectory prior:

>>> ts = jnp.linspace(0.0, 0.05, 6)
>>> vdx.DynTrajectory(model=f)(jnp.arange(5.0), ts).shape
(6, 5)

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

class Lorenz96(eqx.Module):
    r"""Lorenz-96 vector field as an Equinox module.

    Implements the system (with periodic boundary conditions)

    $$
    \dot x_k = (x_{k+1} - x_{k-2})\, x_{k-1} - x_k + F .
    $$

    The module is a diffrax-compatible right-hand side ``f(t, y, args)``.

    Attributes:
        forcing: Forcing constant $F$ (classically 8).

    Examples:
        The uniform state $x_k = F$ is a fixed point, and at the origin the
        field equals the forcing:

        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> f = vdx.Lorenz96(forcing=8.0)
        >>> bool(jnp.allclose(f(0.0, jnp.full(6, 8.0), None), 0.0))
        True
        >>> bool(jnp.allclose(f(0.0, jnp.zeros(6), None), 8.0))
        True

        It maps an $N$-vector to an $N$-vector, for any $N$:

        >>> f(0.0, jnp.arange(5.0), None).shape
        (5,)

        It drives a [`DynTrajectory`][vardax.DynTrajectory] prior:

        >>> ts = jnp.linspace(0.0, 0.05, 6)
        >>> vdx.DynTrajectory(model=f)(jnp.arange(5.0), ts).shape
        (6, 5)
    """

    forcing: float

    def __call__(
        self,
        t: float | Float[Array, ""],
        y: Float[Array, " N"],
        args: None,
    ) -> Float[Array, " N"]:
        r"""Evaluate the vector field at state ``y`` and time ``t``.

        Args:
            t: Time (unused; the system is autonomous).
            y: State $x$ on the periodic ring.
            args: Unused, present for the diffrax signature.

        Returns:
            The time derivative $\dot x$.
        """
        # Periodic shifts using jnp.roll
        y_km2 = jnp.roll(y, 2)  # x_{k-2}
        y_km1 = jnp.roll(y, 1)  # x_{k-1}
        y_kp1 = jnp.roll(y, -1)  # x_{k+1}
        return (y_kp1 - y_km2) * y_km1 - y + self.forcing

    def get_config(self) -> dict[str, Any]:
        return {
            # float(): the forcing is a leaf and may be a 0-d array
            "forcing": float(self.forcing),
        }

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

simulate_lorenz63 ¶

simulate_lorenz63(
    key: Array,
    *,
    sigma: float = 10.0,
    rho: float = 28.0,
    beta: float = 8.0 / 3.0,
    dt: float = 0.01,
    n_steps: int = 5000,
    n_burn_in: int = 1000,
    x0: Float[Array, " 3"] | None = None,
) -> tuple[Float[Array, " T"], Float[Array, "T 3"]]

Simulate the Lorenz-63 system and return state trajectory.

Integrates \(\dot x = \sigma (y - x)\), \(\dot y = x (\rho - z) - y\), \(\dot z = x y - \beta z\) with diffrax and discards a burn-in so the trajectory lies on the attractor.

Algorithm

Build the initial condition, \((1, 1, 1)\) perturbed by key, or x0.
Integrate with Tsit5 over n_burn_in + n_steps steps of size dt, saving every step.
Drop the first n_burn_in states and shift time to start at 0.

Parameters:

Name	Type	Description	Default
`key`	`Array`	JAX PRNG key used to perturb the initial condition when `x0` is `None`.	required
`sigma`	`float`	Prandtl number \(\sigma\).	`10.0`
`rho`	`float`	Rayleigh number \(\rho\).	`28.0`
`beta`	`float`	Geometric factor \(\beta\).	`8.0 / 3.0`
`dt`	`float`	Integration time step.	`0.01`
`n_steps`	`int`	Total number of integration steps after burn-in.	`5000`
`n_burn_in`	`int`	Number of initial steps to discard (burn-in).	`1000`
`x0`	`Float[Array, ' 3'] \| None`	Optional explicit initial condition. When `None`, the classic off-attractor point `(1, 1, 1)` is used, perturbed by `key`. From there the transient onto the attractor takes a couple of time units, so the default burn-in lands on it. (Starting at the unstable fixed point instead, as older versions did, leaves the trajectory within a fraction of a unit of it for ~40 time units, far longer than any burn-in the tutorials use.)	`None`

Returns:

Type	Description
`Float[Array, ' T']`	`(time_coords, states)`: the time coordinates (starting at 0) and the
`Float[Array, 'T 3']`	Lorenz-63 state trajectory, with `n_steps + 1` rows.

Examples:

A short run without burn-in:

>>> import jax
>>> import vardax as vdx
>>> t, states = vdx.simulate_lorenz63(
...     jax.random.PRNGKey(0), n_steps=10, n_burn_in=0
... )
>>> t.shape, states.shape
((11,), (11, 3))
>>> float(t[0])
0.0

An explicit initial condition is used as given:

>>> import jax.numpy as jnp
>>> x0 = jnp.array([1.0, 2.0, 3.0])
>>> _, states = vdx.simulate_lorenz63(
...     jax.random.PRNGKey(0), n_steps=5, n_burn_in=0, x0=x0
... )
>>> bool(jnp.allclose(states[0], x0))
True

A burn-in is dropped from the output but the time axis restarts at 0:

>>> t, states = vdx.simulate_lorenz63(
...     jax.random.PRNGKey(0), n_steps=5, n_burn_in=5
... )
>>> states.shape, float(t[0])
((6, 3), 0.0)

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

def simulate_lorenz63(
    key: Array,
    *,
    sigma: float = 10.0,
    rho: float = 28.0,
    beta: float = 8.0 / 3.0,
    dt: float = 0.01,
    n_steps: int = 5000,
    n_burn_in: int = 1000,
    x0: Float[Array, " 3"] | None = None,
) -> tuple[Float[Array, " T"], Float[Array, "T 3"]]:
    r"""Simulate the Lorenz-63 system and return state trajectory.

    Integrates $\dot x = \sigma (y - x)$, $\dot y = x (\rho - z) - y$,
    $\dot z = x y - \beta z$ with diffrax and discards a burn-in so the
    trajectory lies on the attractor.

    Algorithm:
        1. Build the initial condition, $(1, 1, 1)$ perturbed by ``key``, or
           ``x0``.
        2. Integrate with ``Tsit5`` over ``n_burn_in + n_steps`` steps of size
           ``dt``, saving every step.
        3. Drop the first ``n_burn_in`` states and shift time to start at 0.

    Args:
        key: JAX PRNG key used to perturb the initial condition when
            ``x0`` is ``None``.
        sigma: Prandtl number $\sigma$.
        rho: Rayleigh number $\rho$.
        beta: Geometric factor $\beta$.
        dt: Integration time step.
        n_steps: Total number of integration steps *after* burn-in.
        n_burn_in: Number of initial steps to discard (burn-in).
        x0: Optional explicit initial condition.  When
            ``None``, the classic off-attractor point ``(1, 1, 1)`` is used,
            perturbed by ``key``. From there the transient onto the
            attractor takes a couple of time units, so the default burn-in
            lands on it. (Starting at the unstable fixed point instead, as
            older versions did, leaves the trajectory within a fraction of
            a unit of it for ~40 time units, far longer than any burn-in
            the tutorials use.)

    Returns:
        ``(time_coords, states)``: the time coordinates (starting at 0) and the
        Lorenz-63 state trajectory, with ``n_steps + 1`` rows.

    Examples:
        A short run without burn-in:

        >>> import jax
        >>> import vardax as vdx
        >>> t, states = vdx.simulate_lorenz63(
        ...     jax.random.PRNGKey(0), n_steps=10, n_burn_in=0
        ... )
        >>> t.shape, states.shape
        ((11,), (11, 3))
        >>> float(t[0])
        0.0

        An explicit initial condition is used as given:

        >>> import jax.numpy as jnp
        >>> x0 = jnp.array([1.0, 2.0, 3.0])
        >>> _, states = vdx.simulate_lorenz63(
        ...     jax.random.PRNGKey(0), n_steps=5, n_burn_in=0, x0=x0
        ... )
        >>> bool(jnp.allclose(states[0], x0))
        True

        A burn-in is dropped from the output but the time axis restarts at 0:

        >>> t, states = vdx.simulate_lorenz63(
        ...     jax.random.PRNGKey(0), n_steps=5, n_burn_in=5
        ... )
        >>> states.shape, float(t[0])
        ((6, 3), 0.0)
    """
    model = Lorenz63(sigma=sigma, rho=rho, beta=beta)

    if x0 is None:
        noise = jax.random.normal(key, shape=(3,)) * 0.01
        x0 = jnp.array([1.0, 1.0, 1.0]) + noise

    total_steps = n_burn_in + n_steps
    start_time = 0.0
    end_time = total_steps * dt

    save_times = jnp.linspace(start_time, end_time, total_steps + 1)

    solution = diffeqsolve(
        ODETerm(model),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
        Tsit5(),
        t0=start_time,
        t1=end_time,
        dt0=dt,
        y0=x0,
        saveat=SaveAt(ts=save_times),
        # Fixed-step integration takes exactly ``total_steps`` steps; diffrax's
        # default cap (4096) would abort long simulations (e.g. 5000 + 1000).
        max_steps=total_steps + 1,
    )

    # Discard burn-in steps; keep indices n_burn_in .. total_steps (inclusive)
    states = solution.ys[n_burn_in:]
    time_coords = save_times[n_burn_in:] - save_times[n_burn_in]
    return time_coords, states

simulate_lorenz96 ¶

simulate_lorenz96(
    key: Array,
    *,
    n_sites: int = 40,
    forcing: float = 8.0,
    dt: float = 0.01,
    n_steps: int = 5000,
    n_burn_in: int = 1000,
) -> tuple[Float[Array, " T"], Float[Array, "T N"]]

Simulate the Lorenz-96 system and return state trajectory.

Integrates \(\dot x_k = (x_{k+1} - x_{k-2})\, x_{k-1} - x_k + F\) on a periodic ring with diffrax and discards a burn-in.

Algorithm

Start from the uniform state \(x_k = F\) plus a small perturbation drawn from key.
Integrate with Tsit5 over n_burn_in + n_steps steps of size dt, saving every step.
Drop the first n_burn_in states and shift time to start at 0.

Parameters:

Name	Type	Description	Default
`key`	`Array`	JAX PRNG key used to perturb the initial condition.	required
`n_sites`	`int`	Number of variables (spatial dimension).	`40`
`forcing`	`float`	Forcing constant.	`8.0`
`dt`	`float`	Integration time step.	`0.01`
`n_steps`	`int`	Total number of integration steps after burn-in.	`5000`
`n_burn_in`	`int`	Number of initial steps to discard (burn-in).	`1000`

Returns:

Type	Description
`Float[Array, ' T']`	`(time_coords, states)`: the time coordinates (starting at 0) and the
`Float[Array, 'T N']`	state trajectory, with `n_steps + 1` rows.

Examples:

>>> import jax
>>> import vardax as vdx
>>> t, states = vdx.simulate_lorenz96(
...     jax.random.PRNGKey(0), n_sites=8, n_steps=10, n_burn_in=0
... )
>>> t.shape, states.shape
((11,), (11, 8))

The perturbation of the initial state depends on the key:

>>> import jax.numpy as jnp
>>> _, other = vdx.simulate_lorenz96(
...     jax.random.PRNGKey(1), n_sites=8, n_steps=10, n_burn_in=0
... )
>>> bool(jnp.allclose(states[0], other[0]))
False

The result feeds a plotting helper directly:

>>> import matplotlib
>>> matplotlib.use("Agg")
>>> fig, ax = vdx.plot_l96_grid(states, t)
>>> type(ax).__name__
'Axes'

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

def simulate_lorenz96(
    key: Array,
    *,
    n_sites: int = 40,
    forcing: float = 8.0,
    dt: float = 0.01,
    n_steps: int = 5000,
    n_burn_in: int = 1000,
) -> tuple[Float[Array, " T"], Float[Array, "T N"]]:
    r"""Simulate the Lorenz-96 system and return state trajectory.

    Integrates $\dot x_k = (x_{k+1} - x_{k-2})\, x_{k-1} - x_k + F$ on a periodic
    ring with diffrax and discards a burn-in.

    Algorithm:
        1. Start from the uniform state $x_k = F$ plus a small perturbation
           drawn from ``key``.
        2. Integrate with ``Tsit5`` over ``n_burn_in + n_steps`` steps of size
           ``dt``, saving every step.
        3. Drop the first ``n_burn_in`` states and shift time to start at 0.

    Args:
        key: JAX PRNG key used to perturb the initial condition.
        n_sites: Number of variables (spatial dimension).
        forcing: Forcing constant.
        dt: Integration time step.
        n_steps: Total number of integration steps *after* burn-in.
        n_burn_in: Number of initial steps to discard (burn-in).

    Returns:
        ``(time_coords, states)``: the time coordinates (starting at 0) and the
        state trajectory, with ``n_steps + 1`` rows.

    Examples:
        >>> import jax
        >>> import vardax as vdx
        >>> t, states = vdx.simulate_lorenz96(
        ...     jax.random.PRNGKey(0), n_sites=8, n_steps=10, n_burn_in=0
        ... )
        >>> t.shape, states.shape
        ((11,), (11, 8))

        The perturbation of the initial state depends on the key:

        >>> import jax.numpy as jnp
        >>> _, other = vdx.simulate_lorenz96(
        ...     jax.random.PRNGKey(1), n_sites=8, n_steps=10, n_burn_in=0
        ... )
        >>> bool(jnp.allclose(states[0], other[0]))
        False

        The result feeds a plotting helper directly:

        >>> import matplotlib
        >>> matplotlib.use("Agg")
        >>> fig, ax = vdx.plot_l96_grid(states, t)
        >>> type(ax).__name__
        'Axes'
    """
    model = Lorenz96(forcing=forcing)

    # Standard L96 initialization: uniform forcing with a small random perturbation
    initial_state = jnp.full((n_sites,), forcing)
    noise = jax.random.normal(key, shape=(n_sites,)) * 0.01
    initial_state = initial_state + noise

    total_steps = n_burn_in + n_steps
    start_time = 0.0
    end_time = total_steps * dt

    save_times = jnp.linspace(start_time, end_time, total_steps + 1)

    solution = diffeqsolve(
        ODETerm(model),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
        Tsit5(),
        t0=start_time,
        t1=end_time,
        dt0=dt,
        y0=initial_state,
        saveat=SaveAt(ts=save_times),
        # Fixed-step integration takes exactly ``total_steps`` steps; diffrax's
        # default cap (4096) would abort long simulations (e.g. 5000 + 1000).
        max_steps=total_steps + 1,
    )

    states = solution.ys[n_burn_in:]
    time_coords = save_times[n_burn_in:] - save_times[n_burn_in]
    return time_coords, states

time_patches ¶

time_patches(
    ts: Float[Array, " T"],
) -> Float[Array, "T-1 2"]

Overlapping consecutive time pairs for one-step increment losses.

Reimplements the legacy kernex.kmap(kernel_size=(2,), relative=True) sliding window without the kernex dependency: for a 1-D array of times [t_0, t_1, ..., t_{T-1}] it returns the T - 1 overlapping pairs [[t_0, t_1], [t_1, t_2], ..., [t_{T-2}, t_{T-1}]].

Parameters:

Name	Type	Description	Default
`ts`	`Float[Array, ' T']`	Monotonic time coordinates \(t_0, \dots, t_{T-1}\).	required

Returns:

Type	Description
`Float[Array, 'T-1 2']`	The consecutive pairs \((t_i, t_{i+1})\), one row per pair.

Examples:

>>> import jax.numpy as jnp
>>> from vardax import time_patches
>>> time_patches(jnp.arange(4.0))
Array([[0., 1.],
       [1., 2.],
       [2., 3.]], dtype=float32)

A single time yields no pairs:

>>> time_patches(jnp.array([0.0])).shape
(0, 2)

The pairs drive the one-step solves of DynIncrements:

>>> time_patches(jnp.linspace(0.0, 1.0, 5)).shape
(4, 2)

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

def time_patches(ts: Float[Array, " T"]) -> Float[Array, "T-1 2"]:
    r"""Overlapping consecutive time pairs for one-step increment losses.

    Reimplements the legacy ``kernex.kmap(kernel_size=(2,), relative=True)``
    sliding window without the ``kernex`` dependency: for a 1-D array of
    times ``[t_0, t_1, ..., t_{T-1}]`` it returns the ``T - 1`` overlapping
    pairs ``[[t_0, t_1], [t_1, t_2], ..., [t_{T-2}, t_{T-1}]]``.

    Args:
        ts: Monotonic time coordinates $t_0, \dots, t_{T-1}$.

    Returns:
        The consecutive pairs $(t_i, t_{i+1})$, one row per pair.

    Examples:
        >>> import jax.numpy as jnp
        >>> from vardax import time_patches
        >>> time_patches(jnp.arange(4.0))
        Array([[0., 1.],
               [1., 2.],
               [2., 3.]], dtype=float32)

        A single time yields no pairs:

        >>> time_patches(jnp.array([0.0])).shape
        (0, 2)

        The pairs drive the one-step solves of
        [`DynIncrements`][vardax.DynIncrements]:

        >>> time_patches(jnp.linspace(0.0, 1.0, 5)).shape
        (4, 2)
    """
    # jnp, not einx: a single time gives zero pairs, which einx cannot size
    return jnp.stack([ts[:-1], ts[1:]], axis=-1)

Forward-model adapters¶

SoftBoundedForward wraps any pipekit_cycle.ForwardModel so that the line search of a variational solver cannot drive an explicit ODE step off the attractor into overflow: inside a box the adapter is the identity, outside it saturates smoothly. Notebooks 13 and 14 use it around Lorenz-96 and Lorenz-63 forwards.

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

SoftBoundedForward ¶

Bases: Module

ForwardModel adapter that softly confines states to a box.

Explicit Runge–Kutta steps of dynamics with quadratic terms (Lorenz-63, Lorenz-96, …) overflow, or drive an adaptive step-size controller past its step cap, when the state is far outside the attractor. A line-search minimiser such as the BFGS inside StrongFourDVar will try such states, and one non-finite cost is enough to poison the whole solve. This adapter maps the state through

\[ s(x) = \operatorname{sign}(x)\,\begin{cases} |x| & |x| \le b \\ b + b\,\dfrac{u}{1 + u}, \quad u = (|x| - b)/b & |x| > b \end{cases} \]

on the way into and out of every step, so that

states inside the box \(|x_i| \le b\) are untouched (choose bound so the attractor sits well inside it — the adapter is then exactly the wrapped model wherever the analysis can end up);
states outside are confined to \(|x_i| < 2b\), keeping every trial cost finite;
the map is continuously differentiable (slope one at the box edge from both sides) and its gradient outside, \(1/(1+u)^2\), stays representable in single precision for trial states up to about \(10^{19}\) times the bound, unlike a hard clip (zero gradient) or a tanh saturation (gradient underflows to zero a few bounds out); a rejected trial therefore still carries curvature information if the minimiser asks for it.

Algorithm

Confine the incoming state with \(s\).
Advance it with the wrapped model's step.
Confine the result with \(s\) again.

Attributes:

Name	Type	Description
`forward`	`Any`	The wrapped `pipekit_cycle.ForwardModel`.
`bound`	`float`	Half-width \(b\) of the box on which the adapter is the identity.

Examples:

>>> import jax.numpy as jnp
>>> import vardax
>>> class Double:
...     dt = 1.0
...     state_signature = None
...
...     def step(self, x, dt):
...         return 2.0 * x
>>> fwd = vardax.SoftBoundedForward(Double(), bound=10.0)
>>> fwd.step(jnp.array([1.0, -3.0]), 1.0)  # inside the box: untouched
Array([ 2., -6.], dtype=float32)
>>> bool(jnp.all(jnp.abs(fwd.step(jnp.array([1e6, -1e6]), 1.0)) < 20.0))
True

The step length and signature are forwarded from the wrapped model:

>>> fwd.dt, fwd.state_signature
(1.0, None)

Wrapped around a diffrax-backed model it keeps a far-off state finite:

>>> def decay(t, y, args):
...     return -y
>>> inner = vardax.DynTrajectory(model=decay).as_forward_model(0.1)
>>> safe = vardax.SoftBoundedForward(inner, bound=5.0)
>>> bool(jnp.all(jnp.isfinite(safe.step(jnp.array([1e9]), 0.1))))
True

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

class SoftBoundedForward(eqx.Module):
    r"""``ForwardModel`` adapter that softly confines states to a box.

    Explicit Runge–Kutta steps of dynamics with quadratic terms (Lorenz-63,
    Lorenz-96, …) overflow, or drive an adaptive step-size controller past
    its step cap, when the state is far outside the attractor. A
    line-search minimiser such as the BFGS inside
    [`StrongFourDVar`][vardax.StrongFourDVar] will try such states, and one
    non-finite cost is enough to poison the whole solve. This adapter maps
    the state through

    $$
    s(x) = \operatorname{sign}(x)\,\begin{cases}
      |x| & |x| \le b \\
      b + b\,\dfrac{u}{1 + u}, \quad u = (|x| - b)/b & |x| > b
    \end{cases}
    $$

    on the way into and out of every step, so that

    - states inside the box $|x_i| \le b$ are untouched (choose ``bound``
      so the attractor sits well inside it — the adapter is then exactly
      the wrapped model wherever the analysis can end up);
    - states outside are confined to $|x_i| < 2b$, keeping every trial
      cost finite;
    - the map is continuously differentiable (slope one at the box edge
      from both sides) and its gradient outside, $1/(1+u)^2$, stays
      representable in single precision for trial states up to about
      $10^{19}$ times the bound, unlike a hard clip (zero gradient) or a
      ``tanh`` saturation (gradient underflows to zero a few bounds out);
      a rejected trial therefore still carries curvature information if
      the minimiser asks for it.

    Algorithm:
        1. Confine the incoming state with $s$.
        2. Advance it with the wrapped model's ``step``.
        3. Confine the result with $s$ again.

    Attributes:
        forward: The wrapped ``pipekit_cycle.ForwardModel``.
        bound: Half-width $b$ of the box on which the adapter is the
            identity.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax
        >>> class Double:
        ...     dt = 1.0
        ...     state_signature = None
        ...
        ...     def step(self, x, dt):
        ...         return 2.0 * x
        >>> fwd = vardax.SoftBoundedForward(Double(), bound=10.0)
        >>> fwd.step(jnp.array([1.0, -3.0]), 1.0)  # inside the box: untouched
        Array([ 2., -6.], dtype=float32)
        >>> bool(jnp.all(jnp.abs(fwd.step(jnp.array([1e6, -1e6]), 1.0)) < 20.0))
        True

        The step length and signature are forwarded from the wrapped model:

        >>> fwd.dt, fwd.state_signature
        (1.0, None)

        Wrapped around a diffrax-backed model it keeps a far-off state finite:

        >>> def decay(t, y, args):
        ...     return -y
        >>> inner = vardax.DynTrajectory(model=decay).as_forward_model(0.1)
        >>> safe = vardax.SoftBoundedForward(inner, bound=5.0)
        >>> bool(jnp.all(jnp.isfinite(safe.step(jnp.array([1e9]), 0.1))))
        True
    """

    forward: Any
    bound: float = eqx.field(static=True)

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

    @property
    def state_signature(self) -> Any:
        return getattr(self.forward, "state_signature", None)

    def _confine(self, state: Array) -> Array:
        half_width = self.bound
        abs_state = jnp.abs(state)
        # The in-box branch returns ``state`` itself (not ``sign * abs``) so the
        # gradient is exactly one everywhere inside, including at zero, where
        # ``jnp.sign`` would otherwise zero it. Outside, the algebraic
        # saturation u / (1 + u) keeps the derivative 1 / (1 + u)^2 finite and
        # non-zero far beyond any state a line search will try (tanh would
        # underflow to a zero gradient a few bounds out).
        # u = max(|x| − b, 0) / b
        excess = jnp.maximum(abs_state - half_width, 0.0) / half_width
        # s(x) = sign(x) (b + b u / (1 + u))
        outside = jnp.sign(state) * (half_width + half_width * excess / (1.0 + excess))
        return jnp.where(abs_state <= half_width, state, outside)

    def step(self, state: Array, dt: float) -> Array:
        r"""Advance ``state`` by ``dt`` with confinement before and after.

        Args:
            state: State $x$ to advance.
            dt: Step length $\Delta t$.

        Returns:
            $s(\Phi_{\Delta t}(s(x)))$, with $\Phi$ the wrapped model's step.
        """
        return self._confine(self.forward.step(self._confine(state), dt))

step ¶

step(state: Array, dt: float) -> Array

Advance state by dt with confinement before and after.

Parameters:

Name	Type	Description	Default
`state`	`Array`	State \(x\) to advance.	required
`dt`	`float`	Step length \(\Delta t\).	required

Returns:

Type	Description
`Array`	\(s(\Phi_{\Delta t}(s(x)))\), with \(\Phi\) the wrapped model's step.

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

def step(self, state: Array, dt: float) -> Array:
    r"""Advance ``state`` by ``dt`` with confinement before and after.

    Args:
        state: State $x$ to advance.
        dt: Step length $\Delta t$.

    Returns:
        $s(\Phi_{\Delta t}(s(x)))$, with $\Phi$ the wrapped model's step.
    """
    return self._confine(self.forward.step(self._confine(state), dt))

Validation gates¶

The six-step methodology's go/no-go checks (Decision D12): simulation-based calibration ranks the truth within posterior samples and must be uniform; assert_posterior_agreement cross-checks two posterior adapters against each other; and assert_adjoint_calibrated verifies that a cheap adjoint tracks the exact gradient before it is used for training. Run these before believing any uncertainty estimate.

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

simulation_based_calibration ¶

simulation_based_calibration(
    sample_posterior: Callable[
        [Array, PRNGKeyArray, int], Array
    ],
    sample_prior: Callable[[PRNGKeyArray], Array],
    simulate_obs: Callable[[Array, PRNGKeyArray], Array],
    *,
    key: PRNGKeyArray,
    n_runs: int = 200,
    n_samples: int = 200,
) -> Int[Array, " n_runs"]

Per Talts et al. 2018: rank histogram of true draws within posterior samples.

Procedure for each of n_runs independent draws:

Sample \(x^{(j)} \sim p(x)\) from the prior.
Simulate \(y^{(j)} = \mathrm{simulate\_obs}(x^{(j)})\).
Draw n_samples posterior samples from \(q_\phi(\cdot \mid y^{(j)})\).
Compute the rank of (a flattened scalar reduction of) \(x^{(j)}\) in the sample set.

A well-calibrated posterior produces uniformly distributed ranks over [0, n_samples]. Bumps near the edges → over-confident; centre-mass bump → under-confident.

Algorithm

For each run, split the key into prior, observation and posterior keys.
Draw \(x\), simulate \(y\) and draw the posterior samples.
Reduce \(x\) and every sample to its \(\ell_2\) norm.
Record how many sample norms fall below the norm of \(x\).

References

Talts, Betancourt, Simpson, Vehtari & Gelman (2018). "Validating Bayesian inference algorithms with simulation-based calibration."

Parameters:

Name	Type	Description	Default
`sample_posterior`	`Callable[[Array, PRNGKeyArray, int], Array]`	`(y, key, n) -> Array` of posterior samples with `n` leading rows, conditioned on `y`. For an `AmortizedPosterior`, the natural wrapper is `lambda y, k, n: model.sample(Batch(input=y[None], ...), k, n)[0]`.	required
`sample_prior`	`Callable[[PRNGKeyArray], Array]`	`key -> x`. Single prior draw.	required
`simulate_obs`	`Callable[[Array, PRNGKeyArray], Array]`	`(x, key) -> y`. Forward + noise.	required
`key`	`PRNGKeyArray`	Top-level PRNG key.	required
`n_runs`	`int`	Number of (prior, obs, posterior) triples.	`200`
`n_samples`	`int`	Posterior samples per run (defines the rank histogram resolution).	`200`

Returns:

Type	Description
`Int[Array, ' n_runs']`	Integer ranks in `[0, n_samples]`, one per run.

Examples:

Ranks of the truth among samples drawn around the observation lie in [0, n_samples]:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> prior = lambda k: jax.random.normal(k, (2,))
>>> observe = lambda x, k: x + 0.1 * jax.random.normal(k, (2,))
>>> posterior = lambda y, k, n: y + 0.1 * jax.random.normal(k, (n, 2))
>>> ranks = vdx.simulation_based_calibration(
...     posterior,
...     prior,
...     observe,
...     key=jax.random.PRNGKey(0),
...     n_runs=6,
...     n_samples=10,
... )
>>> ranks.shape, bool(jnp.all((ranks >= 0) & (ranks <= 10)))
((6,), True)

A posterior that ignores the data puts the truth below every sample, so every rank is zero:

>>> wrong = lambda y, k, n: 100.0 + jax.random.normal(k, (n, 2))
>>> bad = vdx.simulation_based_calibration(
...     wrong,
...     prior,
...     observe,
...     key=jax.random.PRNGKey(0),
...     n_runs=6,
...     n_samples=10,
... )
>>> bool(jnp.all(bad == 0))
True

An AmortizedPosterior plugs in through a one-element batch:

>>> key = jax.random.PRNGKey(1)
>>> amortized = vdx.AmortizedPosterior(
...     encoder=vdx.MLPObsEncoder(input_size=6, context_dim=4, key=key),
...     head=vdx.RegressionHead(context_dim=4, state_shape=(2, 3), key=key),
...     config=vdx.AmortizedConfig(),
... )
>>> def sample_posterior(y, k, n):
...     batch = vdx.Batch1D(input=y[None], mask=jnp.ones_like(y)[None])
...     return amortized.sample(batch, k, n)[0]
>>> ranks = vdx.simulation_based_calibration(
...     sample_posterior,
...     lambda k: jax.random.normal(k, (2, 3)),
...     lambda x, k: x,
...     key=jax.random.PRNGKey(2),
...     n_runs=4,
...     n_samples=5,
... )
>>> ranks.shape
(4,)

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

def simulation_based_calibration(
    sample_posterior: Callable[[Array, PRNGKeyArray, int], Array],
    sample_prior: Callable[[PRNGKeyArray], Array],
    simulate_obs: Callable[[Array, PRNGKeyArray], Array],
    *,
    key: PRNGKeyArray,
    n_runs: int = 200,
    n_samples: int = 200,
) -> Int[Array, " n_runs"]:
    r"""Per Talts et al. 2018: rank histogram of true draws within
    posterior samples.

    Procedure for each of ``n_runs`` independent draws:

    1. Sample $x^{(j)} \sim p(x)$ from the prior.
    2. Simulate $y^{(j)} = \mathrm{simulate\_obs}(x^{(j)})$.
    3. Draw ``n_samples`` posterior samples from
       $q_\phi(\cdot \mid y^{(j)})$.
    4. Compute the rank of (a flattened scalar reduction of)
       $x^{(j)}$ in the sample set.

    A well-calibrated posterior produces uniformly distributed ranks
    over ``[0, n_samples]``. Bumps near the edges → over-confident;
    centre-mass bump → under-confident.

    Algorithm:
        1. For each run, split the key into prior, observation and posterior keys.
        2. Draw $x$, simulate $y$ and draw the posterior samples.
        3. Reduce $x$ and every sample to its $\ell_2$ norm.
        4. Record how many sample norms fall below the norm of $x$.

    References:
        Talts, Betancourt, Simpson, Vehtari & Gelman (2018). "Validating
        Bayesian inference algorithms with simulation-based calibration."

    Args:
        sample_posterior: ``(y, key, n) -> Array`` of posterior samples with
            ``n`` leading rows, conditioned on ``y``. For an
            ``AmortizedPosterior``, the natural wrapper is
            ``lambda y, k, n: model.sample(Batch(input=y[None], ...), k, n)[0]``.
        sample_prior: ``key -> x``. Single prior draw.
        simulate_obs: ``(x, key) -> y``. Forward + noise.
        key: Top-level PRNG key.
        n_runs: Number of (prior, obs, posterior) triples.
        n_samples: Posterior samples per run (defines the rank
            histogram resolution).

    Returns:
        Integer ranks in ``[0, n_samples]``, one per run.

    Examples:
        Ranks of the truth among samples drawn around the observation lie in
        ``[0, n_samples]``:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> prior = lambda k: jax.random.normal(k, (2,))
        >>> observe = lambda x, k: x + 0.1 * jax.random.normal(k, (2,))
        >>> posterior = lambda y, k, n: y + 0.1 * jax.random.normal(k, (n, 2))
        >>> ranks = vdx.simulation_based_calibration(
        ...     posterior,
        ...     prior,
        ...     observe,
        ...     key=jax.random.PRNGKey(0),
        ...     n_runs=6,
        ...     n_samples=10,
        ... )
        >>> ranks.shape, bool(jnp.all((ranks >= 0) & (ranks <= 10)))
        ((6,), True)

        A posterior that ignores the data puts the truth below every sample, so
        every rank is zero:

        >>> wrong = lambda y, k, n: 100.0 + jax.random.normal(k, (n, 2))
        >>> bad = vdx.simulation_based_calibration(
        ...     wrong,
        ...     prior,
        ...     observe,
        ...     key=jax.random.PRNGKey(0),
        ...     n_runs=6,
        ...     n_samples=10,
        ... )
        >>> bool(jnp.all(bad == 0))
        True

        An [`AmortizedPosterior`][vardax.AmortizedPosterior] plugs in through a
        one-element batch:

        >>> key = jax.random.PRNGKey(1)
        >>> amortized = vdx.AmortizedPosterior(
        ...     encoder=vdx.MLPObsEncoder(input_size=6, context_dim=4, key=key),
        ...     head=vdx.RegressionHead(context_dim=4, state_shape=(2, 3), key=key),
        ...     config=vdx.AmortizedConfig(),
        ... )
        >>> def sample_posterior(y, k, n):
        ...     batch = vdx.Batch1D(input=y[None], mask=jnp.ones_like(y)[None])
        ...     return amortized.sample(batch, k, n)[0]
        >>> ranks = vdx.simulation_based_calibration(
        ...     sample_posterior,
        ...     lambda k: jax.random.normal(k, (2, 3)),
        ...     lambda x, k: x,
        ...     key=jax.random.PRNGKey(2),
        ...     n_runs=4,
        ...     n_samples=5,
        ... )
        >>> ranks.shape
        (4,)
    """
    keys = jax.random.split(key, n_runs)
    ranks = []
    for run_key in keys:
        key_prior, key_obs, key_post = jax.random.split(run_key, 3)
        x_true = sample_prior(key_prior)
        observation = simulate_obs(x_true, key_obs)
        samples = sample_posterior(observation, key_post, n_samples)
        # Flatten and use the L2 norm as the scalar reduction.
        scalar_true = float(jnp.linalg.norm(x_true))
        # ‖xᵢ‖₂ over every state dimension of each sample
        squared_norms = einx.sum("n [s...]", samples**2)
        scalar_samples = jnp.sqrt(squared_norms)
        rank = int(jnp.sum(scalar_samples < scalar_true))
        ranks.append(rank)
    return jnp.asarray(ranks, dtype=jnp.int32)

assert_posterior_agreement ¶

assert_posterior_agreement(
    p_fast: Posterior,
    p_oracle: Posterior,
    *,
    tolerance_sigma: float = 1.0,
) -> None

Check that p_fast.mean lies within tolerance_sigma standard deviations of p_oracle.mean.

Marginal-only test: each component of the mean must satisfy

\[ \frac{|x^*_\text{fast} - x^*_\text{oracle}|}{\sigma_\text{post, oracle}} \le \text{tolerance\_sigma}. \]

The oracle marginal \(\sigma_i\) is extracted from p_oracle.cov by probing \(e_i^T \Sigma e_i\) via one matvec per component. Cheap for moderate state size; for very large state sizes use Hutchinson estimation upstream and supply the diagonal directly via Posterior.cov materialised as a diagonal operator.

Algorithm

Extract the oracle marginal standard deviations \(\sigma_i\) from p_oracle.cov.
Form the mean difference \(\Delta x^* = x^*_\text{fast} - x^*_\text{oracle}\).
Compute the z-scores \(z_i = |\Delta x^*_i| / \sigma_i\).
Fail if any \(z_i\) exceeds tolerance_sigma.

Parameters:

Name	Type	Description	Default
`p_fast`	`Posterior`	Posterior produced by the amortized / fast model.	required
`p_oracle`	`Posterior`	Posterior produced by the oracle (e.g. `StrongFourDVar` + `LaplaceCovariance`).	required
`tolerance_sigma`	`float`	Allowed deviation in units of \(\sigma\).	`1.0`

Raises:

Type	Description
`MissingInputError`	If `p_oracle.cov` is `None`.
`ShapeError`	If the two means have different shapes.
`AssertionError`	If any component exceeds the tolerance.

Examples:

A fast mean within one oracle standard deviation passes silently:

>>> import jax
>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import vardax as vdx
>>> cov = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
>>> oracle = vdx.Posterior(mean=jnp.zeros(3), cov=cov)
>>> fast = vdx.Posterior(mean=jnp.array([0.5, -0.5, 0.0]))
>>> vdx.assert_posterior_agreement(fast, oracle, tolerance_sigma=1.0)

A deviation of three standard deviations fails:

>>> far = vdx.Posterior(mean=jnp.array([3.0, 0.0, 0.0]))
>>> vdx.assert_posterior_agreement(
...     far, oracle
... )
Traceback (most recent call last):
...
AssertionError: posterior agreement failed: max z = 3.000 > 1.000.

An oracle without a covariance cannot supply the scale:

>>> vdx.assert_posterior_agreement(
...     fast, fast
... )
Traceback (most recent call last):
...
MissingInputError: assert_posterior_agreement requires p_oracle.cov ...

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

def assert_posterior_agreement(
    p_fast: Posterior,
    p_oracle: Posterior,
    *,
    tolerance_sigma: float = 1.0,
) -> None:
    r"""Check that ``p_fast.mean`` lies within ``tolerance_sigma`` standard
    deviations of ``p_oracle.mean``.

    Marginal-only test: each component of the mean must satisfy

    $$
    \frac{|x^*_\text{fast} - x^*_\text{oracle}|}{\sigma_\text{post, oracle}}
    \le \text{tolerance\_sigma}.
    $$

    The oracle marginal $\sigma_i$ is extracted from
    ``p_oracle.cov`` by probing $e_i^T \Sigma e_i$ via one matvec per
    component. Cheap for moderate state size; for very large state
    sizes use Hutchinson estimation upstream and supply the diagonal
    directly via ``Posterior.cov`` materialised as a diagonal operator.

    Algorithm:
        1. Extract the oracle marginal standard deviations $\sigma_i$ from
           ``p_oracle.cov``.
        2. Form the mean difference $\Delta x^* = x^*_\text{fast} - x^*_\text{oracle}$.
        3. Compute the z-scores $z_i = |\Delta x^*_i| / \sigma_i$.
        4. Fail if any $z_i$ exceeds ``tolerance_sigma``.

    Args:
        p_fast: Posterior produced by the amortized / fast model.
        p_oracle: Posterior produced by the oracle (e.g. ``StrongFourDVar``
            + [`LaplaceCovariance`][vardax.LaplaceCovariance]).
        tolerance_sigma: Allowed deviation in units of $\sigma$.

    Raises:
        MissingInputError: If ``p_oracle.cov`` is ``None``.
        ShapeError: If the two means have different shapes.
        AssertionError: If any component exceeds the tolerance.

    Examples:
        A fast mean within one oracle standard deviation passes silently:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import vardax as vdx
        >>> cov = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> oracle = vdx.Posterior(mean=jnp.zeros(3), cov=cov)
        >>> fast = vdx.Posterior(mean=jnp.array([0.5, -0.5, 0.0]))
        >>> vdx.assert_posterior_agreement(fast, oracle, tolerance_sigma=1.0)

        A deviation of three standard deviations fails:

        >>> far = vdx.Posterior(mean=jnp.array([3.0, 0.0, 0.0]))
        >>> vdx.assert_posterior_agreement(
        ...     far, oracle
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        AssertionError: posterior agreement failed: max z = 3.000 > 1.000.

        An oracle without a covariance cannot supply the scale:

        >>> vdx.assert_posterior_agreement(
        ...     fast, fast
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        MissingInputError: assert_posterior_agreement requires p_oracle.cov ...
    """
    if p_oracle.cov is None:
        raise MissingInputError(
            "assert_posterior_agreement requires p_oracle.cov to be set "
            "(used to extract the marginal standard deviations)."
        )
    if p_fast.mean.shape != p_oracle.mean.shape:
        raise ShapeError(
            "p_fast.mean and p_oracle.mean must have the same shape; got "
            f"{p_fast.mean.shape} vs {p_oracle.mean.shape}."
        )
    oracle_std = _marginal_std(p_oracle.cov, p_oracle.mean)
    # z = |x*_fast − x*_oracle| / σ_oracle
    z_score = jnp.abs(p_fast.mean - p_oracle.mean) / jnp.maximum(oracle_std, 1e-12)
    if jnp.any(z_score > tolerance_sigma):
        worst = float(jnp.max(z_score))
        raise AssertionError(
            f"posterior agreement failed: max z = {worst:.3f} > {tolerance_sigma:.3f}."
        )

assert_adjoint_calibrated ¶

assert_adjoint_calibrated(
    fn_fast: Callable[[Array], Array],
    fn_oracle: Callable[[Array], Array],
    y: Float[Array, ...],
    *,
    key: PRNGKeyArray,
    threshold: float = 0.05,
    n_probes: int = 10,
) -> None

Random-vector probe of the Jacobian agreement at y.

Tests

\[ \frac{\|J_\text{fast} v - J_\text{oracle} v\|} {\|J_\text{oracle} v\|} < \text{threshold} \]

for n_probes random unit vectors \(v\). Avoids materialising either Jacobian — uses jax.jvp to apply each as needed. Cheaper than dense Jacobian comparison and is the operational test used by the six-step cycle.

Algorithm

Draw a random unit vector \(v\) of the shape of y.
Apply both Jacobians to it with jax.jvp.
Compute the relative error of the two products.
Repeat for n_probes vectors and fail if the worst error exceeds threshold.

Parameters:

Name	Type	Description	Default
`fn_fast`	`Callable[[Array], Array]`	Callable `y → analysis` (e.g. `lambda y_: amortized(Batch1D(input=y_, mask=mask, target=None))`).	required
`fn_oracle`	`Callable[[Array], Array]`	Same signature, but the oracle.	required
`y`	`Float[Array, ...]`	Observation tensor at which the Jacobians are probed.	required
`key`	`PRNGKeyArray`	PRNG key for the probe vectors.	required
`threshold`	`float`	Maximum allowed relative error.	`0.05`
`n_probes`	`int`	Number of probe vectors.	`10`

Raises:

Type	Description
`AssertionError`	If any probe exceeds the threshold.

Examples:

Two maps with the same Jacobian agree, even if their values differ:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> y = jnp.arange(4.0)
>>> vdx.assert_adjoint_calibrated(
...     lambda v: 2.0 * v + 1.0,
...     lambda v: 2.0 * v,
...     y,
...     key=jax.random.PRNGKey(0),
... )

A fast model whose sensitivity is off by a third fails:

>>> vdx.assert_adjoint_calibrated(
...     lambda v: 2.0 * v, lambda v: 3.0 * v, y, key=jax.random.PRNGKey(0)
... )
Traceback (most recent call last):
...
AssertionError: adjoint calibration failed: ...

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

def assert_adjoint_calibrated(
    fn_fast: Callable[[Array], Array],
    fn_oracle: Callable[[Array], Array],
    y: Float[Array, ...],
    *,
    key: PRNGKeyArray,
    threshold: float = 0.05,
    n_probes: int = 10,
) -> None:
    r"""Random-vector probe of the Jacobian agreement at ``y``.

    Tests

    $$
    \frac{\|J_\text{fast} v - J_\text{oracle} v\|}
         {\|J_\text{oracle} v\|} < \text{threshold}
    $$

    for ``n_probes`` random unit vectors $v$. Avoids
    materialising either Jacobian — uses ``jax.jvp`` to apply each as
    needed. Cheaper than dense Jacobian comparison and is the
    operational test used by the six-step cycle.

    Algorithm:
        1. Draw a random unit vector $v$ of the shape of ``y``.
        2. Apply both Jacobians to it with ``jax.jvp``.
        3. Compute the relative error of the two products.
        4. Repeat for ``n_probes`` vectors and fail if the worst error exceeds
           ``threshold``.

    Args:
        fn_fast: Callable ``y → analysis`` (e.g.
            ``lambda y_: amortized(Batch1D(input=y_, mask=mask, target=None))``).
        fn_oracle: Same signature, but the oracle.
        y: Observation tensor at which the Jacobians are probed.
        key: PRNG key for the probe vectors.
        threshold: Maximum allowed relative error.
        n_probes: Number of probe vectors.

    Raises:
        AssertionError: If any probe exceeds the threshold.

    Examples:
        Two maps with the same Jacobian agree, even if their values differ:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> y = jnp.arange(4.0)
        >>> vdx.assert_adjoint_calibrated(
        ...     lambda v: 2.0 * v + 1.0,
        ...     lambda v: 2.0 * v,
        ...     y,
        ...     key=jax.random.PRNGKey(0),
        ... )

        A fast model whose sensitivity is off by a third fails:

        >>> vdx.assert_adjoint_calibrated(
        ...     lambda v: 2.0 * v, lambda v: 3.0 * v, y, key=jax.random.PRNGKey(0)
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        AssertionError: adjoint calibration failed: ...
    """
    keys = jax.random.split(key, n_probes)
    worst_rel = 0.0
    for probe_key in keys:
        probe = jax.random.normal(probe_key, y.shape)
        probe = probe / jnp.linalg.norm(probe)
        _, jvp_fast = jax.jvp(fn_fast, (y,), (probe,))
        _, jvp_oracle = jax.jvp(fn_oracle, (y,), (probe,))
        denom = jnp.linalg.norm(jvp_oracle) + 1e-12
        # ‖J_fast v − J_oracle v‖ / ‖J_oracle v‖
        relative_error = float(jnp.linalg.norm(jvp_fast - jvp_oracle) / denom)
        worst_rel = max(worst_rel, relative_error)
    if worst_rel > threshold:
        raise AssertionError(
            f"adjoint calibration failed: max relative error {worst_rel:.4f} > "
            f"{threshold:.4f} (over {n_probes} probes)."
        )

Visualization¶

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

plot_l96_trajectories ¶

plot_l96_trajectories(
    states: ndarray,
    time_coords: ndarray,
    *,
    n_vars: int = 5,
    ax: Axes | None = None,
) -> tuple[Figure, Axes]

Line plot of selected Lorenz-96 variables over time.

Parameters:

Name	Type	Description	Default
`states`	`ndarray`	State trajectory, indexed `[time, variable]`.	required
`time_coords`	`ndarray`	Time coordinate of each row of `states`.	required
`n_vars`	`int`	Number of evenly-spaced variables to plot (capped at the number of variables).	`5`
`ax`	`Axes \| None`	Optional existing `Axes` to draw on.	`None`

Returns:

Type	Description
`tuple[Figure, Axes]`	`(fig, ax)`: the Matplotlib figure and axes.

Raises:

Type	Description
`ConfigurationError`	If `n_vars` is not positive.

Examples:

>>> import matplotlib
>>> matplotlib.use("Agg")
>>> import numpy as np
>>> import vardax as vdx
>>> states = np.random.default_rng(0).normal(size=(20, 8))
>>> time = np.linspace(0.0, 1.0, 20)
>>> fig, ax = vdx.plot_l96_trajectories(states, time, n_vars=3)
>>> len(ax.lines)
3

Asking for more variables than exist plots all of them:

>>> fig, ax = vdx.plot_l96_trajectories(states, time, n_vars=50)
>>> len(ax.lines)
8

A non-positive count is rejected:

>>> vdx.plot_l96_trajectories(
...     states, time, n_vars=0
... )
Traceback (most recent call last):
...
ConfigurationError: n_vars must be a positive integer, got 0.

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

def plot_l96_trajectories(
    states: np.ndarray,
    time_coords: np.ndarray,
    *,
    n_vars: int = 5,
    ax: Axes | None = None,
) -> tuple[Figure, Axes]:
    """Line plot of selected Lorenz-96 variables over time.

    Args:
        states: State trajectory, indexed ``[time, variable]``.
        time_coords: Time coordinate of each row of ``states``.
        n_vars: Number of evenly-spaced variables to plot (capped at the
            number of variables).
        ax: Optional existing ``Axes`` to draw on.

    Returns:
        ``(fig, ax)``: the Matplotlib figure and axes.

    Raises:
        ConfigurationError: If ``n_vars`` is not positive.

    Examples:
        >>> import matplotlib
        >>> matplotlib.use("Agg")
        >>> import numpy as np
        >>> import vardax as vdx
        >>> states = np.random.default_rng(0).normal(size=(20, 8))
        >>> time = np.linspace(0.0, 1.0, 20)
        >>> fig, ax = vdx.plot_l96_trajectories(states, time, n_vars=3)
        >>> len(ax.lines)
        3

        Asking for more variables than exist plots all of them:

        >>> fig, ax = vdx.plot_l96_trajectories(states, time, n_vars=50)
        >>> len(ax.lines)
        8

        A non-positive count is rejected:

        >>> vdx.plot_l96_trajectories(
        ...     states, time, n_vars=0
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: n_vars must be a positive integer, got 0.
    """
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()
        assert fig is not None

    N = states.shape[1]
    if n_vars <= 0:
        raise ConfigurationError(f"n_vars must be a positive integer, got {n_vars}.")
    n_selected = min(n_vars, N)
    indices = np.unique(np.linspace(0, N - 1, n_selected, dtype=int))
    for i in indices:
        ax.plot(time_coords, states[:, i], label=f"x{i}")
    ax.set_xlabel("time")
    ax.set_ylabel("state")
    ax.legend()
    return fig, ax  # type: ignore[return-value]  # ty:ignore[invalid-return-type]

plot_l96_grid ¶

plot_l96_grid(
    states: ndarray,
    time_coords: ndarray,
    *,
    ax: Axes | None = None,
) -> tuple[Figure, Axes]

Hovmöller-style image of Lorenz-96 states over time.

Parameters:

Name	Type	Description	Default
`states`	`ndarray`	State trajectory, indexed `[time, variable]`.	required
`time_coords`	`ndarray`	Time coordinate of each row of `states`.	required
`ax`	`Axes \| None`	Optional existing `Axes` to draw on.	`None`

Returns:

Type	Description
`tuple[Figure, Axes]`	`(fig, ax)`: the Matplotlib figure and axes.

Examples:

>>> import matplotlib
>>> matplotlib.use("Agg")
>>> import numpy as np
>>> import vardax as vdx
>>> states = np.random.default_rng(0).normal(size=(20, 8))
>>> fig, ax = vdx.plot_l96_grid(states, np.linspace(0.0, 1.0, 20))
>>> type(ax).__name__, ax.get_xlabel()
('Axes', 'time')

Drawing onto an existing axes returns that axes:

>>> import matplotlib.pyplot as plt
>>> _, mine = plt.subplots()
>>> _, same = vdx.plot_l96_grid(states, np.linspace(0.0, 1.0, 20), ax=mine)
>>> same is mine
True

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

def plot_l96_grid(
    states: np.ndarray,
    time_coords: np.ndarray,
    *,
    ax: Axes | None = None,
) -> tuple[Figure, Axes]:
    """Hovmöller-style image of Lorenz-96 states over time.

    Args:
        states: State trajectory, indexed ``[time, variable]``.
        time_coords: Time coordinate of each row of ``states``.
        ax: Optional existing ``Axes`` to draw on.

    Returns:
        ``(fig, ax)``: the Matplotlib figure and axes.

    Examples:
        >>> import matplotlib
        >>> matplotlib.use("Agg")
        >>> import numpy as np
        >>> import vardax as vdx
        >>> states = np.random.default_rng(0).normal(size=(20, 8))
        >>> fig, ax = vdx.plot_l96_grid(states, np.linspace(0.0, 1.0, 20))
        >>> type(ax).__name__, ax.get_xlabel()
        ('Axes', 'time')

        Drawing onto an existing axes returns that axes:

        >>> import matplotlib.pyplot as plt
        >>> _, mine = plt.subplots()
        >>> _, same = vdx.plot_l96_grid(states, np.linspace(0.0, 1.0, 20), ax=mine)
        >>> same is mine
        True
    """
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()
        assert fig is not None

    ax.imshow(
        states.T,
        aspect="auto",
        origin="lower",
        extent=(
            float(time_coords[0]),
            float(time_coords[-1]),
            0.0,
            float(states.shape[1]),
        ),
    )
    ax.set_xlabel("time")
    ax.set_ylabel("variable index")
    return fig, ax  # type: ignore[return-value]  # ty:ignore[invalid-return-type]

plot_reconstruction_comparison ¶

plot_reconstruction_comparison(
    target: ndarray,
    masked_input: ndarray,
    reconstruction: ndarray,
    *,
    sample_idx: int = 0,
) -> tuple[Figure, ndarray]

Side-by-side comparison of target, masked input, and reconstruction.

Draws three panels (time on the x-axis, feature on the y-axis) for one batch element.

Parameters:

Name	Type	Description	Default
`target`	`ndarray`	Ground-truth states, indexed `[batch, time, feature]`.	required
`masked_input`	`ndarray`	Masked / noisy observations, same layout.	required
`reconstruction`	`ndarray`	Model reconstruction, same layout.	required
`sample_idx`	`int`	Which batch element to visualize.	`0`

Returns:

Type	Description
`tuple[Figure, ndarray]`	`(fig, axes)`: the Matplotlib figure and the array of three axes.

Examples:

>>> import matplotlib
>>> matplotlib.use("Agg")
>>> import numpy as np
>>> import vardax as vdx
>>> truth = np.random.default_rng(0).normal(size=(2, 6, 4))
>>> fig, axes = vdx.plot_reconstruction_comparison(truth, truth, truth)
>>> axes.shape, type(axes[0]).__name__
((3,), 'Axes')

Another batch element can be selected:

>>> fig, axes = vdx.plot_reconstruction_comparison(
...     truth, truth, truth, sample_idx=1
... )
>>> axes[1].get_title()
'Masked input'

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

def plot_reconstruction_comparison(
    target: np.ndarray,
    masked_input: np.ndarray,
    reconstruction: np.ndarray,
    *,
    sample_idx: int = 0,
) -> tuple[Figure, np.ndarray]:
    """Side-by-side comparison of target, masked input, and reconstruction.

    Draws three panels (time on the x-axis, feature on the y-axis) for one
    batch element.

    Args:
        target: Ground-truth states, indexed ``[batch, time, feature]``.
        masked_input: Masked / noisy observations, same layout.
        reconstruction: Model reconstruction, same layout.
        sample_idx: Which batch element to visualize.

    Returns:
        ``(fig, axes)``: the Matplotlib figure and the array of three axes.

    Examples:
        >>> import matplotlib
        >>> matplotlib.use("Agg")
        >>> import numpy as np
        >>> import vardax as vdx
        >>> truth = np.random.default_rng(0).normal(size=(2, 6, 4))
        >>> fig, axes = vdx.plot_reconstruction_comparison(truth, truth, truth)
        >>> axes.shape, type(axes[0]).__name__
        ((3,), 'Axes')

        Another batch element can be selected:

        >>> fig, axes = vdx.plot_reconstruction_comparison(
        ...     truth, truth, truth, sample_idx=1
        ... )
        >>> axes[1].get_title()
        'Masked input'
    """
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    data = [
        (target[sample_idx], "Target"),
        (masked_input[sample_idx], "Masked input"),
        (reconstruction[sample_idx], "Reconstruction"),
    ]
    for ax, (arr, title) in zip(axes, data, strict=False):
        ax.imshow(arr.T, aspect="auto", origin="lower")
        ax.set_title(title)
        ax.set_xlabel("time")
        ax.set_ylabel("feature")
    fig.tight_layout()
    return fig, axes

Errors¶

Every error vardax raises on purpose derives from VardaxError, and each subclass names what failed. Each one also inherits the builtin it replaced (ValueError, or TypeError for InputTypeError), so existing except ValueError handlers keep working.

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

VardaxError ¶

Bases: Exception

Base class for every error vardax raises on purpose.

Examples:

Catch any vardax error with the base class:

>>> import vardax as vdx
>>> try:
...     vdx.OptimistixMinimiser(max_steps=0)
... except vdx.VardaxError as error:
...     type(error).__name__
'ConfigurationError'

Every subclass is still the builtin it replaced:

>>> issubclass(vdx.ShapeError, ValueError)
True
>>> issubclass(vdx.InputTypeError, TypeError)
True