vardax

Variational data assimilation in JAX — seven analysis methods, one cost function, zero hand-written adjoints.

📖 Documentation · 🧪 Tutorials · 📐 Maths reference · 🧭 Design docs

---

Optimal interpolation, 3DVar, strong- and weak-constraint 4DVar, incremental 4DVar, 4DVarNet and amortized posteriors are the same estimation problem solved under different assumptions. vardax implements all seven as peer equinox modules over a shared set of seams — observation operator, background covariance, forward model, minimiser, adjoint — so switching method is a constructor change, not a rewrite.

Every method is called on a batch the same way and cycles through the same analysis-step interface:

# today (vardax)
# x* = argmin_x J(x)
analysis = model(batch)
# today (vardax)
step = model.as_analysis_step()

Gradients are composed, not coded: dynamics differentiate through diffrax adjoints, inner minimisers through optimistix adjoints, and covariances stay as lineax operators so nothing is ever densified unless you ask.

🌟 Highlights

• 🧮 Seven peer methods, from closed-form BLUE to amortized posteriors, each an equinox module behind one constructor.
• 🔁 One cycling interface: every model becomes a pipekit_cycle.AnalysisStep through .as_analysis_step().
• 🧷 No hand-written adjoints: gradients come from diffrax and optimistix adjoints, plus the one-step rule of Bolte et al. (2023).
• 🧱 Operator covariances: $B$, $R$ and the posterior covariance are lineax operators, never dense matrices you did not ask for.
• 🛰️ Observation operators for gappy grids, off-grid points, retrieval averaging kernels and multi-instrument fusion.
• ✅ Validation gates that decide whether a fast method may replace a trusted one.
• 📓 Seventeen tutorials executed at documentation build time.

📐 The seven methods

All seven ship today, and all seven satisfy pipekit_cycle.AnalysisStep via .as_analysis_step().

Class	Method	Control	Reach for it when
OptimalInterpolation	BLUE / OI, closed form	$x$	Linear $H$, Gaussian $B$/$R$. The right default.
ThreeDVar	3D variational	$x$	Snapshot inversion with nonlinear $H$
StrongFourDVar	Strong-constraint 4DVar	$x_0$	A window of data, dynamics trusted
WeakFourDVar	Weak-constraint 4DVar	$(x_0, \boldsymbol{\eta})$	A window of data, model error active
IncrementalFourDVar	Gauss–Newton outer, CG inner	$\delta x_0$	The operational fast path
FourDVarNet1D / FourDVarNet2D	Learned prior $\varphi_\theta$ + learned solver	$x$	Dense training data, no trusted model
AmortizedPosterior	Direct head $q_\phi(x \mid y)$	—	Real-time, many-event regimes

ReducedBasisFourDVar is StrongFourDVar with the control moved into a reduced basis, $x_0 = x_b + \Phi X$ (see docs/06). It was called Incremental4DVar before 0.1.11; that name is deprecated.

Each is a specialisation of one cost:

$$x^* = \underset{x, \boldsymbol{\eta}}{\arg\min} \underbrace{\frac{1}{2}|x - x_b|^2_{B^{-1}}}{\text{background}} + \underbrace{\frac{1}{2}\sum{t=0}^{T} |y_t - H_t(M_t(x; \boldsymbol{\eta}))|^2_{R_t^{-1}}}{\text{observations}} \left[ + \underbrace{\frac{1}{2}\sum{t=1}^{T} |\eta_t|^2_{Q_t^{-1}}}_{\text{model error}} \right].$$

Which terms are active, and how the minimisation proceeds (closed form, iterative, learned, amortized), is the only thing that differs. The problem setting chapter derives it.

📥 Install

vardax is not on PyPI yet. Install from the repository with uv, which resolves the git dependencies declared in [tool.uv.sources]. gaussx and geonnax are pinned to a revision there; the pipekit packages track their default branch and uv.lock is not committed, so a fresh install can pick up newer pipekit commits than the last one tested:

git clone https://github.com/jejjohnson/vardax.git
cd vardax
make install          # uv sync --all-extras + pre-commit hooks

To depend on it from another uv project, map vardax's own git dependencies too. uv does not read a dependency's [tool.uv.sources], so a consumer sees plain gaussx, geonnax, pipekit and pipekit-cycle requirements and tries its indexes: three of those are not on PyPI, and pipekit there is an unrelated project, so without these mappings the install resolves wrongly or not at all.

[project]
dependencies = ["vardax"]

[tool.uv.sources]
vardax = { git = "https://github.com/jejjohnson/vardax" }
gaussx = { git = "https://github.com/jejjohnson/gaussx", rev = "cd6264c" }
geonnax = { git = "https://github.com/jejjohnson/geonnax", rev = "5318173" }
pipekit = { git = "https://github.com/jejjohnson/pipekit", subdirectory = "packages/pipekit" }
pipekit-cycle = { git = "https://github.com/jejjohnson/pipekit", subdirectory = "packages/pipekit-cycle" }

Add the matching entries for pipekit-jax, pipekit-experiment or pipekit-train if you enable the persist or train extras.

Requires Python ≥ 3.12, < 3.14. Extras: data, viz, exp, jlab, examples, persist, train, all.

🌪️ Quickstart — 4DVar on Lorenz-63

A window of noisy, half-missing observations; the true chaotic dynamics as the forward model; BFGS on the variational cost with the adjoint supplied by jax.grad through diffrax.

Here $x_b$ is the background, $B = \sigma_b^2 I$ and $R = \sigma_o^2 I$ the background- and observation-error covariances, $M_t$ the Lorenz-63 flow to time $t$ and $m_t$ the observation mask.

import einx
import jax
import jax.numpy as jnp
import lineax as lx

import vardax as vdx

# Shapes: S = 2,000 simulated steps; T = 21 window times; D = 3 state variables

key_truth, key_obs, key_background = jax.random.split(jax.random.PRNGKey(0), 3)

# today (vardax)
# dx/dt = f(x), sampled every Δt = 0.01 after a burn-in on the attractor
_, trajectory = vdx.simulate_lorenz63(  # → (S, D)
    key_truth, dt=0.01, n_steps=2000, n_burn_in=1000
)
# x_true(t) for t = 0, 0.05, …, 1.0
true_window = trajectory[500 : 500 + 5 * 21 : 5]  # (S, D) → (T, D)

# m_t = 1 on every other time, 0 otherwise
mask = jnp.zeros((21, 3)).at[::2].set(1.0)  # → (T, D)
# ε_t ~ 𝒩(0, I)
obs_noise = jax.random.normal(key_obs, (21, 3))  # → (T, D)
# y_t = x_true(t) + σ_o ε_t
observations = true_window + 0.5 * obs_noise  # (T, D) → (T, D)
# η ~ 𝒩(0, I)
background_noise = jax.random.normal(key_background, (3,))  # → (D,)
# x_b = x_true(0) + σ_b η
background = true_window[0] + 2.0 * background_noise  # (D,) → (D,)


# user
# Σ = σ² I, tagged positive semi-definite for the CG solve
def diagonal_covariance(variance: float) -> lx.AbstractLinearOperator:
    diagonal = jnp.full(3, variance)
    operator = lx.DiagonalLinearOperator(diagonal)
    return lx.TaggedLinearOperator(operator, lx.positive_semidefinite_tag)


# today (vardax)
# dx/dt = (σ(x₂ − x₁), x₁(ρ − x₃) − x₂, x₁x₂ − βx₃)
lorenz = vdx.Lorenz63(sigma=10.0, rho=28.0, beta=8.0 / 3.0)
# today (vardax)
# M_t : x₀ ↦ x(t), integrated with diffrax
forward = vdx.DynTrajectory(model=lorenz).as_forward_model(dt=0.05)

# today (vardax)
# x₀ᵃ = argmin ½‖x₀ − x_b‖²_B⁻¹ + ½ Σ_t ‖m_t ⊙ (y_t − M_t(x₀))‖²_R⁻¹
model = vdx.StrongFourDVar(
    forward=forward,
    obs_op=vdx.MaskedIdentity(),
    prior_mean=background,
    prior_cov_op=diagonal_covariance(2.0**2),
    obs_cov_op=diagonal_covariance(0.5**2),
)

# m ⊙ y
masked_observations = observations * mask  # (T, D) → (T, D)
# a batch of one window
batch_input, batch_mask, batch_target = einx.id(  # (T, D) → (1, T, D)
    "t d, t d, t d -> 1 t d, 1 t d, 1 t d",
    masked_observations,
    mask,
    true_window,
)
# today (vardax)
batch = vdx.Batch1D(input=batch_input, mask=batch_mask, target=batch_target)
# today (vardax)
analysis = model(batch)[0]  # (1, T, D) → (D,)

# ‖x_b − x_true(0)‖
background_error = jnp.linalg.norm(background - true_window[0])  # → ()
# ‖x₀ᵃ − x_true(0)‖
analysis_error = jnp.linalg.norm(analysis - true_window[0])  # → ()
print(f"background error {background_error:.2f}, analysis error {analysis_error:.2f}")

This prints background error 3.84, analysis error 0.81.

Correlated background errors are a one-line change: swap diagonal_covariance(...) for a gaussx Matérn operator and the conjugate-gradient solve adapts.

♻️ Swap the method, keep the pipeline

Every model exposes .as_analysis_step(), so the cycling code never sees which method it is running: model can be any of the seven. In the sketch below, $x_c^f$ and $x_c^a$ are the forecast and analysis at cycle $c$, $M$ the forward model over one cycle and $y_c$ that cycle's observations; load_window, n_cycles, obs_cov and initial_analysis are yours to supply.

import pipekit
import pipekit_cycle as pc

# Shapes: C cycles; D = 3 state variables

# today (pipekit)
# c ↦ y_c, with load_window the user's reader for cycle c
obs_source = pipekit.Lambda(load_window)

# today* (vardax)
# x_c^f = M(x_{c−1}^a),  x_c^a = analysis(x_c^f, y_c)
cycle = vdx.VarDACycle(
    forward=forward,
    obs_op=vdx.MaskedIdentity(),
    model=model,
    obs_source=obs_source,
    n_steps=n_cycles,
)

# today (pipekit_cycle)
initial_state = pc.DAState(t=0.0, cycle_count=0, obs_err_cov=obs_cov)
final_analysis, final_state = cycle(initial_analysis, initial_state)  # → (D,)
# [(x_c^f, x_c^a)] for c = 1, …, C
forecasts, analyses = zip(*cycle.history)  # C pairs of (D,)

Each cycle forecasts the carrier one step, hands the forecast and that cycle's observations (NaN marks a gap) to the analysis step, and keeps the pair in cycle.history. VarSmootherCycle does the same over fixed-lag windows. Both return plain pipekit_cycle objects, so whatever orchestrates pipekit operators orchestrates vardax. Tutorial 16 runs the loop for real, including the forecast-as-background analysis step that makes each cycle's $x_b$ the previous forecast.

🔬 What else is in the box

Area	Symbols	Notes
Observation operators	MaskedIdentity, LinearObs, InterpObs, AveragingKernel, MultiInstrumentFusion, interp_obs_from_coords	Gappy grids, off-grid interpolation with a passing adjoint test, retrieval averaging kernels, and multi-instrument fusion that satisfies the precision-addition identity
Posterior	LaplaceCovariance, GaussNewtonHessian, EnsembleCovariance, Posterior, GaussianMarkLikelihood	Covariance as a lineax operator, never a dense matrix you did not ask for
Adjoints	OneStepAdjoint, KStepAdjoint, ImplicitAdjoint, RecursiveCheckpointAdjoint, to_optimistix_adjoint	The Bolte (2023) one-step rule packaged as an optimistix.AbstractAdjoint, targeting upstream contribution
Dynamical priors	DynTrajectory, DynIncrements, DynamicalPrior	An ODE prior that doubles as a forward model (.as_forward_model(dt)) or a trajectory loss (.bind(ts))
Reduced bases	LinearBasis, CompositeBasis, eof_basis, fourier_basis, rbf_basis, wavelet_basis	Control-variable transforms for incremental 4DVar
Amortized inference	AmortizedPosterior, MLPObsEncoder, IdentityObsEncoder, RegressionHead	Train once on simulations, assimilate in a forward pass. ConditionalFlowHead and ScoreDiffusionHead are declared but raise NotImplementedError on construction, pending gauss_flows and the reverse-SDE sampler
Validation gates	assert_posterior_agreement, assert_adjoint_calibrated, simulation_based_calibration	The go/no-go checks that decide whether a fast method may replace a trusted one
Forward adapters	SoftBoundedForward

📓 Tutorials

Seventeen executable walkthroughs, all jupytext sources under docs/notebooks/ that are run at documentation build time — the figures on the site are produced by the code in this repository, not committed snapshots.

#	Tutorial	What it shows
01–02	Model-based 4DVar, unrolling vs fixed point	The reference 4DVar picture, then the learned solver against it
03–04	4DVarNet on L63, 2-D demo	Training the learned variant in 1-D and on spatiotemporal fields
05–07	L63 pipeline, L96 pipeline, prior pre-training	Simulation → patching → masking → training, end to end
08	Classical 4DVar vs 4DVarNet	Gradient descent on the cost against the learned solver
09–11	Parameter estimation, bilevel optimisation, gradient learning	Learning ODE parameters, cost weights, and the update rule itself

📖 Documentation

• Maths reference — 21 chapters: the Bayesian foundation (1–3), each analysis method (4–10), cross-cutting concerns (11–14), worked examples on Lorenz, sea-surface height, methane and latent-space DA (15–18), then physical models, uncertainty quantification and OceanBench (19–21).
• API reference — every symbol vardax exports, grouped by seam and generated from the docstrings. The two it re-exports from optimistix, ImplicitAdjoint and RecursiveCheckpointAdjoint, are described there and linked upstream rather than regenerated.
• Design docs — architecture, ownership boundaries, the decision log, and the pipekit composition story.

🧩 Ecosystem

vardax owns the DA algorithms and nothing else. Forward models, optimisers, ODE solvers, ensemble methods and structured linear algebra are composed in:

Concern	Owner
Geophysical forward models (SWM, QG, primitive equations)	somax
Atmospheric transport and methane retrieval	plumax
ODE / SDE integration and adjoints	diffrax
Minimisers and their adjoints	optimistix
Linear solvers (CG, GMRES, Lanczos)	lineax
Structured operators (Matérn, Kronecker, low rank)	gaussx
Ensemble filters and smoothers (EnKF / EnKS / EnKI)	filterax
Operator composition and DA cycle protocols	pipekit + pipekit-cycle
Neural building blocks	geonnax

See boundaries for the full ownership map.

🌳 Repository layout

src/vardax/
├── __init__.py              ← the public API; everything below is private
└── _src/
    ├── models/              ← OI, 3DVar, strong / weak / incremental 4DVar
    ├── model.py             ← FourDVarNet1D / FourDVarNet2D (equinox)
    ├── amortized/           ← encoders, heads, AmortizedPosterior
    ├── obs_operators/       ← masked, linear, interp, averaging kernel, fusion
    ├── posterior/           ← Laplace, Gauss–Newton, ensemble, container
    ├── adjoints/            ← one-step, k-step, optimistix mapping
    ├── basis/               ← reduced bases and control-variable transforms
    ├── cycle/               ← VarDACycle, VarSmootherCycle
    ├── priors.py            ← autoencoder and Lorenz priors
    ├── priors_dynamical.py  ← DynTrajectory, DynIncrements
    ├── costs.py             ← the variational cost and its decompositions
    ├── solver.py            ← unrolled / one-step / fixed-point solvers
    ├── training.py          ← train and eval steps
    └── utils/               ← Lorenz simulators, SoftBoundedForward, validation gates, plotting
docs/                        ← 21 maths chapters, API pages, design docs, 17 executed tutorials
tests/                       ← 376 tests across 29 modules

⚙️ Development

make install        # uv sync --all-extras + pre-commit
make uv-test        # pytest -v
make uv-test-cov    # pytest with coverage XML
make uv-format      # ruff format + ruff check --fix
make uv-lint        # ruff check + ty check
make uv-pre-commit  # all hooks on all files
make docs-serve     # MyST site + API reference at http://127.0.0.1:8000

Four checks must pass before a commit lands:

uv run pytest -q
uv run --group lint ruff check .            # entire repo, not just src/
uv run --group lint ruff format --check .
uv run --group typecheck ty check src/vardax

Dependency groups follow PEP 735 (dev, test, lint, typecheck, docs); activate one with uv sync --group <name>. Commits are conventional — release-please cuts the versions. Committed .ipynb files are rejected by a pre-commit hook: edit the jupytext .py source instead.

🎓 Citation and references

• Fablet et al. (2021). Learning Variational Data Assimilation Models and Solvers. JAMES.
• Fablet et al. (2023). Multimodal 4DVarNets for the reconstruction of sea surface dynamics from NADIR and wide-swath altimetry. IEEE TGRS.
• Bolte, Pauwels & Vaiter (2023). One-step differentiation of iterative algorithms. NeurIPS.
• Courtier, Thépaut & Hollingsworth (1994). A strategy for operational implementation of 4D-Var, using an incremental approach. QJRMS.
• Carrassi et al. (2018). Data assimilation in the geosciences: an overview of methods, issues, and perspectives. WIREs Climate Change.
• Full bibliography in docs/references.md.

📄 License

MIT. See LICENSE.