Forward model

v0's layered forward model, ForwardModel: an AtmosphericState through an OpacityProvider, an RTSolver (with a Surface and a ViewingGeometry) and an Instrument, to one value per band. Built on it, unit_absorption_target (in radtrax.retrieval.matched_filter) is the matched-filter target t = ∂ ln y / ∂ΔX for a scene, by one forward-mode derivative. Both are importable from the top level.

A swath's Jacobians are eqx.filter_vmap of a per-pixel jax.jacrev:

def y_of(theta):  # (ΔX [mol/m²], albedo, shift [cm⁻¹]) → y (B,)
    fm = eqx.tree_at(lambda m: m.instrument.shift_cm, model, theta[2])
    return fm(atm.add_column("CH4", theta[0]), Lambertian(theta[1]), geometry)

K = eqx.filter_jit(eqx.filter_vmap(jax.jacrev(y_of)))(thetas)  # (N, B, 3)

forward_model

The layered forward model.

ForwardModel composes the stages of radtrax · 3: an atmospheric state through an OpacityProvider into per-layer optical depth, through an RTSolver (with a Surface and a ViewingGeometry) into radiance at the top of the atmosphere, through an Instrument into one value per band:

.. math:: y = I\bigl(S(O(x), R, g)\bigr) \in \mathbb{R}^B .

Every leaf is differentiable, so a retrieval's Jacobian is jax.jacrev of the model, and a swath is eqx.filter_vmap of that per-pixel Jacobian (never a Jacobian of a batch).

:func:radtrax.unit_absorption_target (in :mod:radtrax.retrieval.matched_filter) is v0's deliverable built on it: the matched-filter target t = ∂ ln y / ∂ΔX for a scene.

Shapes: L layers, G wavenumbers, B bands.

ForwardModel

Bases: Module

State → opacity → solver → instrument: band values for one pixel.

The surface and the geometry are arguments, not fields, because they change from pixel to pixel; the opacity, solver and instrument are shared by a scene. To differentiate a module leaf (the spectral shift, say) and not the tables, eqx.tree_at it in from a parameter, or split with eqx.partition. The opacity's and the instrument's grids are checked equal when they are concrete (not when the model is built inside jit). It is an observation operator H(x), not a time-stepping pipekit_cycle.ForwardModel (no step).

Attributes:

Name	Type	Description
opacity	OpacityProvider	Per-layer optical depth from the state.
solver	RTSolver	Radiance from optical depth, surface and geometry.
instrument	Instrument	Band values from radiance; its wavenumber grid is the opacity's.

Examples:

>>> import jax.numpy as jnp
>>> from radtrax.atmosphere import ViewingGeometry, layers_from_levels
>>> from radtrax.instrument import BandInstrument
>>> from radtrax.solvers import BeerLambert
>>> from radtrax.spectroscopy import CrossSectionTable, TableOpacity
>>> from radtrax.surface import Lambertian
>>> nu = jnp.array([4300.0, 4300.1, 4300.2])
>>> table = CrossSectionTable(
...     jnp.log(jnp.array([1e4, 1e5])),
...     jnp.array([200.0, 300.0]),
...     jnp.full((2, 2, 3), 1e-21),
... )
>>> model = ForwardModel(
...     TableOpacity(nu, {"CH4": table}),
...     BeerLambert(),
...     BandInstrument(nu, jnp.full((1, 3), 1 / 3)),
... )
>>> atm = layers_from_levels(
...     jnp.array([1e5, 5e4]),
...     jnp.array([288.0, 255.0]),
...     {"CH4": jnp.full(2, 1.9e-6)},
... )
>>> model(atm, Lambertian(0.3), ViewingGeometry(1.0, 1.0)).shape
(1,)

optics

optics(state: AtmosphericState) -> OpticalProperties

Per-layer optical depth on the model's grid, with the layers' altitudes.

Source code in src/radtrax/forward_model.py

def optics(self, state: AtmosphericState) -> OpticalProperties:
    """Per-layer optical depth on the model's grid, with the layers' altitudes."""
    return OpticalProperties(
        self.opacity.nu_cm, self.opacity(state), z_m=state.z_m, dz_m=state.dz_m
    )

radiance

radiance(state: AtmosphericState, surface: Surface, geometry: ViewingGeometry) -> Float[Array, ' G']

Radiance at the top of the atmosphere, before the instrument.

Source code in src/radtrax/forward_model.py

def radiance(
    self, state: AtmosphericState, surface: Surface, geometry: ViewingGeometry
) -> Float[Array, " G"]:
    """Radiance at the top of the atmosphere, before the instrument."""
    return self.solver(self.optics(state), surface, geometry)