ExoJAX Spectrum Code for Exoplanet Retrieval

• ExoJAX Spectrum Code is a Python-based, differentiable framework that synthesizes high-resolution spectra for exoplanetary atmospheres and brown dwarfs using GPU/TPU acceleration.
• The methodology integrates end-to-end auto-differentiability in molecular opacity, radiative transfer, and observational response to enable efficient gradient-based Bayesian inference.
• ExoJAX2 extends capabilities to emission, transmission, and reflection spectroscopy, incorporating diverse opacity sources such as HITRAN, ExoMol, and Mie scattering for robust retrievals.

ExoJAX, or the ExoJAX Spectrum Code, is an open-source, Python-based framework for high-dispersion spectral synthesis and retrieval for exoplanets and brown dwarfs. Its defining design choice is end-to-end auto-differentiability of molecular opacity, line-shape computation, radiative transfer, and observational response, implemented in JAX so that the same forward model can be accelerated on GPU/TPU hardware and embedded directly in gradient-based Bayesian inference workflows such as Hamiltonian Monte Carlo and the No-U-Turn Sampler in NumPyro. The original framework was developed for narrowband high-dispersion emission spectroscopy, while ExoJAX2 extends the approach to emission, transmission, and reflection spectroscopy at native spectral resolution (Kawahara et al., 2021, Kawahara et al., 2024).

1. Origins, remit, and conceptual identity

ExoJAX was introduced to make high-dispersion atmospheric modeling computationally compatible with full Bayesian inference. The core problem identified in the original work is that high-resolution spectra encode temperature, composition, gravity, rotation, winds, and atmospheric structure in individual molecular lines, but repeated line-by-line forward modeling inside posterior samplers is computationally expensive. ExoJAX addresses that problem by making the forward model differentiable, so that gradient-based samplers can be used directly rather than relying only on random-walk methods, grid searches, or template-based fitting (Kawahara et al., 2021).

The scope of the framework broadened substantially with ExoJAX2. ExoJAX1 focused mainly on narrowband emission and direct opacity evaluation as a function of temperature and pressure, thereby avoiding interpolation errors from pretabulated cross-section grids. ExoJAX2 retains that differentiable philosophy but generalizes it to emission, transmission, and reflection, and expands the supported opacity ecosystem to ExoMol, HITRAN, HITEMP, Kurucz, VALD3, CIA, Rayleigh scattering, H$^-$ continuum, and Mie scattering with a log-normal particle distribution (Kawahara et al., 2024).

A recurrent clarification in the literature is that ExoJAX is not merely a retrieval wrapper. ExoJAX2 explicitly describes it as a differentiable spectral model that can be wrapped by whatever inference engine the user prefers. This distinction matters because the code’s scientific role lies in the forward model itself: opacity synthesis, radiative transfer, and differentiable propagation of observational effects are the primary objects of development, while retrieval pipelines built around NumPyro, JAXopt, or related tools are downstream consumers of that model (Kawahara et al., 2024).

2. Numerical architecture and differentiable programming

The numerical core of ExoJAX is organized around differentiable opacity calculation. In the original implementation, the code provides a direct line-profile formulation for relatively small line lists and MODIT, a modified discrete integral transform, for larger lists. In the direct mode, the cross section of a line is written as

$\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$

where $S$ is the line strength, $V$ is the Voigt profile, $\beta$ is the Gaussian width from thermal broadening, and $\gamma_L$ is the Lorentzian half-width. The thermal width is given as

$$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$

To keep these operations compatible with automatic differentiation, the code implements JAX-compatible Voigt-function evaluation and custom derivatives (Kawahara et al., 2021).

ExoJAX2 introduces PreMODIT, described as a fast, memory-efficient opacity algorithm. Its central idea is to precompute a line basis density at a reference temperature and then apply fast temperature corrections rather than rebuilding the full line density at each step. The line strength is factorized as

$$S_l(T;E_l,\hat{\nu}_l)=\frac{g(\hat{\nu}_l,T)}{q_r(T)}\,f(E_l,T)\,S_{\mathrm{ref},l},$$

which allows temperature dependence to be handled without storing many temperature slices. In the GL229 B benchmark, the paper states that PreMODIT reduces device memory from tens of GB to fractions of a GB, making broadband retrievals feasible on current GPUs (Kawahara et al., 2024).

Differentiability is used both for optimization and for posterior exploration. ExoJAX-based studies employ HMC/NUTS in NumPyro, and in some cases JAXopt/Adam is used for initialization or direct $\chi^2$ minimization. Published retrievals also show that stable sampling in high-dimensional settings may require transformed coordinates: the Luhman 16AB analysis rescales parameters affinely to $\mathcal{O}(1)$, and for the Gaussian-process temperature-profile case it uses a whitened latent representation to improve convergence. That study also adopts a dense mass matrix in NUTS, illustrating how ExoJAX’s differentiability is coupled to careful probabilistic numerics rather than treated as a purely software-level convenience (Yama et al., 28 Nov 2025).

3. Opacity sources, atmospheric structure, and radiative transfer

ExoJAX’s physical model begins with line-by-line opacity from molecular databases. The original framework supports HITRAN/HITEMP and ExoMol line lists, along with collision-induced absorption from the HITRAN CIA database. ExoJAX2 enlarges that set to include atomic and metal databases such as Kurucz and VALD3, and extends the continuum treatment to Rayleigh scattering and H$\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$0 continuum (Kawahara et al., 2021, Kawahara et al., 2024).

In application papers, the layer optical depth is commonly expressed as a sum of molecular absorption, CIA, and cloud opacity: $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$1 For molecular opacity, the Luhman 16AB retrieval writes

$\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$2

with line strengths and pressure broadening taken from HITRAN/HITEMP or ExoMol under the assumed $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$3 structure. In that study, $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$4, $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$5, $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$6, and $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$7 are retrieved as vertically uniform VMRs, while $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$8 and $\sigma(\nu)=S\,V(\nu-\nu_0,\beta,\gamma_L),$9 make up the remainder in solar ratio by mass, converted to a volume ratio (Yama et al., 28 Nov 2025).

Radiative transfer in ExoJAX1 is formulated for pure-absorption emission spectra. The formal equation is written as

$S$0

and the emission solver is implemented in a JAX-compatible form. ExoJAX-based retrievals have used both two-stream and eight-stream approximations. For example, the Gl 229 B retrieval uses a pure-absorption, intensity-based scheme with 8 streams, and the Luhman 16AB retrieval uses an eight-stream radiative transfer approximation followed by rotational broadening, radial-velocity shifting, per-chip normalization, and a Gaussian-process noise model (Kawahara et al., 2021, Kawashima et al., 2024, Yama et al., 28 Nov 2025).

ExoJAX2 generalizes this physics beyond thermal emission. For transmission spectroscopy it computes chord optical depths through a geometric matrix,

$S$1

and evaluates the atmospheric annulus through $S$2. For reflection spectroscopy it implements differentiable two-stream radiative transfer with scattering. The paper also introduces the linear source approximation for non-isothermal layers, Opart for layer-by-layer opacity plus radiative transfer, $S$3-stitching for broadband calculations, and FFT-based rotational or instrumental broadening with overlap-and-add for segmented convolution (Kawahara et al., 2024).

4. Retrieval workflows and model composition

Published ExoJAX studies use the code as the forward-modeling core inside larger Bayesian retrieval pipelines. A representative workflow is given explicitly for Luhman 16AB: build a line-by-line atmosphere model in ExoJAX, compute emergent spectra for a 1D plane-parallel atmosphere, apply rotation, Doppler shift, per-chip normalization, and a Gaussian-process noise model, and then sample the posterior with HMC/NUTS in NumPyro. That analysis discretizes the atmosphere on 101 pressure layers from $S$4 to $S$5 bar and compares two thermal parameterizations: a power-law profile,

$S$6

and a Gaussian-process profile defined on 21 temperature nodes and interpolated to the 101-layer radiative-transfer grid (Yama et al., 28 Nov 2025).

ExoJAX can also be embedded in more elaborate composite forward models. The HR 7672 B retrieval is explicitly based on ExoJAX v1.5.1 and simultaneously models the companion atmosphere, telluric absorption in Earth’s atmosphere, leakage of host-star light into the companion fiber, cloud opacity, instrumental broadening, and rotational broadening. In that study, the brown-dwarf atmosphere is modeled with 300 layers; J-band opacities include $S$7 and $S$8, H-band opacities include $S$9 only, telluric transmission is fitted rather than removed empirically, and the likelihood jointly fits the normalized high-resolution spectrum and the J-band photometric magnitude (Kasagi et al., 2 Aug 2025).

A third workflow, used for Gl 229 B, shows how ExoJAX handles normalized but not absolutely flux-calibrated data. That retrieval uses ExoJAX v1.5 on a 200-layer atmosphere from $V$0 to $V$1 bar, computes the internal spectrum at $V$2, solves for a multiplicative scale factor at each HMC step, and combines the normalized H-band spectrum with the observed H-band magnitude in the likelihood. The study also reports successful HMC convergence with 500 warmup steps and 1000 samples and states that the computations were run in 32-bit JAX mode on an NVIDIA A100 80 GB PCIe GPU (Kawashima et al., 2024).

Across these retrievals, the parameterizations remain deliberately broad but structured: temperature-profile parameters, $V$3, RV, $V$4, molecular VMRs, cloud-top pressure or cloud-deck parameters, normalization or leakage terms, limb-darkening coefficients, and GP hyperparameters are all represented as differentiable parameters inside a single forward model. This suggests that ExoJAX’s main methodological contribution lies in unifying atmospheric physics and nuisance modeling within the same autodiff-compatible computation graph (Yama et al., 28 Nov 2025, Kasagi et al., 2 Aug 2025, Kawashima et al., 2024).

5. Benchmark applications and scientific outputs

The published literature uses ExoJAX primarily on benchmark brown dwarfs and directly imaged substellar companions, where high-resolution spectra make line-list quality, radiative-transfer approximations, and nuisance-model specification empirically testable.

Target and study	Observational setting	Principal ExoJAX-based result
Luhman 16 A (Kawahara et al., 2021)	CRIRES/VLT, $V$5–$V$6, $V$7	Retrieved $V$8 at 1 bar and $V$9
Gl 229 B (Kawashima et al., 2024)	Subaru/IRD H band, $\beta$0	Retrieved $\beta$1; HITEMP and ExoMol MM methane fit better than YT34to10; results favored an unresolved equal-mass binary interpretation
HR 7672 B (Kasagi et al., 2 Aug 2025)	REACH/Subaru Y, J, H bands, $\beta$2	Identified $\beta$3 and $\beta$4 as the primary absorbers; cloudy model favored with $\beta$5; $\beta$6 km s$\beta$7
Luhman 16AB (Yama et al., 28 Nov 2025)	VLT/CRIRES K band retrieval of both components	Inferred $\beta$8 ratios of about 0.67 for both components; CO line-list systematics shifted C/O by about 7%

These applications also serve as code validations. The first ExoJAX paper demonstrated that a model including water, carbon monoxide, and $\beta$9 CIA fit the Luhman 16 A high-dispersion spectrum well and showed that line ratios can act as temperature diagnostics (Kawahara et al., 2021). The Gl 229 B study used the retrieved high-S/N H-band spectrum to test methane line lists, found NH$\gamma_L$0 above thermochemical equilibrium expectations, and argued that CIA can help break abundance–gravity degeneracy sufficiently to yield a spectroscopic mass estimate, albeit one still differing from the astrometric mass by about $\gamma_L$1 in the mass-free run (Kawashima et al., 2024). The HR 7672 B analysis established ExoJAX as the backbone of a realistic end-to-end retrieval under cloud opacity, host-star leakage, telluric absorption, blaze or continuum residuals, and limited wavelength coverage (Kasagi et al., 2 Aug 2025). The Luhman 16AB work, in turn, emphasized systematic testing, especially the effect of CO broadening prescriptions on inferred C/O (Yama et al., 28 Nov 2025).

ExoJAX2 adds three broader demonstrations outside the original narrowband regime. It reports a GL229 B high-dispersion emission retrieval with $\gamma_L$2, a WASP-39 b JWST/NIRSpec G395H transmission retrieval at original resolution $\gamma_L$3 with a radial velocity of $\gamma_L$4 km/s, and a Jupiter reflection retrieval yielding a methane enrichment of roughly 2.7–4.1 times solar (Kawahara et al., 2024).

6. Uncertainties, limitations, and relation to adjacent software

A recurrent finding in ExoJAX-based studies is that retrieval sophistication does not remove spectroscopic systematics. The clearest example is the Luhman 16AB analysis, which compared ExoMol CO, HITEMP CO with air broadening, and HITEMP CO with H$\gamma_L$5 broadening inside the same ExoJAX framework and found that the CO line list is the dominant systematic, shifting inferred C/O by about 7%. The same study found that changing the $\gamma_L$6–$\gamma_L$7 parameterization from a power law to a Gaussian process had only minor impact on C/O, while photometric variability was subdominant, with epoch-wise results agreeing with the stacked-spectrum retrieval within at most $\gamma_L$8 in C/O and $\gamma_L$9 in cloud-deck pressure (Yama et al., 28 Nov 2025).

Other published limitations are model-specific but methodologically instructive. The Gl 229 B retrieval used a simple power-law $$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$0–$$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$1 profile, ignored clouds, and explicitly described the mass retrieval as model dependent and sensitive to deep thermal structure and reduction systematics. The same paper derived the line and CIA optical-depth scalings to explain the abundance–gravity degeneracy and used Appendix tests to show how CIA helps break it (Kawashima et al., 2024). The HR 7672 B study identified cloud opacity as useful and physically interpretable, but also noted the need for broader wavelength coverage and improved treatment of periodic instrumental noise and cloud or $$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$2–$$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$3 flexibility to exploit future REACH data fully (Kasagi et al., 2 Aug 2025). ExoJAX2 can be read as a response to earlier broadband and memory limitations, since it explicitly targets native-resolution analysis without aggressive binning and introduces PreMODIT, Opart, $$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$4-stitching, and FFT-based broadening to reduce the computational bottleneck (Kawahara et al., 2024).

The position of ExoJAX within atmospheric-spectrum software is most clearly seen by comparison with adjacent tools. Exo-Transmit is an open-source C code for exoplanet transmission spectra that performs a pure-absorption transmission calculation with Rayleigh opacity and simple cloud or haze parameterizations; it is a forward model for transmission spectroscopy rather than an autodifferentiable high-resolution retrieval engine (Kempton et al., 2016). PyExoCross is a Python adaptation of ExoCross for post-processing very large molecular line lists, generating cross-sections, stick spectra, partition functions, and format conversions; it is described as a line-list postprocessor with CPU parallelization and no GPU support, conceptually complementary to ExoJAX rather than a replacement for its differentiable forward modeling (Zhang et al., 2024). Exo Skryer is a separate JAX-accelerated sub-stellar retrieval framework with YAML-based model assembly and a direct retrieval method for aerosol optical constants $$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$5 and $$\beta_T=\sqrt{\frac{k_B T}{M m_u}}\,\frac{\nu_0}{c}.$$6; it occupies a related JAX-native ecosystem but emphasizes modular retrieval construction and cloud or aerosol inference rather than ExoJAX’s line-by-line differentiable spectrum engine (Lee, 23 Feb 2026).

Taken together, these studies define the ExoJAX Spectrum Code as a differentiable atmospheric spectrum engine whose technical identity is anchored in JAX-based opacity calculation, radiative transfer, and gradient-compatible inference. Its scientific record to date is strongest in high-resolution brown-dwarf and directly imaged companion spectroscopy, where it has been used both to retrieve atmospheric properties and to expose the limits imposed by line lists, nuisance models, and thermal-structure parameterization (Kawahara et al., 2021, Kawahara et al., 2024).