vardax Tutorial Master List

A reconciled, exhaustive curriculum for learning data assimilation through vardax — JAX-native variational DA shipping seven peer pipekit_cycle.AnalysisStep methods (OI / 3DVar / strong / weak / incremental 4DVar / FourDVarNet / AmortizedPosterior) on a single set of eqx.Module primitives.

Ensemble-DA tutorials live in the sister list TUTORIAL_MASTER_LIST_FILTER.md; GP / SVI tutorials live in ../../gaussian_processes/TUTORIAL_MASTER_LIST.md. Cross-listed items (Bayesian update, structured covariances, sigma points, validation gates, latent-space generative priors) are flagged 🔁.

Legend — Source columns:

• V = exists in vardax (docs/notebooks/<name>)
• G = exists in gaussx (docs/notebooks/<name>)
• K = exists in pipekit (docs/notebooks/<name>)
• R = exists in research_notebook (projects/assimilation/notebooks/<path>)
• — = does not exist yet (gap)

Scope tag: 🧱 fundamental · 🔬 research · 🌉 bridge · 🔁 cross-listed (filterax / GP master list)

Refs column: gh#N = open GitHub issue · dd:path = vardax docs/design/<path> · math:N = vardax docs/0N_<chapter>.md · api:foo = vardax exported symbol.

Math policy: every section’s first notebook (00_*) ports the relevant math chapter from vardax/docs/ directly into the tutorial, so each part is self-contained for someone learning DA from scratch. Subsequent notebooks in the section are API-driven.

Curriculum at a glance¶

• Part 0 — Foundations (learning DA from scratch)
  • 0.A — The Bayesian update
  • 0.B — Linear-Gaussian closed form (BLUE / Kalman)
  • 0.C — Variational reformulation (cost ⇄ posterior)
  • 0.D — pipekit-cycle protocols (the vardax shape)
  • 0.E — Your first assimilation
• Part 1 — Observation operators
  • 1.A — MaskedIdentity and the symmetric-masking invariant
  • 1.B — LinearObs with explicit H
  • 1.C — AveragingKernel (satellite footprints, RTM intuition)
  • 1.D — MultiInstrumentFusion and heterogeneous obs
  • 1.E — Writing a custom ObservationOperator
• Part 2 — Forward models
  • 2.A — The pipekit_cycle.ForwardModel protocol
  • 2.B — Lorenz family wrappers (L63, L96, L96-2L)
  • 2.C — Shallow water (somax)
  • 2.D — Quasi-geostrophic (somax)
  • 2.E — Plume models (plumax)
  • 2.F — Writing your own forward
• Part 3 — Static priors & covariances
  • 3.A — Diagonal B (the simplest case)
  • 3.B — Structured priors via gaussx (Matérn, Kronecker)
  • 3.C — Dynamical priors (forward-as-φ)
• Part 4 — Encoders & decoders (latent-space DA)
  • 4.A — AE priors overview (BilinAE, ConvAE, MLP)
  • 4.B — Pretraining an AE prior
  • 4.C — Latent-space DA (assimilate in z, decode to x)
  • 4.D — Observation encoders for amortized inference
• Part 5 — Optimal Interpolation (BLUE)
  • 5.A — BLUE derivation (closed-form Kalman update)
  • 5.B — OptimalInterpolation API walkthrough
  • 5.C — OI with structured B
• Part 6 — 3DVar
  • 6.A — The 3DVar cost
  • 6.B — ThreeDVar API + minimiser choice
  • 6.C — Nonlinear H with 3DVar
• Part 7 — Strong-constraint 4DVar
  • 7.A — The 4DVar cost (adjoint of M)
  • 7.B — StrongFourDVar API
  • 7.C — forward_adjoint choices (diffrax variants)
  • 7.D — Lorenz demo (deep dive, links to assimilation/)
• Part 8 — Weak-constraint 4DVar
  • 8.A — Model-error term
  • 8.B — WeakFourDVar API
  • 8.C — Imbalance failure modes
• Part 9 — Incremental 4DVar
  • 9.A — GN outer + CG inner
  • 9.B — IncrementalFourDVar API
  • 9.C — Control-variable transform (gaussx √B preconditioning)
• Part 10 — FourDVarNet (learned solvers)
  • 10.A — Unrolled iteration math
  • 10.B — FourDVarNet1D / FourDVarNet2D API
  • 10.C — Training loop (train_step composition)
  • 10.D — solver_adjoint dispatch (D15)
• Part 11 — Amortized posterior
  • 11.A — Simulation-based inference framing
  • 11.B — Regression head
  • 11.C — Conditional-flow head (stub — pending gauss_flows)
  • 11.D — Score-diffusion head (stub — pending diffrax reverse-SDE)
• Part 12 — Adjoints deep dive
  • 12.A — What is an adjoint?
  • 12.B — diffrax adjoint variants
  • 12.C — optimistix adjoint variants
  • 12.D — Writing your own AbstractAdjoint
• Part 13 — Posterior uncertainty
  • 13.A — Laplace at MAP
  • 13.B — Gauss-Newton Hessian
  • 13.C — Ensemble covariances
  • 13.D — GaussianMarkLikelihood serialisation
• Part 14 — Validation gates (six-step cycle)
  • 14.A — Six-step cycle motivation (D12)
  • 14.B — assert_posterior_agreement
  • 14.C — assert_adjoint_calibrated
  • 14.D — simulation_based_calibration (Talts et al.)
• Part 15 — Orchestration (pipekit composition)
  • 15.A — VarDACycle for sequential DA
  • 15.B — VarSmootherCycle for retrospective windows
  • 15.C — obs_source operators
  • 15.D — Composing with pipekit.Sequential
• Part 16 — Training at scale (pipekit-train)
  • 16.A — train_step / amortized_train_step primitives
  • 16.B — pipekit_train.Loss adapter
  • 16.C — TrainingLoop integration
  • 16.D — Checkpoints & sweeps
• Part 17 — Performance
  • 17.A — jit basics
  • 17.B — eqx.filter_vmap batching
  • 17.C — Memory profiling adjoint variants
  • 17.D — When to unroll vs scan
• Part 18 — Debugging
  • 18.A — Cost decomposition (decomposed_loss)
  • 18.B — Inspecting residuals
  • 18.C — Common failure modes
• Part 19 — Extending vardax
  • 19.A — Writing your own AnalysisStep
  • 19.B — Protocol-conformance tests
  • 19.C — Publishing as a plugin
• Part 20 — Applied (cross-link to applied projects)
  • 20.A — Lorenz benchmarks (→ projects/assimilation/)
  • 20.B — SSH satellite interpolation
  • 20.C — Methane single-overpass retrieval
  • 20.D — Multi-instrument atmospheric fusion

Part 0 — Foundations¶

0.A — The Bayesian update¶

Key equations / models:

• Prior, likelihood, posterior: $p(x \mid y) \propto p(y \mid x)\,p(x)$
• Gaussian-Gaussian closure: $p(x) = \mathcal{N}(x_b, B)$, $p(y \mid x) = \mathcal{N}(Hx, R)$
• Posterior in information form: $\Lambda^a = B^{-1} + H^\top R^{-1} H$, $\eta^a = B^{-1} x_b + H^\top R^{-1} y$

0.B — Linear-Gaussian closed form (BLUE)¶

Key equations / models:

• Kalman gain: $K = B H^\top (H B H^\top + R)^{-1}$
• Update: $x^a = x_b + K(y - H x_b)$
• Posterior cov: $P^a = (I - K H) B$ (or Joseph form)

0.C — Variational reformulation (cost ⇄ posterior)¶

Key equations / models:

• 3DVar cost: $J(x) = \tfrac{1}{2}\|x - x_b\|^2_{B^{-1}} + \tfrac{1}{2}\|y - H x\|^2_{R^{-1}}$
• Argmin ↔ posterior mean equivalence (linear-Gaussian)
• Why we minimise: nonlinear $H$, structured $B$, large state dim

0.D — pipekit-cycle protocols (the vardax shape)¶

Key equations / models:

• ForwardModel: step(state, dt) → state, dt property
• ObservationOperator: __call__(state, mask=...) → obs, linearize(x)
• AnalysisStep: __call__(forecast, obs, *, obs_op, obs_err_cov) → analysis

0.E — Your first assimilation¶

Part 1 — Observation operators¶

Every method in vardax accepts an ObservationOperator. This part walks through the shipped catalogue and the protocol.

1.A — MaskedIdentity and the symmetric-masking invariant¶

Key equations / models:

• $H(x) = m \odot x$, $\mathrm{linearize}(x) = \mathrm{diag}(m)$
• Symmetric masking: $r = m \odot (y - H(x))$ — both sides masked
• Why naive m \odot y - H(x) biases the analysis (PR #41 review story)

1.B — LinearObs with explicit H¶

Key equations / models:

• $H \in \mathbb{R}^{N_y \times N_x}$ as a lineax.AbstractLinearOperator
• Structured H via gaussx operators (Toeplitz, block-diagonal)

1.C — AveragingKernel (satellite footprints, RTM intuition)¶

Key equations / models:

• $H(x) = A(h \odot x + (1 - h) x_a)$
• Smoothing-kernel + a-priori interpretation (Rodgers 2000)
• Why $\mathrm{linearize}$ gives a structured Jacobian

1.D — MultiInstrumentFusion and heterogeneous obs¶

Key equations / models:

• Per-instrument $H_i$, $R_i$; output = dict[str, Array]
• InstrumentRegistry / InstrumentSpec plumbing

1.E — Writing a custom ObservationOperator¶

Part 2 — Forward models¶

2.A — The pipekit_cycle.ForwardModel protocol¶

Key equations / models:

• step(state, dt) → state — one integration step
• state_signature — JAX shape/dtype declaration
• One-step + jax.lax.scan for rollouts

2.B — Lorenz family wrappers (L63, L96, L96-2L)¶

2.C — Shallow water (somax)¶

Key equations / models:

• $\partial_t h + \nabla \cdot (hv) = 0$, $\partial_t v + v \cdot \nabla v + g \nabla h = -fk \times v + \nu \Delta v$
• Arakawa C-grid via finitevolx

2.D — Quasi-geostrophic (somax)¶

Key equations / models:

• QG potential vorticity: $q = \nabla^2 \psi + \beta y - F\psi$
• Layered QG for ocean / atmosphere

2.E — Plume models (plumax)¶

Key equations / models:

• Gaussian plume: $C(x, y, z) = \frac{Q}{2\pi u \sigma_y \sigma_z} \exp(\ldots)$
• Gaussian puff (time-resolved): superposed instantaneous releases

2.F — Writing your own forward¶

Part 3 — Static priors & covariances¶

3.A — Diagonal B (the simplest case)¶

Key equations / models:

• $B = \sigma_b^2 I$
• The PSD-tag requirement for lineax.CG consumers
• Why lx.DiagonalLinearOperator must be wrapped in TaggedLinearOperator

3.B — Structured priors via gaussx (Matérn, Kronecker)¶

Key equations / models:

• Matérn-ν covariance with length-scale $\ell$
• Kronecker structure for space × time
• Half-operator $\sqrt{B}$ for control-variable transform

3.C — Dynamical priors (forward-as-φ)¶

Key equations / models:

• $\varphi(x) = M(x)$ — N applications of the forward
• DynamicalPrior(forward, n_steps, forward_adjoint)

Part 4 — Encoders & decoders (latent-space DA)¶

The “latent-space” angle on vardax — autoencoder priors as a regulariser, latent-space assimilation, and observation encoders for the amortized head.

4.A — AE priors overview (BilinAE, ConvAE, MLP)¶

Key equations / models:

• $\varphi_\theta(x) = D_\theta(E_\theta(x))$ — encode-then-decode
• Bottleneck dim $\ll N_x$ enforces a manifold prior
• Bilinear, conv, MLP architectures shipped in vardax.priors

4.B — Pretraining an AE prior¶

Key equations / models:

• Reconstruction loss: $\mathcal{L} = \mathbb{E}_x\|x - \varphi_\theta(x)\|^2$
• Simulation-based pretraining via vmap’d generate_problem

4.C — Latent-space DA (assimilate in z, decode to x)¶

Key equations / models:

• Latent posterior: $q(z \mid y) \propto p(y \mid D(z))\,p(z)$
• Decode-then-observe: $H \circ D$ as the effective obs op
• Why this conditions the problem (lower-dim z = better-posed)

4.D — Observation encoders for amortized inference¶

Key equations / models:

• Context encoder $c_\psi(y, m)$ for AmortizedPosterior
• Different role from prior encoder: maps obs not state
• IdentityObsEncoder vs MLPObsEncoder

Part 5 — Optimal Interpolation (BLUE)¶

5.A — BLUE derivation (closed-form Kalman update)¶

5.B — OptimalInterpolation API walkthrough¶

5.C — OI with structured B¶

Part 6 — 3DVar¶

6.A — The 3DVar cost¶

6.B — ThreeDVar API + minimiser choice¶

6.C — Nonlinear H with 3DVar¶

Part 7 — Strong-constraint 4DVar¶

7.A — The 4DVar cost (adjoint of M)¶

7.B — StrongFourDVar API¶

7.C — forward_adjoint choices (diffrax variants)¶

7.D — Lorenz demo (deep dive, links to assimilation/)¶

Part 8 — Weak-constraint 4DVar¶

8.A — Model-error term¶

Key equations / models:

• $x_t = M(x_{t-1}) + \eta_t$, $\eta_t \sim \mathcal{N}(0, Q)$
• Augmented control $(x_0, \eta_1, \ldots, \eta_T)$ in $\mathbb{R}^{N + TN}$

8.B — WeakFourDVar API¶

8.C — Imbalance failure modes¶

Part 9 — Incremental 4DVar¶

9.A — GN outer + CG inner¶

Key equations / models:

• Linearise $M$, $H$ at $x_b^{(k)}$
• GN Hessian: $J''_{GN} = B^{-1} + \sum_t (H'_t M'_t)^\top R^{-1} (H'_t M'_t)$
• Solve $J''_{GN}\,\delta x^* = -\nabla J$ by lineax.CG

9.B — IncrementalFourDVar API¶

9.C — Control-variable transform (gaussx √B preconditioning)¶

Part 10 — FourDVarNet (learned solvers)¶

10.A — Unrolled iteration math¶

Key equations / models:

• $x^{(k+1)} = x^{(k)} - \alpha\,\Phi_\phi(\nabla J(x^{(k)}), x^{(k)}, h^{(k)})$
• Learned ConvLSTM modulator $\Phi_\phi$
• Learned prior $\varphi_\theta$ inside $J$

10.B — FourDVarNet1D / FourDVarNet2D API¶

10.C — Training loop (train_step composition)¶

10.D — solver_adjoint dispatch (D15)¶

Part 11 — Amortized posterior¶

11.A — Simulation-based inference framing¶

11.B — Regression head¶

11.C — Conditional-flow head (stub — pending gauss_flows)¶

11.D — Score-diffusion head (stub — pending diffrax reverse-SDE)¶

Part 12 — Adjoints deep dive¶

12.A — What is an adjoint?¶

12.B — diffrax adjoint variants¶

12.C — optimistix adjoint variants¶

12.D — Writing your own AbstractAdjoint¶

Part 13 — Posterior uncertainty¶

13.A — Laplace at MAP¶

Key equations / models:

• $P^* \approx \bigl((H')^\top R^{-1} H' + B^{-1}\bigr)^{-1}$
• Returned as lineax.AbstractLinearOperator (lazy CG-backed inverse)

13.B — Gauss-Newton Hessian¶

13.C — Ensemble covariances¶

13.D — GaussianMarkLikelihood serialisation¶

Part 14 — Validation gates (six-step cycle)¶

14.A — Six-step cycle motivation (D12)¶

14.B — assert_posterior_agreement¶

14.C — assert_adjoint_calibrated¶

14.D — simulation_based_calibration (Talts et al.)¶

Part 15 — Orchestration (pipekit composition)¶

15.A — VarDACycle for sequential DA¶

15.B — VarSmootherCycle for retrospective windows¶

15.C — obs_source operators¶

15.D — Composing with pipekit.Sequential¶

Part 16 — Training at scale (pipekit-train)¶

16.A — train_step / amortized_train_step primitives¶

16.B — pipekit_train.Loss adapter¶

16.C — TrainingLoop integration¶

16.D — Checkpoints & sweeps¶

Part 17 — Performance¶

17.A — jit basics¶

17.B — eqx.filter_vmap batching¶

17.C — Memory profiling adjoint variants¶

17.D — When to unroll vs scan¶

Part 18 — Debugging¶

18.A — Cost decomposition (decomposed_loss)¶

18.B — Inspecting residuals¶

18.C — Common failure modes¶

Part 19 — Extending vardax¶

19.A — Writing your own AnalysisStep¶

19.B — Protocol-conformance tests¶

19.C — Publishing as a plugin¶

Part 20 — Applied (cross-link)¶

20.A — Lorenz benchmarks (→ projects/assimilation/)¶

20.B — SSH satellite interpolation¶

20.C — Methane single-overpass retrieval¶

20.D — Multi-instrument atmospheric fusion¶

Status snapshot¶

This is a curriculum scaffold. Each row’s Source column reflects what already exists:

• R rows are notebooks that already live in projects/assimilation/notebooks/ (the L63 / L96 / L96-2L tutorials shipped with the assimilation project).
• V / G / K rows are notebooks that exist upstream in the vardax / gaussx / pipekit repos (none at the time of writing — every such row currently reads —, but the legend documents what those tags mean when upstream notebooks land).
• — rows are gaps to be filled by follow-up PRs.

Follow-up PRs will populate projects/assimilation/notebooks/vardax/<section>/<name>.ipynb (a new vardax/ subdirectory under the existing notebooks tree, alongside 00_lorenz63_setup.md and friends) and flip the Source column from — to R. Pure-prose chapters that port math may land as .md instead of .ipynb.

v1 priority (~6 sessions, ~35 notebooks): Parts 0–6 + 13 + 14 + 16 + the L63 entries that already exist in projects/assimilation/. Foundations + components + half the seven methods + posterior + validation + training.

v2 priority: Parts 7–12 dynamics + adjoints depth; Parts 15 cycling demo; Part 17 performance; Part 19 extension; Part 20 applied demos.

v3 priority: stub heads (11.C, 11.D) flipped to real when gauss_flows / diffrax reverse-SDE land; Part 9.C CVT depending on gaussx.sqrt; somax / plumax forwards (2.C / 2.D / 2.E).