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 $\vect{y} = h(\vect{x}) + \gvect{\varepsilon}$ for the state $\vect{x} \in \mathbb{R}^{D_x}$ (trace-gas profiles, temperature, surface, instrument), where $\vect{y} \in \mathbb{R}^{D_y}$ is the measurement and $\gvect{\varepsilon} \in \mathbb{R}^{D_y}$ is its noise. Every serious inversion needs the Jacobian $\mat{J}_h = \partial h / \partial \vect{x} \in \mathbb{R}^{D_y \times D_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 $\mat{J}_h$ 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 interface between spectroscopy and solvers (the optical depth, single-scattering albedo and phase-function moments $k$ of layer $l$, all scalars, and the surface kernels), and to get $\mat{J}_h$ through it with jax.jacrev. Its uses, in the order the roadmap reaches them, are:

1. Matched-filter targets for EMIT, PRISMA, EnMAP and AVIRIS-NG: the unit absorption spectrum $\partial \ln \vect{y} / \partial \Delta\mathrm{VMR} \in \mathbb{R}^{D_y}$ for a scene’s geometry and water vapour, in place of a MODTRAN table. It is the target $\vect{J}_h$ (the one column of $\mat{J}_h$ for the methane enhancement) of the matched filter in radtrax · 1 and Report 16. [CHECK: radtrax · 1 defines $\vect{J}_h$ as the derivative of the reflectance with respect to $\Delta\Omega$, which is this log derivative times $\gvect{\mu}_y$ and a conversion from $\Delta\mathrm{VMR}$ to $\Delta\Omega$, not the unit absorption spectrum itself.]

2. Optimal-estimation retrievals of $\mathrm{XCH_4}$, $\mathrm{XH_2O}$, surface and spectral shift, with $\mat{J}_h$ 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 rate $Q$ to a column, radtrax maps the column to radiance, and one gradient gives $\partial \mathcal{L} / \partial Q$ of the loss $\mathcal{L}$. 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, and the table below keeps what radtrax changes in each.

We have the pieces radtrax needs. What is missing is the layer dimension, a single Equinox forward model which 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         (N_L,), dict
      │  OpacityProvider: σ(ν; p_l, T_l) per gas, Σ n VMR σ Δz
      ▼
OpticalProperties     τ, ω (N_L, N_ν), χ (N_L, N_ν, N_χ)     ← J_h flows here
      │  + Surface: Lambertian | RossLi | RPV | CoxMunk
      │  RTSolver: BeerLambert | SingleScatter | TwoStream
      │            | DiscreteOrdinates(N) | DoublingAdding
      ▼
Radiance              L(ν) at the top of the atmosphere       (N_ν,)
      │  Instrument: SRF, spectral shift, FWHM, PSF, GSD
      ▼
Observation           y_b per band; NoiseModel → Σ_ε          (N_B,)
      │  Retrieval: OE (optimistix LM / GN) | MAP (optax)
      │             | Bayesian (numpyro) | matched filter
      ▼
State estimate        x̂, Σ_x|y, averaging kernel

We take four rules from the note and from 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, whereas 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. So grids, tables and SRFs stay leaves, and we choose what is not differentiated 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, and never a Jacobian of a batch.

The layout we propose:

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. We write it in the notebook with every array as a leaf:

1
2
3
4
5
6
7
8
9
10
11
12
13
# Shapes: N_L = 20 layers, N_ν = 14,000 wavenumbers, N_B = 7 bands (8.5 nm)
# σ_l = bilinear(σ; ln p_l, T_l);  τ_l = σ_l n_l VMR_l Δz_l     (N_L, N_ν)
tau = TableOpacity(log_p, t_grid, sigma)(atm)    # proposed (radtrax)
# L = (μ₀/π) A exp(−AMF Σ_l τ_l),  AMF = 1/μ₀ + 1/μ             (N_ν,)
rad = BeerLambert()(tau, albedo, mu0, mu)        # proposed (radtrax)
# y_b = Σ_g f_bg L(ν_g + δ)                                     (N_B,)
y = Instrument(srf, nu, shift)(rad)              # proposed (radtrax)
# one slab against plumax: L and ∂L/∂x                          (N_ν,)
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
# J_h = ∂y/∂(ΔVMR₀, A, δ) per pixel, 256 pixels    (256, 3) → (256, 7, 3)
j_h = 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 to rounding, for the radiance and the Jacobian alike, and the Jacobian that plumax derives by hand comes out of jax.jacrev unchanged. However, the layers are not decoration. The same 1.9 ppm column ($4.1 \times 10^{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, whereas 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, and 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 where 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 \times 10^{-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 once $\vert \mu - \mu_0 \vert$ is within about 10^-3 Spurr, 2006. In JAX that branch needs the double-where (a safe denominator in the branch not taken), or else jnp.where still propagates the NaN into the gradient.

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

4. The boundary-value solve. Discrete ordinates ends in one block-tridiagonal system per Fourier component, and implicit differentiation through it costs one more solve. gaussx’s BlockTriDiag solves that structure in $O(N d^3)$ (for $N$ blocks of size $d$) and matches a dense solve to $5 \times 10^{-16}$. However, its band is symmetric ($\mat{A}_k$ below the diagonal, $\mat{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.

For the fixed points (successive orders, SHDOM-like iterations) and the Gauss–Newton and Levenberg–Marquardt loops of optimal estimation we use optimistix (FixedPointIteration, GaussNewton, LevenbergMarquardt, with ImplicitAdjoint), which plumax’s assimilation already depends on, and 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, and the table keeps what each code contributes to radtrax’s design. We checked the entries marked † for this page (the class names against ExoJAX’s documentation, the status against JAXopt’s README, the papers against their publication records), and 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, whereas v2’s line-by-line is also differentiable in the line parameters, 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, and it is the one GeoNumerics addition that v4 needs.

4. Which sensor first for v1? EMIT has operational products to compare against, whereas 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.