filterax Tutorial Master List

A reconciled, exhaustive curriculum spanning what currently exists in filterax, gaussx, pipekit, and research_notebook, plus gaps surfaced from the filterax public API, open GitHub issues, and the design_docs/ series under filterax/docs/design_docs/. Goal: the most complete ensemble-DA / data-assimilation tutorial sequence we could ship.

GP / SVI tutorials live in ../../gaussian_processes/TUTORIAL_MASTER_LIST.md. Cross-listed items (state-space GPs, ensemble VI, structured Gaussians, sigma points, shrinkage) are flagged 🔁.

Legend — Source columns:

• F = exists in filterax (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 (GP master list)

Refs column: gh#N = open GitHub issue · dd:path = filterax docs/design_docs/<path> · api:foo = filterax exported symbol.

Curriculum at a glance¶

• Part 0 — Bayesian Filtering & DA Foundations
  • 0.A — The filtering problem
  • 0.B — Linear-Gaussian Kalman filter
  • 0.C — The forecast → analysis → inflate cycle
  • 0.D — Variational DA (3D/4D-Var) contrast
  • 0.E — Information vs covariance form
• Part 1 — Layer 0 Primitives
  • 1.A — Ensemble statistics
  • 1.B — Gain & innovation
  • 1.C — Likelihood & innovation statistics
  • 1.D — Perturbations
  • 1.E — Localisation kernels
  • 1.F — Inflation primitives
  • 1.G — Patches & domain decomposition
• Part 2 — Layer 1 Sequential Filters
  • 2.A — Stochastic / perturbed-observation
  • 2.B — Deterministic square-root family
  • 2.C — Localised
  • 2.D — Symmetry-breaking variants
  • 2.E — Parametric (non-ensemble)
  • 2.F — Selection guide
• Part 3 — Layer 2 Forecast-Analysis Loops
  • 3.A — Protocols & extension points
  • 3.B — L2 model walkthroughs
  • 3.C — Inflator integration
  • 3.D — Stochastic key handling
• Part 4 — Backward-Pass Smoothers
  • 4.A — Sequential smoothers
  • 4.B — Square-root smoothers
  • 4.C — Iterative smoothers
  • 4.D — Selection & memory trade-offs
• Part 5 — Ensemble Kalman Processes (Inversion & Sampling)
  • 5.A — Inversion (EKI family)
  • 5.B — Sampling (EKS family)
  • 5.C — Parametric (UKI)
  • 5.D — Regularised / sparse
  • 5.E — Schedulers
• Part 6 — Localisation, Inflation, Calibration
  • 6.A — Why localisation
  • 6.B — Why inflation
  • 6.C — Adaptive variants
  • 6.D — Shrinkage estimators
• Part 7 — Diagnostics & Verification
  • 7.A — Basic spread / RMSE / rank
  • 7.B — Innovation diagnostics
  • 7.C — Reliability & sharpness
  • 7.D — Predictive likelihood
• Part 8 — Differentiable DA
  • 8.A — Theory & gradient stability
  • 8.B — differentiable_assimilate mechanics
  • 8.C — Training patterns
  • 8.D — Memory & remat
  • 8.E — Loss zoo
• Part 9 — optax Integration
  • 9.A — Process transforms
  • 9.B — Composition with optax chains
  • 9.C — Hybrid SGD + EKI
• Part 10 — Sequential Variational Inference
  • 10.A — Foundations
  • 10.B — Particle filters & SMC
  • 10.C — Variational SMC
  • 10.D — Ensemble VI for SSMs
  • 10.E — Amortised / streaming inference
  • 10.F — Sequential VB comparison
• Part 11 — Ecosystem Integrations
  • 11.A — gaussx (structured covariances)
  • 11.B — pipekit (orchestration)
  • 11.C — somax (SDE dynamics)
  • 11.D — geo_toolz / xr_assimilate (xarray)
  • 11.E — plumax (Tier IV)
• Part 12 — Applied Case Studies (research_notebook)
  • 12.A — Canonical DA benchmarks
  • 12.B — Atmospheric & remote sensing
  • 12.C — Inverse problems
  • 12.D — Online / streaming
• Part 13 — Reference Surfaces (Zoo)
  • 13.A — Continuous-time
  • 13.B — Toy dynamical systems
  • 13.C — Hybrid Var-EnKF

Part 0 — Bayesian Filtering & DA Foundations¶

0.A — The filtering problem¶

Key equations / models:

• Prior, dynamics, likelihood factorisation: $p(x_{0:T}, y_{1:T}) = p(x_0)\prod_t p(x_t \mid x_{t-1})\,p(y_t \mid x_t)$
• Filtering target: $p(x_t \mid y_{1:t})$
• Smoothing target: $p(x_t \mid y_{1:T})$, $T > t$
• Forecasting: $p(x_{t+h} \mid y_{1:t})$
• Chapman-Kolmogorov: $p(x_t \mid y_{1:t-1}) = \int p(x_t \mid x_{t-1})\, p(x_{t-1} \mid y_{1:t-1})\,dx_{t-1}$

0.B — Linear-Gaussian Kalman filter¶

Key equations / models:

• Forecast: $\bar x^f_t = M_t \bar x^a_{t-1}$, $P^f_t = M_t P^a_{t-1} M_t^\top + Q_t$
• Analysis: $K_t = P^f_t H_t^\top S_t^{-1}$, $\bar x^a_t = \bar x^f_t + K_t(y_t - H_t \bar x^f_t)$, $P^a_t = (I - K_t H_t) P^f_t$
• Joseph form: $P^a_t = (I - K_t H_t)\, P^f_t\, (I - K_t H_t)^\top + K_t R_t K_t^\top$
• Log-marginal: $\log p(y_t \mid y_{1:t-1}) = -\tfrac{1}{2}\bigl[N_y \log 2\pi + \log|S_t| + v_t^\top S_t^{-1} v_t\bigr]$

0.C — The forecast → analysis → inflate cycle¶

Key equations / models:

• Cycle structure: forecast (apply dynamics) → analysis (assimilate obs) → inflate (counteract sample collapse)
• Ensemble representation: $X \in \mathbb{R}^{N_e \times N_x}$, rows = members
• Sample covariance: $P_e = (N_e - 1)^{-1} (X - \bar X)^\top (X - \bar X)$, rank $\le N_e - 1$

0.D — Variational DA (3D/4D-Var) contrast¶

Key equations / models:

• 3D-Var cost: $J(x) = \tfrac{1}{2}(x - x^b)^\top B^{-1}(x - x^b) + \tfrac{1}{2}(y - H x)^\top R^{-1}(y - H x)$
• 4D-Var window: $J(x_0) = \tfrac{1}{2}\|x_0 - x^b_0\|^2_{B^{-1}} + \sum_t \tfrac{1}{2}\|y_t - H_t M_{1:t} x_0\|^2_{R_t^{-1}}$
• Adjoint vs autodiff: see Part 8

0.E — Information vs covariance form¶

Key equations / models:

• Information matrix $\Lambda = \Sigma^{-1}$, information vector $\eta = \Sigma^{-1} \mu$
• Conjugate update in natural form: $\eta_{t+1} = \eta_t + H^\top R^{-1} y_t$, $\Lambda_{t+1} = \Lambda_t + H^\top R^{-1} H$
• When to prefer information form: GMRF priors, banded Λ, sequential-by-obs updates

Part 1 — Layer 0 Primitives¶

filterax’s pure-function building blocks. Every L1 / L2 algorithm composes from these.

1.A — Ensemble statistics¶

Key equations / models:

• $\bar x = N_e^{-1} \sum_j x^{(j)}$
• $X' = X - \mathbf{1}\bar x^\top$ (rows sum to zero)
• $P_e = (N_e - 1)^{-1} X'^\top X'$ — Bessel-corrected; returned as gaussx.LowRankUpdate
• Cross-cov $C^{xH} = (N_e - 1)^{-1} X'^\top (HX)'$

1.B — Gain & innovation¶

Key equations / models:

• Ensemble Kalman gain: $K = C^{xH}(C^{HH} + R)^{-1}$
• Innovation covariance: $S = C^{HH} + R$ (low-rank update over $R$)
• Innovation: $v = y - \bar H X$

1.C — Likelihood & innovation statistics¶

Key equations / models:

• $\log p(y \mid \text{forecast}) = -\tfrac{1}{2}[N_y \log 2\pi + \log|S| + v^\top S^{-1} v]$
• InnovationStatistics — packaged $v$, $S$, $\log p$ for diagnostics & training

1.D — Perturbations¶

Key equations / models:

• $\epsilon^{(j)} \sim \mathcal{N}(0, R)$, structure-aware via gaussx.root_decomposition
• Diagonal-$R$ fast path: $\epsilon^{(j)}_k = z^{(j)}_k \sqrt{R_{kk}}$
• Determinism: identical key ⇒ identical draws

1.E — Localisation kernels¶

Key equations / models:

• Gaspari-Cohn (compact, $C^4$ at origin, support $[0, 2r]$)
• Gaussian taper: $\exp(-d^2/2r^2)$
• SOAR: $(1 + d/r)\exp(-d/r)$
• Hard cutoff (non-differentiable!) & adaptive (Anderson 2007 / 2012)

1.F — Inflation primitives¶

Key equations / models:

• Multiplicative: $X'_a \leftarrow \lambda X'_a$, posterior cov $\lambda^2 P_a$
• Additive: $X_a \leftarrow X_a + \xi$, $\xi \sim \mathcal{N}(0, Q_\text{add})$
• RTPS (Whitaker-Hamill 2012): $X'_a \leftarrow X'_a [\alpha\, \sigma^f/\sigma^a + (1-\alpha)]$
• RTPP (Zhang 2004): $X'_a \leftarrow \alpha X'_f + (1-\alpha) X'_a$
• Ledoit-Wolf shrinkage: blend sample cov with structured target

1.G — Patches & domain decomposition¶

Key equations / models:

• Patch index sets $P_i \subseteq \{1, …, N_x\}$, obs-to-patch assignment, blending for overlapping patches
• Per-patch ETKF analysis assembled with smooth weighting in overlaps

Part 2 — Layer 1 Sequential Filters¶

Each filter as its own tutorial. Verified against the closed-form Kalman update on a linear-Gaussian problem; the existing tests/test_filters.py baselines are the template.

2.A — Stochastic / perturbed-observation¶

Key equations / models:

• $X^a_j = X^f_j + K(y + \epsilon^{(j)} - H X^f_j)$, $\epsilon^{(j)} \sim \mathcal{N}(0, R)$
• Monte Carlo gain noise scales as $1/\sqrt{N_e}$
• Per-window PRNG key: jr.fold_in(base_key, step)

2.B — Deterministic square-root family¶

Key equations / models:

• ETKF transform precision: $\tilde C = (N_e - 1) I + Y' R^{-1} Y'^\top$
• Rank-$N_y$ spectrum trick (no eigh of degenerate $(N_e, N_e)$ matrix) — fixed in #82
• $W_a = \sqrt{(N_e - 1) \tilde C^{-1}}$ applied to anomalies via $g(\tilde C) v = g(N_e-1) v + U_y \mathrm{diag}(g_\Delta) U_y^\top v$

2.C — Localised¶

Key equations / models:

• R-localisation (Hunt et al. 2007): inflate local $R^{-1}$ by per-obs taper ρ
• Per-grid-point local ETKF, vmapped across $N_x$ points
• Hard cutoff at radius: obs beyond $r$ excluded entirely (Gaspari-Cohn nonzero out to $2r$)

2.D — Symmetry-breaking variants¶

Key equations / models:

• $W_a^\text{rot} = W_a \cdot \Theta$, $\Theta \in O(N_e)$, $\Theta \mathbf{1} = \mathbf{1}$
• Counteracts preferred-direction drift over many cycles

2.E — Parametric (non-ensemble)¶

Key equations / models:

• Cholesky-form mean & covariance propagation
• Marginal likelihood from the innovation sequence
• PSD-preserving across $T$ steps by construction

2.F — Selection guide¶

Part 3 — Layer 2 Forecast-Analysis Loops¶

3.A — Protocols & extension points¶

Key equations / models:

• AbstractDynamics(state, t0, t1) -> state — vmap over members
• AbstractObsOperator(state) -> obs
• AbstractInflator(particles, forecast, **kwargs) -> particles
• AbstractLocalizer(cov, coords) -> cov
• AbstractScheduler.get_dt(state) for processes

3.B — L2 model walkthroughs¶

3.C — Inflator integration¶

3.D — Stochastic key handling¶

Part 4 — Backward-Pass Smoothers¶

4.A — Sequential smoothers¶

Key equations / models:

• Smoother gain: $G_t = C^{af}_{t,t+1}\,(C^{ff}_{t+1})^{-1}$
• Recursion: $X^s_t = X^a_t + G_t(X^s_{t+1} - X^f_{t+1})$
• Dual ensemble-space form: $G_t (X^s_{t+1} - X^f_{t+1}) = D F^\top (F F^\top)^+ A$ — avoids materialising $(N_x, N_x)$ cov

4.B — Square-root smoothers¶

Key equations / models:

• Decompose mean + perturbation; apply symmetric sqrt to anomalies in ensemble space
• $\Lambda = I + K_e (D^\top D - F^\top F) K_e^\top$, $K_e = (F F^\top)^+ F$
• Smoothed perts $X^{s'}_t = \Lambda^{1/2} A_t$ live in the row span of $A_t$

4.C — Iterative smoothers¶

Key equations / models:

• Chen-Oliver IES update with prior anchor: $\theta^j_{i+1} = (1 - \alpha)\,\theta^j_i + \alpha\bigl[\theta^j_0 + K_i(y + \epsilon^{(j)} - G(\theta^j_i))\bigr]$
• $K_i = C^{\theta G}_i (C^{GG}_i + \Gamma_y)^{-1}$ recomputed each iteration

4.D — Selection & memory trade-offs¶

Part 5 — Ensemble Kalman Processes (Inversion & Sampling)¶

5.A — Inversion (EKI family)¶

Key equations / models:

• EKI update: $\theta^j_{n+1} = \theta^j_n + \Delta t_n\, C^{\theta G}_n (C^{GG}_n + \Delta t_n^{-1} \Gamma)^{-1}(y - G(\theta^j_n))$
• Tempered $\Delta t^{-1}\Gamma$ noise as low-rank update over Γ
• Algo-time $\to 1$ collapse: spread $\to 0$, MAP recovered

5.B — Sampling (EKS family)¶

Key equations / models:

• EKS / interacting Langevin: $\theta^j_{n+1} = \theta^j_n + \Delta t\, C^{\theta G}_n (\ldots) + \sqrt{2\Delta t\, C^{\theta\theta}_n}\,\xi^j_n$
• Approaches the Bayesian posterior as $\Delta t \to 0$, $N_e \to \infty$
• Ergodic — spread doesn’t collapse

5.C — Parametric (UKI)¶

Key equations / models:

• Sigma-point propagation: $\{\theta_k\} = \mu \pm \sqrt{(N_p + \kappa)\Sigma}_i$
• Closed-form mean + cov update; no random perturbations
• Exact second-moment match in linear-Gaussian limit

5.D — Regularised / sparse¶

Covered in 5.A (TEKI, SparseInversion); listed here for navigation.

5.E — Schedulers¶

Key equations / models:

• Fixed: $\Delta t_n = \text{const}$
• Data-misfit controller (Iglesias 2016): adapt $\Delta t$ from current misfit, freeze past $\text{algo\_time}=1$
• EKS-stable: $\Delta t = h / (\text{trace}(C^{GG}) + \delta)$ (avoids ensemble blow-up)

Part 6 — Localisation, Inflation, Calibration¶

6.A — Why localisation¶

Key equations / models:

• Spurious correlations: $|C_{ij}| \sim 1/\sqrt{N_e}$ for unrelated state pairs
• Sample-cov rank ≤ $N_e − 1$ in state space; rank-deficient inverse undefined without regularisation

6.B — Why inflation¶

Key equations / models:

• Filter divergence: posterior cov collapses → obs rejected → estimates drift
• Recovery: $\lambda > 1$ multiplicative or $Q_\text{add}$ additive

6.C — Adaptive variants¶

6.D — Shrinkage estimators¶

Part 7 — Diagnostics & Verification¶

7.A — Basic spread / RMSE / rank¶

Key equations / models:

• Ensemble spread: $\sigma_t = \sqrt{(N_e-1)^{-1}\sum_j \|x^{(j)}_t - \bar x_t\|^2 / N_x}$
• RMSE: $\sqrt{\langle (\bar x_t - x^\text{true}_t)^2\rangle_t}$
• Rank-histogram of obs vs ensemble (Talagrand)

7.B — Innovation diagnostics¶

Key equations / models:

• Mahalanobis: $v^\top S^{-1} v \sim \chi^2_{N_y}$ if filter is consistent
• Desroziers (2005): $\mathbb{E}[(y - H\bar x^a)(y - H\bar x^f)^\top] = R$

7.C — Reliability & sharpness¶

Key equations / models:

• Spread-skill: $\sigma_t \approx \text{RMSE}_t$ when the filter is calibrated
• CRPS: $\int_{-\infty}^\infty \bigl(F(z) - \mathbf{1}\{y \le z\}\bigr)^2 dz$
• DFS: $\text{tr}(I - P^a (P^f)^{-1})$
• ESS: $(\sum w^{(j)})^2 / \sum w^{(j)2}$

7.D — Predictive likelihood¶

Part 8 — Differentiable DA¶

See dd:features/differentiable_da.md.

8.A — Theory & gradient stability¶

Key equations / models:

• Pure-function filter ⇒ jax.grad flows for free
• Stochastic vs deterministic filters: perturbed-obs draws inject non-smooth randomness
• eigh degeneracy: $(N_e, N_e)$ transform precision has $N_e − N_y$ repeated eigenvalues → NaN gradient
• Rank-$N_y$ QR-based fix: stable at any $N_e$

8.B — differentiable_assimilate mechanics¶

Key equations / models:

• $X^a_t = \text{EnKF}(X^f_t = f_\theta(X^a_{t-1}), y_t)$ unrolled as jax.lax.scan
• Memory under reverse-mode: $O(T \cdot N_e \cdot N_x)$ → $O(\sqrt{T}\cdot N_e \cdot N_x)$ with checkpoint

8.C — Training patterns¶

Key equations / models:

• Pattern A: $\theta_f$ via $\nabla_{\theta_f} (-\sum_t \log p(y_t \mid \text{forecast}_t))$
• Pattern B: $\theta_H$ via the same loss, gradient flows through obs operator
• Pattern C: ρ, $r_\text{loc}$, $R$ diag via reparameterised log_factor etc.

8.D — Memory & remat¶

Key equations / models:

• jax.checkpoint(step) placed on the scan body
• Binomial checkpointing schedule: $O(\sqrt{T})$ memory at 2-3× compute
• ROAD-EnKF (Chen et al. 2023): $O(1)$ memory, approximate gradient

8.E — Loss zoo¶

Key equations / models:

• NLL: $\sum_t -\log p(y_t \mid \bar x^f_t)$
• MSE: $\sum_t \|v_t\|^2$
• CRPS: $\sum_t \text{CRPS}(\text{ens}_t, y_t)$
• Spread-skill: $\sum_t (\sigma_t - \text{err}_t)^2$

Part 9 — optax Integration¶

9.A — Process transforms¶

Key equations / models:

• Each process exposed as an optax.GradientTransformation: init(params) -> state, update(grad, state, params) -> (updates, state)
• The “gradient” slot is unused; updates come from the ensemble Kalman process internally
• optax.apply_updates(params, updates) advances the parameter mean

9.B — Composition with optax chains¶

9.C — Hybrid SGD + EKI¶

Part 10 — Sequential Variational Inference¶

The bridge into broader filtering / sequential-VI work. Each tutorial sits next to a filterax / pyrox primitive and points at the GP master list where overlap exists.

10.A — Foundations¶

Key equations / models:

• Sequential VI: $q_t(x_t) \approx p(x_t \mid y_{1:t})$ updated as obs arrive
• Recursive ELBO: $\mathcal{L}_t = \mathbb{E}_{q_t}[\log p(y_t \mid x_t)] - \text{KL}(q_t \Vert q_{t-1}^\text{forecast})$
• Connection to Kalman: linear-Gaussian $q$ ⇒ exact Kalman recursion

10.B — Particle filters & SMC¶

Key equations / models:

• Importance sampling + resampling: $w^{(j)} \propto p(y_t \mid x^{(j)}_t)$
• ESS: $N_\text{eff} = (\sum w^{(j)})^2 / \sum (w^{(j)})^2$
• Resampling schemes: multinomial, stratified, systematic

10.C — Variational SMC¶

Key equations / models:

• VSMC ELBO (Naesseth 2018): tighter than IWAE via SMC structure
• Filtering variational objectives (Maddison 2017): per-step lower-bound

10.D — Ensemble VI for SSMs¶

10.E — Amortised / streaming inference¶

10.F — Sequential VB comparison¶

Part 11 — Ecosystem Integrations¶

11.A — gaussx (structured covariances)¶

Key equations / models:

• $S = C^{HH} + R$ as gaussx.LowRankUpdate → Woodbury solve in $O(N_e^2 N_y + N_e^3)$ vs $O(N_y^3)$
• Diagonal / Toeplitz / Kronecker $R$ never densified

11.B — pipekit (orchestration)¶

Key equations / models:

• D11 wrapper pattern: FilterAsAnalysisStep, DynamicsAsForwardModel
• Sequential / Graph / Cycle composition of analysis steps

11.C — somax (SDE dynamics)¶

11.D — geo_toolz / xr_assimilate (xarray)¶

11.E — plumax (Tier IV)¶

Part 12 — Applied Case Studies (research_notebook)¶

12.A — Canonical DA benchmarks¶

12.B — Atmospheric & remote sensing¶

12.C — Inverse problems¶

12.D — Online / streaming¶

Part 13 — Reference Surfaces (Zoo)¶

Explicitly not maintained as core API; lives under zoo/ (gap, planned for #61) for educational / benchmarking use.

13.A — Continuous-time¶

13.B — Toy dynamical systems¶

13.C — Hybrid Var-EnKF¶

References¶