Forward model: full PDE solved on a grid. Uses finitevolX for spatial discretisation and spectraldiffx for spectral / periodic-domain operators.
This is the gold-standard physics tier: full mass conservation, arbitrary wind fields, support for chemistry and deposition. Also the most expensive — emulators (Step 3) are essential, not optional.
(1) Simple model¶
Conservation equation¶
Operational implementations carry mixing ratio (kg/kg), not mass density, because that’s what the satellite observes and what WRF couples to. The conservation form is then:
with the dry-air density from the met field. When is approximately constant over the integration window (typical at fixed altitude in a regional basin) this reduces to the textbook . Commit to the full -weighted form for any work crossing significant pressure-altitude variation.
— wind from WRF / ERA5 / HRRR.
— eddy diffusivity. Commit to K-theory from MO similarity in the PBL ( with stability correction; [monin1954,stull1988]) plus Smagorinsky in the free troposphere. Source: MO prereqs + WRF TKE.
— source field, the unknown we invert for. Per-cell, time-resolved.
— first-order chemical loss. Negligible for CH₄ at <1-day scales; included for portability to CO/HCHO.
Initial and boundary conditions¶
Both are load-bearing for inversion accuracy and currently underspecified in operational code:
Spatial / temporal discretisation¶
Cell-centred FV with flux-limited advection via finitevolX; explicit RK or IMEX for stiff diffusion; spectraldiffx for periodic / global domains.
Build order within Tier III¶
1D column model. Validates numerics — diffusion against analytical Gaussian, advection against cosine pulse / Burgers’.
2D horizontal slab. What most satellite work needs (vertical column already integrated by AK).
Full 3D with vertical layers. Needed when injection altitude matters or when assimilating multi-layer in-situ data alongside columns.
(2) Model-based inference¶
Forward observation operator¶
Same column + AK pipeline as Tiers I and II:
is the observation operator at time . The 4D-Var cost below uses ; the column + AK implementation is shared with Tier I and the prereqs.
4D-Var cost — three terms¶
Three terms — source background, IC background, and time-summed observation mismatch. The IC term drops out only if you commit to a long warm-up that pins . Treating observations as a single instantaneous mismatch (no time index) is incompatible with multi-hour assimilation windows.
Likelihood model¶
— heteroscedastic per-pixel from the L2 retrieval-error map at time .
— representation error (model-vs-observation footprint mismatch). Diagonal addition; rises with terrain complexity and at coarse-instrument boundaries. Don’t omit — naive overweights observations and produces overconfident posteriors.
Block-diagonal across overpasses; cross-time correlation only within met decorrelation scale.
Prior on — spatially correlated, sign-constrained¶
Same structure as Tier II:
Prior specification on the source field .
| Choice | Form | Notes |
|---|---|---|
| Mean | from emission inventory | EDGAR / GFEI / EPA per-cell median |
| Covariance | Matérn-3/2, km | spatial regulariser; tuneable or hierarchical |
| Positivity | (lognormal) | non-negative emissions; conjugate when linearised |
| BC scaling | per-face | absorbs lateral-BC bias |
Adjoint¶
The adjoint of the transport equation (backward in time, conservative form) is:
The term is conservative (matches the forward ); the doc previously had , which is equivalent only for divergence-free and is incorrect for compressible WRF winds. JAX computes the discretised adjoint exactly via reverse-mode autodiff through the FV solver — no hand-derived adjoint code, no separate adjoint-correctness derivation. The mathematical form above is for reading, not implementation.
Incremental 4D-Var (the operational default)¶
Linearise around the current iterate , solve the linear inner minimisation, update outer iterate:
uses the tangent linear of the FV solver, trivially built via jax.linearize. Inner solves are quadratic in → CG or Lanczos via gaussx. Cuts cost by 1–2 orders of magnitude vs. fully-nonlinear 4D-Var. This is the default; full nonlinear is a sanity check.
Control-variable transform¶
Direct optimisation in -space with is infeasible — for a grid is entries. Standard fix:
materialised via the structure-aware Cholesky in gaussx (Kronecker structure for separable correlation). gaussx has no covariance-kernel functions: the spatial Matérn matrix is built in plumax (matern32_covariance), while the temporal axis can use gaussx.MaternSDE + MarkovGaussian, the state-space Matérn with solves. This is the load-bearing trick that makes 4D-Var tractable; should be explicit in assimilation/control.py and on the API.
Posterior covariance — three options¶
4D-Var alone gives MAP; downstream Tier V needs the posterior:
Gauss–Newton Hessian inversion. evaluated at MAP. Tractable for moderate grids via
gaussxKrylov.Laplace around MAP. Sample from ; cheapest path.
En4D-Var. Ensemble around MAP gives sample covariance; couples to
filterax. Best when the posterior is non-Gaussian.
The posterior export to Tier V.A is via the same adapter pattern as Tier I/II — see the posterior-export module.
Cost / performance¶
For a grid, 24-hour assimilation window, with incremental 4D-Var: ~minutes per outer iteration on CPU, seconds on GPU. Inner CG iteration is .
(3) Model emulator¶
Full 3D FV transport is expensive: state, repeated time integration. Emulator is essential here, not optional like at Tier I.
State variable — commit to 2D column for v1¶
Most satellite inversion needs only the column-integrated XCH₄, not the 3D field. Default emulator state is 2D column ; full 3D is a v2 escalation when multi-layer in-situ assimilation enters scope. Order-of-magnitude difference in training cost.
Architecture options¶
CNN UNet. Honest baseline for non-stationary terrain. No translation-invariance assumption.
Graph-network operator. Right structural fit for unstructured analysis grids and basin-scoped problems.
Fourier Neural Operator (FNO). Resolution-agnostic and natural fit with
spectraldiffxphilosophy — but assumes translation invariance in the kernel. Works well for periodic / homogeneous problems; breaks for urban basins with terrain. Use only when translation invariance is plausible.Neural ODE. iterated. Simple but suffers from long-horizon drift.
Reduced-order model (ROM). Proper orthogonal decomposition on simulation snapshots → low-rank basis; learn projected dynamics via Galerkin or DMD. Great when dynamics live on a low-dim manifold.
Training data¶
Active-learning over uniform climatology binning. Sample WRF / ERA5 climatology adaptively — the emulator’s residual error map drives where to run the next FV simulation. Reaches operational accuracy with 100–300 runs vs. ~1000 for uniform sampling.
Sample met conditions from the operational distribution (facility locations of interest, overpass times) rather than a uniform-bin climatology — same critique as Tier II’s emulator.
Emulator adjoint must be calibrated¶
Backprop through a trained emulator gives some gradient — whether it matches the true PDE adjoint is empirical.
(4) Emulator-based inference¶
Replace the FV integrator with the FNO / neural ODE in the 4D-Var loop:
Gradient via autodiff through the emulator (not the PDE solver). Both are
jax.grad-able; only the cost-per-iteration changes.Orders of magnitude faster per iteration (typically 100–1000×).
Adjoint validation: emulator gradient ≈ FV adjoint (Step 3 calibration test).
Posterior validation: posterior from emulator-based 4D-Var ≈ posterior from adjoint 4D-Var on the same observations. If they don’t agree, the emulator is biased — diagnose before trusting downstream.
(5) Amortized inference (predictor)¶
Output grid commitment¶
is on the inversion grid (~1–10 km, basin-scoped). Predictor outputs a fixed-size 2D field per basin tile per time slice. Per-instrument predictor heads, dispatched by instrument_id — same pattern as Tier II.
Irregular temporal sampling¶
Satellite passes are intermittent (1–3 per day per instrument over a basin) with gaps. Input is not regularly sampled — are observation times, not a fixed grid. Architecture: set-transformer with time-encoding over irregular passes, fused with a regular-grid encoder over from the met field. Vanilla seq2seq / ConvLSTM assumes regular sampling and will silently underperform.
Posterior over the spatial source field¶
Conditional flow over images vs. score-based diffusion — same trade-off as Tier II:
gauss_flowsis currently 1D-only; 2D coupling layers are a multi-month extension.Score-based diffusion is the safer v1 path for the spatial posterior.
Context conditioning via FiLM / hypernet primitives in
pyrox.nn— same pattern as Tier I/II.
(6) Improve¶
Multi-species coupling. Add CO and CO₂ tracers; their source ratios constrain CH₄ source attribution (e.g. fossil vs. agricultural).
Adaptive grid refinement. Refine near sources, coarsen elsewhere.
finitevolXmay need primitives for this.Online DA. Assimilate observations as they arrive (rolling-window 4D-Var, or
filterax’s ensemble Kalman smoother).Sub-grid plume reconstruction. When a source is sub-grid, the FV solver smears it; pair Tier III with a Tier I plume in the near-field for better point-source representation.
Hierarchical Matérn length-scale. Promote to a hyperparameter with its own posterior — let the data choose the regularisation scale (mirrors Tier II).
Module layout¶
Tier III module layout — step, concern, target module, status.
| Step | Concern | Module | Status |
|---|---|---|---|
| 1 | FV grid | les_fvm/grid.py | ✓ |
| 1 | Advection | les_fvm/advection.py | ✓ |
| 1 | Diffusion | les_fvm/diffusion.py | ✓ |
| 1 | Source injection | les_fvm/source.py | ✓ |
| 1 | Boundary conditions | les_fvm/boundary.py | ✓ |
| 1 | Time integration | les_fvm/dynamics.py, simulate.py | ✓ |
| 1 | Eddy diffusivity | les_fvm/diffusion.py (EddyDiffusivity, pg_eddy_diffusivity) | 🚧 PG ✓; MO + Smagorinsky ☐ #94 |
| 1 | Column + AK pipeline | reuse gauss_plume.observation from Tier I | ☐ #83 |
| 2 | 4D-Var cost (prior + time-summed obs) over the FV transport | les_fvm/fourdvar.py (FourDVarProblem) | ✓ |
| 2 | Differentiable forward (emission → column-obs series) + exact adjoint via AD | les_fvm/fourdvar.py (EulerianForward4DVar) | ✓ |
| 2 | Column observation operator H_t | les_fvm/fourdvar.py (ColumnObservationOperator) | ✓ |
| 2 | Likelihood + temporal prior on the emission series (Matérn-3/2, R = R_retr + R_repr) | reuse lagrangian.inversion.matern32_covariance | ✓ |
| 2 | Spatial source prior (per-cell B) | SpatialControl in fourdvar.py — v1 control is a scalar emission rate at a known source | ☐ #91 |
| 2 | Control vector + whitening transform | les_fvm/fourdvar.py (Cholesky whitening of B) | ✓ |
| 2 | 4D-Var solver (L-BFGS in whitened space) | les_fvm/fourdvar.py (solve_4dvar) | ✓ |
| 2 | Incremental (Gauss-Newton inner) 4D-Var solver | assimilation/solve.py (run_gauss_newton, run_dual_psas exist; not wired to the FV forward) | 🚧 #93 |
| 2 | Background (, , BC scaling) | assimilation/background.py (B builders exist; not in the 4D-Var cost) | 🚧 #92 |
| 2 | Diagnostics | assimilation/diagnostics.py | 🚧 |
| 2 | Posterior covariance (Hessian / Laplace / En4D-Var) | les_fvm/fourdvar.py (posterior_covariance, laplace_sample) | 🚧 Gauss-Newton Laplace ✓; En4D-Var ☐ |
| 2 | Posterior export → Tier V | samples / prior_logpdf on FourDVarResult | ☐ #107 |
| 3 | Emulator (UNet / GNN / FNO) | plumax.les_fvm.emulator | ☐ #95 |
| 3 | Emulator-adjoint calibration harness | plumax.les_fvm.emulator_adjoint_test | ☐ #95 |
| 5 | Sequence predictor (set-transformer) | plumax.les_fvm.predictor | ☐ |
| 6 | Multi-species coupling | plumax.les_fvm.multispecies | ☐ |
Validation strategy¶
1D diffusion. Initial Dirac → at time , solution is Gaussian with variance . Check error vs. analytical.
1D advection. Cosine pulse advected at constant → after one period, recover initial condition. Tests upwind / flux-limiter consistency.
CFL stability. Artificially exceed CFL → expect blow-up; stay within → expect stability. Catches silent flux-limiter regressions.
Mass conservation — three regimes.
No sources, periodic BCs: exact to floating-point.
With sources, no deposition: summed over the window.
With deposition: include term.
Adjoint correctness. JAX
vjpof the discrete forward should satisfy for random . Cheap, catches differentiation bugs.Tier I limit. For a single point source in a uniform wind with constant , the steady-state FV solution should match the Gaussian-plume formula at downwind distances much greater than the grid spacing.
Emulator-adjoint calibration. Emulator-autodiff gradients vs. FV-autodiff gradients on a held-out met set, relative error in operator norm. Hard test — failure means Step 4 inversion is biased.
Emulator OOD generalization. Train on one met regime (e.g. summer Permian), evaluate on another (winter Permian, or a different basin). Resists overfit-to-training-distribution.
Real-data benchmark. Compare 4D-Var output to existing CAMS / GEOS-Chem-Adjoint inversions on a published time window (e.g. [maasakkers2023ghgi,jacob2022quantifying] Permian inversions). Posterior credible interval should overlap the published estimate. Without this, the inversion is a synthetic exercise.