# How to Perform Reflection Spectrum Calculations Using ExoJAX

> Learn to perform reflection spectrum calculations with ExoJAX. This guide explores its JAX pipeline and radiative transfer solvers for efficient, differentiable atmospheric retrieval.

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

---

**ExoJAX computes reflected-light spectra through a modular JAX pipeline that combines user-defined opacity layers with two-stream radiative transfer solvers like `OpartReflectPure` and `OpartReflectEmis`, enabling fully differentiable reflection spectrum calculations for atmospheric retrieval.**

ExoJAX is a GPU-accelerated radiative transfer code built on JAX for exoplanet atmospheric analysis. When performing reflection spectrum calculations using ExoJAX, you work with a layered architecture that separates opacity generation from flux computation, ensuring the entire pipeline remains JIT-compilable and auto-differentiable for gradient-based retrievals.

## Architecture of ExoJAX Reflection Calculations

ExoJAX models reflected-light spectra with a modular design that separates **opacity generation**, **radiative-transfer solving**, and **layer-by-layer flux accumulation**. This architecture is implemented across several core modules in the `src/exojax/rt/` directory.

### Opacity Layer Generation

The foundation of any reflection calculation is the opacity layer class. In [`src/exojax/rt/layeropacity.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/layeropacity.py), the function `single_layer_optical_depth` transforms molecular cross-sections into per-layer optical depth (`dtau`). Your custom opacity layer must return three JAX arrays: `dtau` (optical depth), `single_scattering_albedo`, and `asymmetric_parameter` (asymmetry factor g).

### Radiative Transfer Solvers

The two-stream radiative transfer implementation resides in [`src/exojax/rt/reflect.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/reflect.py). For pure reflection scenarios (no thermal emission), use `ArtReflectPure` or its optimized wrapper `OpartReflectPure`. When modeling both reflected stellar light and planetary thermal emission, use `ArtReflectEmis` or `OpartReflectEmis`. These solvers employ the Toon hemispheric mean coefficients via `setrt_toonhm` for accurate flux adding.

### Flux-Adding Algorithm

The core reflection calculation uses the flux-adding method implemented in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) within `rtrun_reflect_fluxadding_toonhm`. This algorithm propagates effective reflectivity (`R̂`) and source terms (`Ŝ`) from the surface to the top of the atmosphere using the recurrence relations:

```

R̂_i = σ_i + τ_i^2 * R̂_{i-1} / (1 - σ_i * R̂_{i-1})
Ŝ_i = π_i + τ_i * (Ŝ_{i-1} + π_i * R̂_{i-1}) / (1 - σ_i * R̂_{i-1})

```

Where `σ_i` is the scattering coefficient, `τ_i` is the transmission coefficient, and `π_i` is the source term. The final reflected spectrum is computed as `F_out = R̂_top * F_inc + Ŝ_top`.

## Step-by-Step Implementation

To perform reflection spectrum calculations using ExoJAX, you must define an opacity layer class, instantiate a reflector object, and execute the flux calculation with a layer update function.

### Pure Reflection Calculation

The following complete example demonstrates how to compute a pure reflection spectrum using `OpartReflectPure` with a mock CO opacity layer:

```python
from exojax.rt import OpartReflectPure
from exojax.opacity import OpaPremodit
from exojax.rt.layeropacity import single_layer_optical_depth
from exojax.test.emulate_mdb import mock_wavenumber_grid, mock_mdbExomol

import jax.numpy as jnp

# Define the opacity layer class

class OpaLayer:
    def __init__(self):
        self.nu_grid, self.wav, self.resolution = mock_wavenumber_grid()
        self.gravity = 2478.57  # cm s^-2

        self.mdb_co = mock_mdbExomol()
        self.opa_co = OpaPremodit(
            self.mdb_co, 
            self.nu_grid,
            auto_trange=[400.0, 1500.0]
        )

    def __call__(self, params):
        temperature, pressure, dP, mixing_ratio = params
        # Compute cross-section and convert to optical depth

        xsv = self.opa_co.xsvector(temperature, pressure)
        dtau = single_layer_optical_depth(
            dP, xsv, mixing_ratio,
            self.mdb_co.molmass, self.gravity
        )
        # Scattering parameters

        ssa = jnp.ones_like(dtau) * 0.3
        g = jnp.ones_like(dtau) * 0.01
        return dtau, ssa, g

# Initialize the reflection calculator

opart = OpartReflectPure(
    OpaLayer(),
    pressure_top=1e-6,  # bar

    pressure_btm=1.0,   # bar

    nlayer=200
)

# Define the layer update function for the JAX scan

def layer_update_function(carry, params):
    carry = opart.update_layer(carry, params)
    return carry, None

# Build atmospheric profiles

temperature = opart.powerlaw_temperature(T0=1300.0, alpha=0.1)
mixing_ratio = opart.constant_profile(3e-4)
layer_params = [temperature, opart.pressure, opart.dParr, mixing_ratio]

# Define incident flux and surface albedo

incoming_flux = jnp.ones_like(opart.opalayer.nu_grid)
surface_albedo = 1.0
reflectivity_surface = surface_albedo * jnp.ones_like(incoming_flux)

# Compute reflected spectrum

reflected_flux = opart(
    layer_params,
    layer_update_function,
    reflectivity_surface,
    incoming_flux
)

```

### Adding Thermal Emission

For atmospheres where thermal emission contributes significantly, use `OpartReflectEmis` instead. This extends the pure reflection calculation to include layer source terms:

```python
from exojax.rt import OpartReflectEmis

opart_emis = OpartReflectEmis(
    OpaLayer(),
    pressure_top=1e-6,
    pressure_btm=1.0,
    nlayer=200
)

# The layer_update_function remains identical

reflected_emission_flux = opart_emis(
    layer_params,
    layer_update_function,
    source_bottom=jnp.zeros_like(opart_emis.nu_grid),
    reflectivity_bottom=reflectivity_surface,
    incoming_flux=incoming_flux
)

```

This returns a spectrum containing both **stellar reflection** and **thermal emission** components, essential for modeling hot Jupiters or warm Neptunes where both processes contribute.

## Key Source Files in ExoJAX

Understanding the codebase structure helps when customizing reflection calculations:

| File | Purpose |
|------|---------|
| [`src/exojax/rt/reflect.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/reflect.py) | Core reflection classes (`OpartReflectPure`, `OpartReflectEmis`, `ArtReflectPure`, `ArtReflectEmis`) |
| [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) | Low-level flux-adding solver `rtrun_reflect_fluxadding_toonhm` |
| [`src/exojax/rt/opart.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/opart.py) | Generic `Opart` framework for JAX `scan` operations over atmospheric layers |
| [`src/exojax/rt/common.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/common.py) | Base class `ArtCommon` providing atmospheric profile utilities (`powerlaw_temperature`, `constant_profile`) |
| [`src/exojax/rt/layeropacity.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/layeropacity.py) | `single_layer_optical_depth` helper for converting cross-sections to optical depth |
| [`src/exojax/opacity/opapremodit.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/opacity/opapremodit.py) | Fast opacity calculator `OpaPremodit` for molecular cross-sections |
| [`tests/unittests/multi/opart/opart_reflection_test.py`](https://github.com/hajimekawahara/exojax/blob/main/tests/unittests/multi/opart/opart_reflection_test.py) | Unit tests demonstrating valid reflection calculations |

## Summary

- ExoJAX implements **reflection spectrum calculations** through a modular JAX architecture separating opacity generation from radiative transfer solving.
- The `OpartReflectPure` class in [`src/exojax/rt/reflect.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/reflect.py) handles pure stellar reflection, while `OpartReflectEmis` adds thermal emission capabilities.
- Users must define an **opacity layer class** that returns `dtau`, `single_scattering_albedo`, and `asymmetric_parameter` for each atmospheric layer.
- The **flux-adding algorithm** with Toon hemispheric mean coefficients propagates effective reflectivity from the surface to the top of the atmosphere.
- The entire pipeline is JAX-native, supporting **JIT compilation** and **automatic differentiation** for gradient-based atmospheric retrievals.

## Frequently Asked Questions

### What is the difference between OpartReflectPure and OpartReflectEmis?

`OpartReflectPure` computes only the reflected stellar component, assuming the atmosphere does not emit thermal radiation. It is ideal for modeling cold planets or high-contrast direct imaging scenarios. `OpartReflectEmis` extends this functionality to include thermal emission from atmospheric layers, making it suitable for warm exoplanets where both reflected starlight and planetary thermal radiation contribute to the observed spectrum. Both classes are implemented in [`src/exojax/rt/reflect.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/reflect.py) and share identical APIs except for the emission-specific parameters in `OpartReflectEmis`.

### How do I define a custom opacity layer for reflection calculations?

You must create a callable Python class (typically named `OpaLayer`) that accepts atmospheric parameters and returns three JAX arrays: optical depth (`dtau`), single-scattering albedo (`ssa`), and asymmetry parameter (`g`). Inside the `__call__` method, compute molecular cross-sections using opacity calculators like `OpaPremodit` from [`src/exojax/opacity/opapremodit.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/opacity/opapremodit.py), then convert these to optical depth using `single_layer_optical_depth` from [`src/exojax/rt/layeropacity.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/layeropacity.py). This modular design allows you to swap opacity models without modifying the radiative transfer code in [`src/exojax/rt/reflect.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/reflect.py).

### Can I compute gradients of reflected spectra for atmospheric retrieval?

Yes. Because ExoJAX is built entirely on JAX, all reflection spectrum calculations are fully differentiable. The `OpartReflectPure` and `OpartReflectEmis` classes return JAX arrays, allowing you to apply `jax.grad` or `jax.jit` to the entire pipeline. This enables gradient-based optimization and Hamiltonian Monte Carlo methods for atmospheric retrieval, where you can directly optimize temperature profiles, mixing ratios, or cloud properties by differentiating through the flux-adding algorithm implemented in [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py).

### What radiative transfer method does ExoJAX use for reflection?

ExoJAX implements the **flux-adding method** with Toon hemispheric mean coefficients for two-stream radiative transfer. Specifically, [`src/exojax/rt/rtransfer.py`](https://github.com/hajimekawahara/exojax/blob/main/src/exojax/rt/rtransfer.py) contains `rtrun_reflect_fluxadding_toonhm`, which propagates effective reflectivity (`R̂`) and source terms (`Ŝ`) from the surface to the top of the atmosphere using recurrence relations. This method accurately handles multiple scattering and arbitrary surface albedos while maintaining computational efficiency through JAX's `scan` operations over atmospheric layers.