radtrax

Differentiable radiative transfer in JAX, for methane and multi-gas remote sensing.

Write the forward model once, from an atmospheric state to what a sensor measures, $y = F(x)$. Every Jacobian a retrieval needs, $K = \partial F / \partial x$, then comes from automatic differentiation, with no adjoint written by hand: today the matched-filter target and the band tables that invert the multi-pass signals are built from the model and its derivatives. 3D-Var still runs on its single-slab observation operator, and optimal estimation on the layered model is v1 work in progress. Formerly rtmx.

Quick start

uv add "radtrax @ git+https://github.com/jejjohnson/radtrax.git"

import jax

jax.config.update("jax_enable_x64", True)  # float64: float32 shifts line centres

import equinox as eqx
import jax.numpy as jnp
import numpy as np
from jaxtyping import Array, Float

from radtrax import ForwardModel, unit_absorption_target
from radtrax.atmosphere import AtmosphericState, ViewingGeometry, afgl_us_standard
from radtrax.instrument import BandInstrument, SpectralResponseFunction
from radtrax.solvers import BeerLambert
from radtrax.spectroscopy import TableOpacity, load_opacity
from radtrax.surface import Lambertian

# Shapes: L layers, G wavenumbers, B bands, N pixels.

# ── the scene: state x, surface R, geometry g ───────────────────────────────
# x = (p_l, T_l, x_k,l): AFGL US Standard to 20 km, 1 km layers           L = 20
atm: AtmosphericState = afgl_us_standard(z_top_km=20.0, gases=["CH4", "H2O"])
# R(ν) = A, a Lambertian albedo
surface: Lambertian = Lambertian(0.25)  # A = 0.25
# μ₀ = cos θ₀, μ = cos θ: solar and viewing zenith cosines
geometry: ViewingGeometry = ViewingGeometry.from_angles(sza_deg=35.0, vza_deg=10.0)

# ── O: opacity ───────────────────────────────────────────────────────────────
# τ_l(ν) = ∑_k σ_k(ν; ln p_l, T_l) · x_k,l · n_l Δz_l   (k: gas, σ from HITRAN tables)
# σ in cm²/molecule; n_l Δz_l the layer's air column in molecules/cm² (Δz in cm)
opacity: TableOpacity = load_opacity(["CH4", "H2O"])  # build once with radtrax[hitran]
nu: Float[Array, " G"] = opacity.nu_cm  # ν [cm⁻¹]                             (G,)

# ── I: instrument ────────────────────────────────────────────────────────────
# λ = 10⁷ / ν [nm], reversed so it ascends
wl_nm: Float[np.ndarray, " G"] = 1e7 / np.asarray(nu)[::-1]  #                 (G,)
# Gaussian bands across the 2.3 µm CH₄ window: centres λ_b, FWHM 8.5 nm
centers_nm: Float[np.ndarray, " B"] = np.arange(2300.0, 2350.0, 7.5)  #       (B,)
srf: SpectralResponseFunction = SpectralResponseFunction(
    wavelengths_hr_nm=wl_nm,
    band_centers_nm=centers_nm,
    band_widths_nm=np.full(centers_nm.size, 8.5),
    band_names=tuple(f"{c:.1f}" for c in centers_nm),
)
# f_bg: each band's response on ν, rows summing to 1                        (B, G)
instrument: BandInstrument = BandInstrument.from_srf(srf, nu)

# ── S and the model ──────────────────────────────────────────────────────────
# L(ν) = (E₀ μ₀ / π) · R(ν) · exp(−m ∑_l τ_l(ν)),   m = 1/μ₀ + 1/μ
# y_b = ∑_g f_bg · L(ν_g),   i.e.  y = I(S(O(x), R, g))
model: ForwardModel = ForwardModel(opacity, BeerLambert(), instrument)
y: Float[Array, " B"] = model(atm, surface, geometry)  # E₀ = 1               (B,)

# ── what autodiff gives ──────────────────────────────────────────────────────
# t_b = ∂ ln y_b / ∂ΔX at ΔX = 0: the matched-filter target [per mol/m²]
t: Float[Array, " B"] = unit_absorption_target(model, atm, surface, geometry)  # (B,)


# K_n = ∂y/∂θ |_θₙ for θ = (ΔX [mol/m²], A). The model is an argument, so its
# tables stay array leaves rather than constants baked into the compiled code.
@eqx.filter_jit
def jacobians(model: ForwardModel, theta: Float[Array, "N 2"]) -> Float[Array, "N B 2"]:
    def bands(t: Float[Array, " 2"]) -> Float[Array, " B"]:  # y(θ), one pixel
        plume: AtmosphericState = atm.add_column("CH4", t[0])
        return model(plume, Lambertian(t[1]), geometry)  #                     (B,)

    # vmap over pixels outside the Jacobian, never a Jacobian of a batch;
    # jacfwd, as θ has two entries (forward mode costs one pass per input)
    return eqx.filter_vmap(jax.jacfwd(bands))(theta)


# θ_n = (ΔX_n, A): a short transect from 0 to 0.1 mol/m²                    (N, 2)
theta: Float[Array, "N 2"] = jnp.stack(
    [jnp.linspace(0.0, 0.1, 8), jnp.full(8, 0.25)], axis=1
)
K: Float[Array, "N B 2"] = jacobians(model, theta)

load_opacity reads tables from your cache. If one is missing, the error names the command that builds it (uv run python -m radtrax.spectroscopy.hitran CH4 H2O, with the [hitran] extra).

What is inside

Every layer uses only the ones beneath it, and every seam between components is a protocol (OpacityProvider, RTSolver, Surface, Instrument, NoiseModel), so you can swap one component without touching the others:

Seam	Implementations today
Atmosphere	AtmosphericState, afgl_us_standard, layers_from_levels; ViewingGeometry with plane-parallel or spherical-shell (Chapman) air mass
Opacity	TableOpacity: per-layer τ (L, G), bilinear in (ln p, T) on HITRAN tables
Solver	BeerLambert
Surface	Lambertian, PolynomialAlbedo (Legendre in ν), RossLi (RossThick + LiSparse-Reciprocal)
Instrument	BandInstrument (any SRF matrix, shift on the grid), GaussianInstrument (analytic shift in cm⁻¹ or nm)
Noise	DiagonalNoise, SNRNoise, as lineax operators with log_prob through gaussx
Retrieval	the matched filter (backgrounds, clusters, streaming, a target from autodiff), the multi-pass family (mbmp_signal, sbmp_signal, mbsp_signal), BandTable, 3D-Var, the absorption-shape test (radtrax.retrieval.absorption: pair_pixels, fit_transmittance, d_norm)
[geo]	the same retrievals as pipekit operators on georeader GeoTensors, with Sentinel-2 SRFs

Status

Stage	What lands	State
Port	plumax's band-integrated Beer–Lambert models, HITRAN tables, matched filter and 3D-Var, ported with their tests	done
One of each	one matched filter, one SRF integration, one air-mass factor; GeoTensor operators behind radtrax[geo]	done
v0	the layered forward model and matched-filter targets from autodiff	done; gate V2 against HITRAN runs in CI next (#116)
v1	spherical air mass, kernel BRDF, spectral albedo, analytic instrument, noise models, multi-pass family, band tables	on main
v1	optimal estimation on vardax; sensors and real-scene gates on EMIT, Sentinel-2 and Landsat 8/9	in progress: Wave 4

Where it sits

radtrax is a GeoModels package, and every dependency points down a layer: to Numerics (gaussx, lineax, optimistix), to GeoML (geonnax now, vardax for v1's optimal estimation) and, through [geo], to the GeoStack readers and pipelines. It never imports plumax: plume physics uses radiative transfer, not the other way round. The rules are in geostack's dependency graph, and tests/test_dependency_rules.py checks them here.

Install

The base install is JAX-only: JAX, Equinox, jaxtyping, lineax, optimistix, gaussx and geonnax, plus NumPy, xarray and scikit-learn for the matched filter's backgrounds. The extras add the rest:

Extra	Adds	For
[hitran]	HAPI	building the cross-section tables offline
[geo]	georeader, geotoolz, pipekit	operators on GeoTensors, Sentinel-2 SRFs
[bayes]	reserved: NumPyro and optax as those retrievals land	Bayesian and L-BFGS retrievals

Develop

make install     # uv sync --all-groups + pre-commit hooks
make test        # tests + doctests (float64 on)
make lint typecheck
make docs        # MyST prose + MkDocs API reference into public/

The rules for contributors and coding agents are in AGENTS.md; code review follows CODE_REVIEW.md. The figures are generated: edit docs/assets/make_assets.py and run it.

License

MIT.