Skip to content

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

dx/dt = σ(y − x) dy/dt = x(ρ − z) − y dz/dt = xy − βz

Attributes:

Name Type Description
sigma float

Prandtl number (default 10.0).
rho float

Rayleigh number (default 28.0).
beta float

Geometric factor (default 8/3).
Source code in src/vardax/_src/utils/dynamical_systems.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Lorenz63(eqx.Module):
    """Lorenz-63 vector field as an Equinox module.

    Implements the system:
        dx/dt = σ(y − x)
        dy/dt = x(ρ − z) − y
        dz/dt = xy − βz

    Attributes:
        sigma: Prandtl number (default 10.0).
        rho: Rayleigh number (default 28.0).
        beta: Geometric factor (default 8/3).
    """

    sigma: float
    rho: float
    beta: float

    def __call__(
        self,
        t: float,
        y: Float[Array, 3],
        args: None,
    ) -> Float[Array, 3]:
        """Evaluate the vector field at state ``y`` and time ``t``."""
        x_, y_, z_ = y[0], y[1], y[2]
        dx = self.sigma * (y_ - x_)
        dy = x_ * (self.rho - z_) - y_
        dz = x_ * y_ - self.beta * z_
        return jnp.array([dx, dy, dz])

Lorenz96

Bases: Module

Lorenz-96 vector field as an Equinox module.

Implements the system (with periodic boundary conditions): dx_k/dt = (x_{k+1} - x_{k-2}) * x_{k-1} - x_k + F

Attributes:

Name Type Description
F float

Forcing constant (default 8.0).
Source code in src/vardax/_src/utils/dynamical_systems.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
class Lorenz96(eqx.Module):
    """Lorenz-96 vector field as an Equinox module.

    Implements the system (with periodic boundary conditions):
        dx_k/dt = (x_{k+1} - x_{k-2}) * x_{k-1} - x_k + F

    Attributes:
        F: Forcing constant (default 8.0).
    """

    F: float

    def __call__(
        self,
        t: float,
        y: Float[Array, N],  # type: ignore[unresolved-reference]  # ty:ignore[unresolved-reference]
        args: None,
    ) -> Float[Array, N]:  # type: ignore[unresolved-reference]  # ty:ignore[unresolved-reference]
        """Evaluate the vector field at state ``y`` and time ``t``."""
        # 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.F

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.

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.

 10.0
rho float

Rayleigh number.

 28.0
beta float

Geometric factor.

 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 of shape (3,). 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) — time coordinates of shape (T,)
Float[Array, 'T 3']

(starting at 0) and the Lorenz-63 state trajectory of shape
tuple[Float[Array, T], Float[Array, 'T 3']]

(T, 3).
Source code in src/vardax/_src/utils/dynamical_systems.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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"]]:  # type: ignore[unresolved-reference]  # ty:ignore[unresolved-reference]
    """Simulate the Lorenz-63 system and return state trajectory.

    Args:
        key: JAX PRNG key used to perturb the initial condition when
            ``x0`` is ``None``.
        sigma: Prandtl number.
        rho: Rayleigh number.
        beta: Geometric factor.
        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 of shape ``(3,)``.  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)`` — time coordinates of shape ``(T,)``
        (starting at 0) and the Lorenz-63 state trajectory of shape
        ``(T, 3)``.
    """
    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
    t0 = 0.0
    t1 = total_steps * dt

    save_times = jnp.linspace(t0, t1, total_steps + 1)

    sol = diffeqsolve(
        ODETerm(model),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
        Tsit5(),
        t0=t0,
        t1=t1,
        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 = sol.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: int = 40,
    F: 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.

Parameters:

Name Type Description Default
key Array

JAX PRNG key used to perturb the initial condition.

 required
N int

Number of variables (spatial dimension).

 40
F 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) — time coordinates of shape (T,)
Float[Array, 'T N']

(starting at 0) and the state trajectory of shape
tuple[Float[Array, T], Float[Array, 'T N']]

(n_steps + 1, N).
Source code in src/vardax/_src/utils/dynamical_systems.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def simulate_lorenz96(
    key: Array,
    *,
    N: int = 40,
    F: float = 8.0,
    dt: float = 0.01,
    n_steps: int = 5000,
    n_burn_in: int = 1000,
) -> tuple[Float[Array, T], Float[Array, "T N"]]:  # type: ignore[unresolved-reference]  # ty:ignore[unresolved-reference]
    """Simulate the Lorenz-96 system and return state trajectory.

    Args:
        key: JAX PRNG key used to perturb the initial condition.
        N: Number of variables (spatial dimension).
        F: 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)`` — time coordinates of shape ``(T,)``
        (starting at 0) and the state trajectory of shape
        ``(n_steps + 1, N)``.
    """
    model = Lorenz96(F=F)

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

    total_steps = n_burn_in + n_steps
    t0 = 0.0
    t1 = total_steps * dt

    save_times = jnp.linspace(t0, t1, total_steps + 1)

    sol = diffeqsolve(
        ODETerm(model),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
        Tsit5(),
        t0=t0,
        t1=t1,
        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,
    )

    states = sol.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 of shape (T,).

 required

Returns:

Type Description
Float[Array, 'T-1 2']

Array of shape (T - 1, 2) of consecutive time pairs.

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)
Source code in src/vardax/_src/utils/patches.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def time_patches(ts: Float[Array, T]) -> Float[Array, "T-1 2"]:  # type: ignore[unresolved-reference]  # ty:ignore[unresolved-reference]
    """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 of shape ``(T,)``.

    Returns:
        Array of shape ``(T - 1, 2)`` of consecutive time pairs.

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

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
Source code in src/vardax/_src/utils/forward.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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.

    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
    """

    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:
        b = self.bound
        a = 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 = jnp.maximum(a - b, 0.0) / b
        outside = jnp.sign(state) * (b + b * u / (1.0 + u))
        return jnp.where(a <= b, state, outside)

    def step(self, state: Array, dt: float) -> Array:
        """Advance ``state`` by ``dt`` with confinement before and after."""
        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.

Source code in src/vardax/_src/utils/forward.py 
90
91
92
def step(self, state: Array, dt: float) -> Array:
    """Advance ``state`` by ``dt`` with confinement before and after."""
    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,
) -> Float[Array, " n_runs"]

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.

Parameters:

Name Type Description Default
sample_posterior Callable[[Array, PRNGKeyArray, int], Array]

(y, key, n) -> Array of shape (n, *state_shape) — posterior samples 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
Float[Array, ' n_runs']

(n_runs,) int array of ranks in [0, n_samples].
Source code in src/vardax/_src/utils/validation.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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,
) -> Float[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.

    Args:
        sample_posterior: ``(y, key, n) -> Array`` of shape
            ``(n, *state_shape)`` — posterior samples 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:
        ``(n_runs,)`` int array of ranks in ``[0, n_samples]``.
    """
    keys = jax.random.split(key, n_runs)
    ranks = []
    for k in keys:
        k_prior, k_obs, k_post = jax.random.split(k, 3)
        x_true = sample_prior(k_prior)
        y = simulate_obs(x_true, k_obs)
        samples = sample_posterior(y, k_post, n_samples)
        # Flatten and use the L2 norm as the scalar reduction.
        scalar_true = float(jnp.linalg.norm(x_true))
        scalar_samples = jnp.linalg.norm(samples.reshape(n_samples, -1), axis=-1)
        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.

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
AssertionError

If any component exceeds the tolerance.
Source code in src/vardax/_src/utils/validation.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
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.

    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:
        AssertionError: If any component exceeds the tolerance.
    """
    if p_oracle.cov is None:
        raise ValueError(
            "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 ValueError(
            "p_fast.mean and p_oracle.mean must have the same shape; got "
            f"{p_fast.mean.shape} vs {p_oracle.mean.shape}."
        )
    sigma = _marginal_std(p_oracle.cov, p_oracle.mean)
    z = jnp.abs(p_fast.mean - p_oracle.mean) / jnp.maximum(sigma, 1e-12)
    if jnp.any(z > tolerance_sigma):
        worst = float(jnp.max(z))
        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.

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.

 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.
Source code in src/vardax/_src/utils/validation.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
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.

    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.
        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.
    """
    keys = jax.random.split(key, n_probes)
    worst_rel = 0.0
    for k in keys:
        v = jax.random.normal(k, y.shape)
        v = v / jnp.linalg.norm(v)
        _, jv_fast = jax.jvp(fn_fast, (y,), (v,))
        _, jv_oracle = jax.jvp(fn_oracle, (y,), (v,))
        denom = jnp.linalg.norm(jv_oracle) + 1e-12
        rel = float(jnp.linalg.norm(jv_fast - jv_oracle) / denom)
        worst_rel = max(worst_rel, rel)
    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 array of shape (T, N).

 required
time_coords ndarray

Time coordinate array of shape (T,).

 required
n_vars int

Number of evenly-spaced variables to plot.

 5
ax Axes | None

Optional existing Axes to draw on.

 None

Returns:

Type Description
tuple[Figure, Axes]

(fig, ax) — Matplotlib figure and axes.
Source code in src/vardax/_src/utils/viz.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
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 array of shape ``(T, N)``.
        time_coords: Time coordinate array of shape ``(T,)``.
        n_vars: Number of evenly-spaced variables to plot.
        ax: Optional existing ``Axes`` to draw on.

    Returns:
        ``(fig, ax)`` — Matplotlib figure and axes.
    """
    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 ValueError(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 array of shape (T, N).

 required
time_coords ndarray

Time coordinate array of shape (T,).

 required
ax Axes | None

Optional existing Axes to draw on.

 None

Returns:

Type Description
tuple[Figure, Axes]

(fig, ax) — Matplotlib figure and axes.
Source code in src/vardax/_src/utils/viz.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
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 array of shape ``(T, N)``.
        time_coords: Time coordinate array of shape ``(T,)``.
        ax: Optional existing ``Axes`` to draw on.

    Returns:
        ``(fig, ax)`` — Matplotlib figure and axes.
    """
    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.

Parameters:

Name Type Description Default
target ndarray

Ground-truth states of shape (B, T, N).

 required
masked_input ndarray

Masked / noisy observations of shape (B, T, N).

 required
reconstruction ndarray

Model reconstruction of shape (B, T, N).

 required
sample_idx int

Which batch element to visualize.

 0

Returns:

Type Description
tuple[Figure, ndarray]

(fig, axes) — Matplotlib figure and array of axes.
Source code in src/vardax/_src/utils/viz.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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.

    Args:
        target: Ground-truth states of shape ``(B, T, N)``.
        masked_input: Masked / noisy observations of shape ``(B, T, N)``.
        reconstruction: Model reconstruction of shape ``(B, T, N)``.
        sample_idx: Which batch element to visualize.

    Returns:
        ``(fig, axes)`` — Matplotlib figure and array of axes.
    """
    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