radtrax: design and roadmap

What radtrax is for¶

TL;DR. One forward model, every derivative of it, and four retrievals that need them.

Problem. A retrieval inverts $y = F(x) + \varepsilon$ for the state $x$ (trace-gas profiles, temperature, surface, instrument), and every serious inversion needs $K = \partial F / \partial x$: optimal estimation iterates on it Rodgers, 2000, a matched filter’s target spectrum is one column of it, and an emission-rate fit chains it with a plume model. The heritage codes either ship $K$ by hand-derived linearisation (VLIDORT Spurr, 2006) or not at all (MODTRAN tables). radtrax’s one bet is to make $\{\tau_l, \omega_l, \chi_{l,k}, \text{surface}\}$, the per-layer optical depth, single-scattering albedo, phase-function moments and surface kernels, the interface between spectroscopy and solvers, and to get $K$ through it by jax.jacrev. Its uses, in the order the roadmap reaches them:

1. Matched-filter targets for EMIT, PRISMA, EnMAP and AVIRIS-NG: the unit absorption spectrum $t = \partial \ln y / \partial \Delta X$ for a scene’s geometry and water vapour, replacing a MODTRAN table; it is the $t$ of the matched filter in radtrax · 1 and Report 16.

2. Optimal-estimation retrievals of $X_{\text{CH}_4}$, $X_{\text{H}_2\text{O}}$, surface and spectral shift, with $K$ from autodiff.

3. Multispectral retrievals (Sentinel-2, Landsat): the band tables and the multi-pass family of radtrax · 2.

4. Emission rates by maximum likelihood: plumax maps an emission $e$ to a column, radtrax maps the column to radiance, and one gradient gives $\partial \text{loss} / \partial e$; the IME and the cross-sectional flux of plumax · 1 become its starting point.

What exists today¶

TL;DR. Most of v0 is in plumax already, in NumPy, on one slab.

The package boundaries inventory lists every module and its size; the table keeps what radtrax changes in each.

The pieces radtrax needs are present; what is missing is the layer dimension, a single Equinox forward model they share, and everything past Beer–Lambert.

The architecture¶

TL;DR. Five stages, typed pytrees between them, and protocols where the alternatives plug in.

AtmosphericState      p, T, z per layer; VMR per gas         (L,), dict
      │  OpacityProvider: σ(ν; p_l, T_l) per gas, Σ n x σ Δz
      ▼
OpticalProperties     τ (L, G), ω (L, G), χ (L, G, K)        ← K flows here
      │  + Surface: Lambertian | RossLi | RPV | CoxMunk
      │  RTSolver: BeerLambert | SingleScatter | TwoStream
      │            | DiscreteOrdinates(N) | DoublingAdding
      ▼
Radiance              L(ν) at the top of the atmosphere       (G,)
      │  Instrument: SRF, spectral shift, FWHM, PSF, GSD
      ▼
Observation           y_b per band; NoiseModel → Σ_y          (B,)
      │  Retrieval: OE (optimistix LM / GN) | MAP (optax)
      │             | Bayesian (numpyro) | matched filter
      ▼
State estimate        x̂, Ŝ, averaging kernel

Four rules come out of the note and the checks below:

1. Modules carry state; functions carry transforms. A table, a line list, an SRF or a noise covariance is an eqx.Module; a Voigt profile, a Chapman factor or a delta-M scaling is a function.

2. Arrays are leaves. A JAX array in a static field makes Equinox warn and breaks hashing under filter_jit. Grids, tables and SRFs stay leaves; what is not differentiated is chosen with eqx.partition or jax.lax.stop_gradient.

3. Protocols at the seams. OpacityProvider, RTSolver, Surface, Instrument and NoiseModel are abstract Equinox modules (eqx.AbstractVar and eqx.abstractmethod on the base, concrete fields on the subclass), so filter_jit compiles each combination.

4. vmap outside jacrev. A swath is a filter_vmap over pixels of a per-pixel Jacobian, never a Jacobian of a batch.

The proposed layout:

radtrax/
├── atmosphere/   AtmosphericState, AFGL profiles, hydrostatics
├── spectroscopy/ tables (v0), line-by-line (v2): partition functions,
│                 line strengths, Voigt, continua, CIA
├── optics/       OpticalProperties, Rayleigh, HG and Mie, delta-M
├── surface/      Lambertian, kernel BRDFs, Cox–Munk
├── solvers/      BeerLambert, SingleScatter, TwoStream, ...
├── instrument/   SRF, shift and FWHM, PSF, GSD, noise
├── retrieval/    optimal estimation, matched filter, multi-pass
└── geo/          GeoTensor wrappers, sensor SRF tables  [geo] extra

The v0 slice, run¶

TL;DR. A layered Equinox forward model reproduces plumax’s slab exactly, and shows what the slab gets wrong.

Problem. The smallest radtrax that is worth having is the note’s v0: an (ln p, T) table lookup per layer, Beer–Lambert along the two-way path, and a band integration with a spectral shift. The notebook writes it with every array a leaf:

1
2
3
4
5
6
7
8
9
10
11
12
13
# Shapes: L = 20 layers, G = 14,000 wavenumbers, B = 7 bands (8.5 nm)
# σ_l = bilinear(σ; ln p_l, T_l);  τ_l = σ_l n_l x_l Δz_l       (L, G)
tau = TableOpacity(log_p, t_grid, sigma)(atm)    # proposed (radtrax)
# L = (μ₀/π) A exp(−m Σ_l τ_l),  m = 1/μ₀ + 1/μ                  (G,)
rad = BeerLambert()(tau, albedo, mu0, mu)        # proposed (radtrax)
# y_b = Σ_g f_bg L(ν_g + δ)                                      (B,)
y = Instrument(srf, nu, shift)(rad)              # proposed (radtrax)
# one slab against plumax: L and ∂L/∂x                           (G,)
ref = forward_nonlinear(ds, nu, T_K=T0, p_atm=P0, vmr=X0, ...)  # today
# → radiance to 1.3 × 10⁻¹⁵, Jacobian to 6 × 10⁻¹⁶, relative  user
# K = ∂y/∂(ΔX₀, A, δ) per pixel, 256 pixels        (256, 3) → (256, 7, 3)
K = eqx.filter_vmap(jax.jacrev(y_of))(thetas)    # user
# → finite differences to 2.3 × 10⁻⁸; vmap equals a loop exactly

The layered model reduces to plumax’s slab, radiance and Jacobian alike, to rounding; the Jacobian that plumax derives by hand comes out of jax.jacrev unchanged. The layers are not decoration. The same 1.9 ppm column ($4.1\times10^{19}$ molecules/cm²) through 20 one-kilometre layers and through one slab at the surface’s pressure differ by up to 19 % in a band’s optical depth, because the slab broadens every line as if all the methane sat at 1 atm; a slab at the column’s pressure-weighted mean differs by up to 3 %. plumax’s forward model and its observation operator are both single-slab, which is the first thing v0 changes.

The roadmap¶

TL;DR. Six releases, each a working retrieval, each adding one fidelity axis, never a refactor and a feature together.

Each release in Table 1 ships the regression set of the one before; the gradient check, autodiff against central finite differences on a reduced state, is the one test every release runs.

Differentiability, component by component¶

TL;DR. Most of radtrax differentiates as written; four places do not, and the stack has some of the fixes.

Problem. Reverse mode through jnp code is exact where the code is smooth. The places it is not are known from the heritage codes, and the notebook reproduces each one:

1
2
3
4
5
6
7
8
9
10
11
# (a) float32 on a wavenumber grid, near 4300 cm⁻¹
np.spacing(np.float32(4300.0))   # → 4.9e-4 cm⁻¹ (Doppler HWHM 6.6e-3)
# (b) M(μ) = (e^{−τ/μ₀} − e^{−τ/μ}) / (1/μ − 1/μ₀) at μ = μ₀  () → ()
multiplier_naive(mu0)            # → nan, and its gradient nan   user
multiplier_safe(mu0)             # Taylor branch, double-where   user
# → τ e^{−τ/μ₀} exactly; slope = finite differences to 10⁻⁴
# (c) ∂/∂A of eigenvectors, repeated eigenvalue        (3, 3) → (3, 3)
jax.grad(first_eigvec_sum)(diag(1, 1, 2))      # → not finite    user
jax.grad(first_eigvec_sum)(diag(1, 1.001, 2))  # → 500, as 1/gap user
# (d) block-tridiagonal solve, 30 blocks of 8          (240,) → (240,)
gx.solve(gx.BlockTriDiag(D, A), b)   # → dense solve to 5e-16   today

1. Precision. Line-by-line opacity runs in float64. In float32 the wavenumber grid itself is coarse: $4.9\times10^{-4}$ cm⁻¹ at 4300 cm⁻¹, 7 % of the Doppler half-width, so line centres move by a visible fraction of a line.

2. Small denominators. The discrete-ordinates streaming multipliers are 0/0 at $\mu = \mu_0$, in value and in gradient. VLIDORT switches to a Taylor series within $\varepsilon \approx 10^{-3}$ Spurr, 2006; in JAX that branch needs the double-where (a safe denominator in the branch not taken) or jnp.where still propagates the NaN into the gradient.

3. Repeated eigenvalues. jnp.linalg.eigh gradients of eigenvectors scale as one over the eigenvalue gap and are not finite at a repeated one. Discrete ordinates’ reduced eigenproblem Stamnes et al., 1988 needs a custom_jvp that differentiates the invariant subspace, not the individual vectors.

4. The boundary-value solve. Discrete ordinates ends in one block-tridiagonal system per Fourier component; implicit differentiation through it costs one more solve. gaussx’s BlockTriDiag solves that structure in $O(N d^3)$ and matches a dense solve to $5\times10^{-16}$, but its band is symmetric ($A_k$ below, $A_k^\top$ above), and with symmetric=False it falls back to a dense LU. The boundary-value matrix is not symmetric, so v4 needs a general block-tridiagonal solve, which belongs in gaussx.

Fixed points (successive orders, SHDOM-like iterations) and the Gauss–Newton and Levenberg–Marquardt loops of optimal estimation use optimistix (FixedPointIteration, GaussNewton, LevenbergMarquardt, with ImplicitAdjoint), which plumax’s assimilation already depends on; L-BFGS for MAP retrievals is optax.lbfgs.

Dependencies and boundaries¶

TL;DR. radtrax depends on GeoNumerics and nothing above it; geo carriers arrive through an extra; plumax imports radtrax.

What to inherit from the landscape¶

The note surveys the field; the table keeps what each code contributes to radtrax’s design. Entries marked † were checked for this page (the class names against ExoJAX’s documentation, the status against JAXopt’s README, the papers against their publication records); the rest are the note’s.

Open questions¶

1. Where does the 3D-Var retrieval live? plumax’s assimilation is radtrax’s v1 driver in all but name; the package boundaries page asks whether it belongs in radtrax or in vardax, and v1 needs the answer.

2. Tables or line-by-line as the default? v0’s tables are fast and differentiable in the state; v2’s line-by-line is differentiable in the line parameters too, which is what HITRAN-uncertainty propagation needs. Whether retrievals default to v2 once it exists is a cost question.

3. How general must the block-tridiagonal solve be? A non-symmetric BlockTriDiag in gaussx serves discrete ordinates and any other two-point boundary-value problem; it is the one GeoNumerics addition v4 needs.

4. Which sensor first for v1? EMIT has operational products to compare against; Sentinel-2 has the multi-pass family already written down.

Upcoming work¶

Table 2 lists the steps, in order.

References¶

• The design note, “A differentiable JAX/Equinox RTM for methane and multi-gas remote sensing — design and staged roadmap” Johnson, 2026.

• Korkin et al., the practical guide Korkin et al., 2022; ExoJAX Kawahara & others, 2022; VLIDORT Spurr, 2006; DISORT Stamnes et al., 1988; Rodgers, inverse methods Rodgers, 2000.

• PyMieDiff, arXiv 2512.08614; Ray-trax, arXiv 2511.09389; JAXopt’s README (no longer maintained).

• plumax, gaussx, geotoolz, georeader, finitevolX, spectraldiffx, pyrox and somax at the revisions in the header.