# How ExoJAX Implements Radiative Transfer Calculations for Emission and Transmission Spectra

> ExoJAX uses a JAX-first engine for fast radiative transfer calculations. Explore its JIT-compiled solvers for emission and transmission spectra on GPUs and TPUs.

- Repository: [Hajime Kawahara/exojax](https://github.com/hajimekawahara/exojax)
- Tags: how-to-guide
- Published: 2026-03-03

---

**ExoJAX implements radiative transfer calculations for emission and transmission spectra through a modular, JAX-first engine that JIT-compiles flux-based two-stream, intensity-based n-stream, and scattering solvers for GPU/TPU acceleration.**

The ExoJAX library provides a fully differentiable framework for exoplanetary atmospheric forward modeling. Its radiative transfer (RT) module structures computations into three distinct layers—user-facing API objects, common atmospheric utilities, and core JIT-compiled solvers—enabling seamless switching between emission and transmission calculations while maintaining automatic gradient support through JAX.

## Architecture of the Radiative Transfer Engine

ExoJAX organizes its radiative transfer calculations for emission and transmission spectra into a three-layer abstraction stack. This design separates atmospheric geometry handling from the mathematical solvers, allowing users to swap opacity sources or RT methods without restructuring their code.

| Layer | Purpose | Main Classes / Functions | Source Location |
|-------|---------|--------------------------|-----------------|
| **User-facing RT objects** | Provide simple API (`run`) for opacity fields and atmospheric structures | `ArtEmisPure`, `ArtEmisScat`, `OpartEmisPure`, `OpartEmisScat`, `ArtTransPure` | [`src/exojax/rt/emis.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/emis.py), [`src/exojax/rt/trans.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/trans.py) |
| **Common RT utilities** | Handle pressure grids, geometry, opacity conversion, Gaussian quadrature | `ArtCommon` and methods (`init_pressure_profile`, `atmosphere_height`, `opacity_profile_xs`, `initialize_gaussian_quadrature`) | [`src/exojax/rt/common.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/common.py) |
| **Core solvers** | JIT-compiled implementations of RT equations | `rtrun_emis_*`, `rtrun_trans_*`, `setrt_toonhm`, `solve_lart_twostream`, `solve_fluxadding_twostream` | [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) |

## Atmospheric Setup and Geometry

All radiative transfer calculations for emission and transmission spectra in ExoJAX begin with the `ArtCommon` base class defined in [`src/exojax/rt/common.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/common.py). This class establishes the computational grid and geometric framework required by both emission and transmission solvers.

During initialization, `ArtCommon` constructs a log-spaced pressure grid spanning `pressure_top` to `pressure_btm` across `nlayer` atmospheric layers. It computes layer midpoint pressures, pressure thicknesses (`dParr`), and the pressure decrease rate necessary for hydrostatic calculations.

```python
class ArtCommon:
    def __init__(self, pressure_top, pressure_btm, nlayer, nu_grid=None):
        # builds pressure, dParr, etc.

        self.init_pressure_profile()          

```

The `atmosphere_height` method calculates the geometric height and radius of each layer given temperature profiles, mean molecular weight, and surface gravity. These values become essential for transmission geometry where light paths traverse chords through the atmospheric annuli.

## Opacity to Optical Depth Conversion

Before executing radiative transfer calculations for emission and transmission spectra, ExoJAX converts molecular cross-sections into layer optical depths. The `ArtCommon.opacity_profile_xs` method in [`src/exojax/rt/common.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/common.py) orchestrates this conversion by calling routines from `exojax.rt.layeropacity`.

Given cross-sections `xs` (with shape `(Nlayer, Nnu)`), mixing ratios, molecular mass, and gravity, the function returns the differential optical depth matrix `dtau`:

```python
def opacity_profile_xs(self, xs, mixing_ratio, molmass, gravity):
    return layer_optical_depth(self.dParr, jnp.abs(xs), mixing_ratio,
                              molmass, gravity)

```

This `dtau` matrix serves as the primary input for both emission and transmission solvers, representing the absorption probability per layer across the spectral grid.

## Emission Spectra: Pure Absorption

For emission radiative transfer calculations without scattering, ExoJAX provides multiple solver strategies through `ArtEmisPure` in [`src/exojax/rt/emis.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/emis.py). The class selects solvers via a dispatch dictionary (`rtsolver_dict`) that maps user-selected methods to JIT-compiled functions in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py).

Available pure-absorption solvers include:

- **`rtrun_emis_pureabs_fbased2st`**: Flux-based two-stream approximation optimized for speed, similar to HELIOS-R1/R2 implementations.
- **`rtrun_emis_pureabs_ibased`**: Intensity-based n-stream using Gaussian quadrature, compatible with NEMESIS/pRT-style calculations.
- **`rtrun_emis_pureabs_linsap`**: Linear source approximation variant for improved accuracy in temperature gradients.

The source function utilizes Planck radiance (`piBarr`) evaluated at the layer temperatures:

```python
sourcef = piBarr(temperature, self.nu_grid)
rtfunc = self.rtsolver_dict[self.rtsolver]
spectrum = rtfunc(dtau, sourcef, self.mus, self.weights)

```

## Emission Spectra: Scattering

When scattering dominates, ExoJAX implements the Toon hemispheric-mean two-stream framework through `ArtEmisScat` in [`src/exojax/rt/emis.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/emis.py). This approach handles single scattering albedo and asymmetric parameters using either the Layer-Adding Radiative Transfer (LART) or flux-adding method.

The scattering workflow in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) proceeds through three stages:

1. **Coefficient generation** (`setrt_toonhm`): Computes transmission coefficients, scattering coefficients, reduced Planck sources, and auxiliary matrices (`zeta_plus`, `zeta_minus`, `lambdan`).

2. **Tridiagonal system formation** (`settridiag_toohm`): Assembles the matrix structure for the two-stream solver.

3. **Solution** (`solve_lart_twostream` or `solve_fluxadding_twostream`): Solves the boundary value problem for the emergent flux.

Users interact with high-level wrappers that accept `dtau`, single scattering albedo (`ssa`), asymmetric parameter (`g`), and Planck source:

```python
spectrum, *_ = rtrun_emis_scat_lart_toonhm(dtau, ssa, g, piBarr(T, nu_grid))

```

## Transmission Spectra

For transmission radiative transfer calculations, ExoJAX computes the effective planetary radius via chord integration through the atmospheric annuli. The `ArtTransPure` class in [`src/exojax/rt/trans.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/trans.py) manages the geometry and numerical integration schemes.

The transmission calculation follows this sequence:

1. **Geometry initialization**: `atmosphere_height` computes normalized heights (`h_norm`) and lower-boundary radii (`r_lower`) for each layer.

2. **Chord optical depth**: `chord_optical_depth` in [`src/exojax/rt/chord.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/chord.py) projects the vertical optical depth matrix `dtau` onto slant paths using the geometric chord matrix.

3. **Integration**: The slant optical depths are integrated along the chord using either trapezoidal or Simpson quadrature to obtain the squared ratio of transit radius to bottom radius `(R/R_b)^2`.

Available integration schemes include:

- **`rtrun_trans_pureabs_trapezoid`**: Uses lower-boundary chord optical depths only, suitable for quick calculations.
- **`rtrun_trans_pureabs_simpson`**: Incorporates both lower-boundary and midpoint chord optical depths for higher accuracy.

```python
cgm = chord_geometric_matrix_lower(h_norm, r_lower)
dtau_chord = chord_optical_depth(cgm, dtau)
transit_sq = transmitter.integration_dict[transmitter.integration](
                dtau_chord, r_lower, r_top)
transit_radius = jnp.sqrt(transit_sq) * radius_btm

```

## JAX-First Performance Optimization

All radiative transfer calculations for emission and transmission spectra in ExoJAX are built on a **JAX-first architecture** that enables hardware acceleration and automatic differentiation. Every core solver function in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) is decorated with `@jit`, ensuring that the Python overhead is eliminated after compilation.

The intensity-based solvers utilize `jax.lax.scan` for angular quadrature loops instead of native Python iteration, allowing the compiler to unroll and optimize the Gaussian quadrature efficiently. Quadrature weights are generated using SciPy's `roots_legendre` wrapped in JAX-compatible operations, ensuring numerical stability for n-stream calculations.

This design makes ExoJAX fully differentiable—users can compute gradients of emission or transmission spectra with respect to temperature profiles, mixing ratios, or molecular abundances directly through `jax.grad`, enabling atmospheric retrieval and optimization workflows.

## Summary

- **ExoJAX** implements radiative transfer calculations for emission and transmission spectra through a three-layer architecture separating user APIs, atmospheric utilities, and JIT-compiled solvers.
- The **`ArtCommon`** base class in [`src/exojax/rt/common.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/common.py) handles pressure grids, geometry, and conversion of cross-sections to optical depths (`dtau`).
- **Emission** calculations support pure-absorption solvers (flux-based two-stream, intensity-based n-stream) in `ArtEmisPure` and scattering-aware Toon two-stream methods (LART, flux-adding) in `ArtEmisScat`.
- **Transmission** spectra use chord integration of slant optical depths via `ArtTransPure`, supporting trapezoidal and Simpson quadrature schemes for computing effective planetary radii.
- All solvers leverage **JAX** for GPU/TPU acceleration and automatic differentiation, enabling gradient-based atmospheric retrievals.

## Frequently Asked Questions

### How does ExoJAX handle scattering in emission spectra?

ExoJAX handles scattering through the Toon hemispheric-mean two-stream approximation implemented in `ArtEmisScat` ([`src/exojax/rt/emis.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/emis.py)). The framework computes transmission and scattering coefficients using `setrt_toonhm`, then solves the radiative transfer equation via either the Layer-Adding Radiative Transfer (LART) method (`solve_lart_twostream`) or flux-adding (`solve_fluxadding_twostream`). Users provide single scattering albedo and asymmetric parameters alongside the optical depth matrix.

### What is the difference between flux-based and intensity-based emission solvers?

Flux-based solvers (`rtrun_emis_pureabs_fbased2st`) compute the emergent flux directly using a two-stream approximation, optimized for speed and similar to HELIOS-R1/R2 implementations. Intensity-based solvers (`rtrun_emis_pureabs_ibased`) integrate the radiative transfer equation over discrete angles using Gaussian quadrature (n-stream), computing specific intensity before converting to flux. The intensity-based approach offers higher accuracy for limb-dependent applications but requires more computational resources.

### How does ExoJAX compute transmission spectra for exoplanet atmospheres?

ExoJAX computes transmission spectra through chord integration implemented in `ArtTransPure` ([`src/exojax/rt/trans.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/trans.py)). The method calculates geometric heights and radii for each atmospheric layer, then projects the vertical optical depth (`dtau`) onto slant light paths using `chord_optical_depth`. The slant optical depths are integrated along the chord using either trapezoidal or Simpson quadrature to compute the squared ratio of transit radius to bottom radius, which converts directly to the observed transit depth.

### Can ExoJAX radiative transfer calculations be differentiated for atmospheric retrieval?

Yes, all ExoJAX radiative transfer calculations for emission and transmission spectra are fully differentiable through JAX. The core solvers in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) use `@jit` decoration and `jax.lax.scan` for loops, enabling automatic differentiation via `jax.grad` or `jax.jacfwd`. This allows users to compute gradients of synthetic spectra with respect to temperature profiles, mixing ratios, or molecular abundances, enabling gradient-based optimization and Bayesian inference for atmospheric retrieval.