How to Define Custom Molecular Opacities or Add New Species to ExoJAX

You can define custom molecular opacities in ExoJAX by wrapping line data in an MDBSnapshot and injecting custom physics providers via the pf_provider and broadening_strategy arguments in OpaPremodit.from_snapshot().

ExoJAX is an open-source radiative transfer library for exoplanet atmospheric retrievals. To define custom molecular opacities or add new species to ExoJAX, you leverage a snapshot-first architecture that separates raw molecular line data from temperature-dependent physics calculations.

Understanding the Snapshot and Provider Architecture

ExoJAX decouples molecular databases into two independent components. The snapshot (MDBSnapshot) stores raw line data as serializable NumPy arrays, while physics providers implement partition functions and pressure broadening calculations.

This design, implemented in src/exojax/database/contracts.py and src/exojax/opacity/providers.py, allows you to inject arbitrary line lists without modifying core opacity code. The OpaPremodit class in src/exojax/opacity/premodit/api.py accepts these snapshots and providers through factory methods from_snapshot() and from_mdb().

Step 1: Create a Minimal Molecular Snapshot

A snapshot requires only NumPy arrays and metadata. In src/exojax/database/contracts.py, the MDBSnapshot dataclass combines MDBMeta (molecular mass, partition function grids) with Lines (transition wavenumbers, lower state energies, reference line strengths).

You can construct a snapshot from any external line list, laboratory measurement, or theoretical calculation:

import numpy as np
from exojax.database.contracts import MDBMeta, Lines, MDBSnapshot

# User-provided line data

nu_lines = np.array([2000.0, 2005.0, 2010.0])          # cm⁻¹

elower   = np.array([  10.0,   15.0,   20.0])          # cm⁻¹

strength = np.array([1e-30, 2e-30, 3e-30])           # at Tref

# Required metadata

meta = MDBMeta(
    dbtype="exomol",
    molmass=44.0,             # g mol⁻¹ (example: CO₂)

    T_gQT=np.array([300., 1000., 2000.]),
    gQT   =np.array([1., 2., 4.]),
)

lines = Lines(
    nu_lines=nu_lines,
    elower=elower,
    line_strength_ref_original=strength,
)

# Optional broadening arrays for ExoMol-type databases

n_Texp   = np.full(nu_lines.size, 0.5)
alpha_ref= np.full(nu_lines.size, 0.1)

snap = MDBSnapshot(meta=meta, lines=lines, n_Texp=n_Texp, alpha_ref=alpha_ref)

Step 2: Define Custom Physics Providers

Providers are callable objects that ExoJAX queries for temperature-dependent quantities. The partition function provider must implement qr_single(T, Tref) and qr_vector(Tarr, Tref), returning the ratio Q(T)/Q(Tref). The broadening provider implements compute(Tref_broadening), returning temperature exponents and reference Lorentz widths.

Default implementations like ExomolPartitionProvider and HitranBroadening live in src/exojax/opacity/providers.py, but you can substitute these with custom classes:

import numpy as np

class FakePF:
    """Implements PartitionFunctionProvider interface"""
    def __init__(self, scalar=2.0):
        self.scalar = scalar

    def qr_single(self, T, Tref):
        return np.asarray(self.scalar)

    def qr_vector(self, Tarr, Tref):
        return np.full_like(Tarr, self.scalar, dtype=float)


class FakeBroad:
    """Implements BroadeningStrategy interface"""
    def __init__(self, n_lines):
        self.n = n_lines

    def compute(self, Tref_broadening):
        n_Texp    = np.full(self.n, 0.123)
        gamma_ref = np.full(self.n, 4.567)
        return n_Texp, gamma_ref

Step 3: Instantiate the Opacity Calculator

With a snapshot and providers ready, pass them to OpaPremodit.from_snapshot(). This factory method, defined in src/exojax/opacity/premodit/api.py, constructs the PreMODIT interpolation grids without loading heavy database objects.

You must configure the PreMODIT parameters via manual_setting() or auto_setting() before computing cross-sections:

from exojax.opacity.premodit.api import OpaPremodit

nu_grid = np.linspace(1995., 2015., 64)

opa = OpaPremodit.from_snapshot(
    snap,
    nu_grid,
    pf_provider=FakePF(scalar=3.0),
    broadening_strategy=FakeBroad(n_lines=snap.lines.nu_lines.size),
    allow_32bit=True,
)

opa.manual_setting(dE=5.0, Tref=1000.0, Twt=1200.0)

# Compute cross-section at specific T, P

T, P = 1500.0, 0.01
xs = opa.xsvector(T, P)

Working with Multiple Custom Species

For atmospheres containing several molecules, use MultiMol from src/exojax/database/multimol.py. This helper automatically iterates over species lists, calls the appropriate factory methods, and handles spectral segment stitching. It accepts both snapshots and legacy MDB objects that expose to_snapshot():

from exojax.database.multimol import MultiMol

molmulti = [["CO2"], ["H2O", "CO"]]
dbmulti  = [["exomol"], ["exomol", "exomol"]]

mm = MultiMol(molmulti, dbmulti, database_root_path="./.database")
mdb_coll = mm.multimdb(nu_grid_list=[nu_grid, nu_grid])

opa_list = mm.multiopa_premodit(
    multimdb=mdb_coll,
    nu_grid_list=[nu_grid, nu_grid],
    auto_trange=[800., 2500.],
)

# Access first CO₂ opacity object

xs_co2 = opa_list[0][0].xsvector(1500., 0.01)

Alternative: Legacy Custom MDB Classes

If you maintain existing database reader classes, add a to_snapshot() method that returns MDBSnapshot. MultiMol.store_single_opa detects this method and routes the object through OpaPremodit.from_mdb(), preserving your existing infrastructure while enabling provider injection:

class MyCustomMdb:
    def __init__(self, filename, nu_range):
        # Load proprietary line list

        pass

    def to_snapshot(self):
        meta = MDBMeta(
            dbtype="custom",
            molmass=28.0,
            T_gQT=np.array([300., 1000.]),
            gQT   =np.array([1., 2.]),
        )
        lines = Lines(
            nu_lines=self.nu_lines,
            elower=self.elower,
            line_strength_ref_original=self.Sij0,
        )
        return MDBSnapshot(
            meta=meta, 
            lines=lines,
            n_Texp=self.n_Texp, 
            alpha_ref=self.alpha_ref
        )

Summary

• Snapshot-first design: MDBSnapshot in src/exojax/database/contracts.py provides a language-agnostic container for any line list using only NumPy arrays.
• Provider injection: Override partition functions and broadening via pf_provider and broadening_strategy in OpaPremodit.from_snapshot(), avoiding edits to src/exojax/opacity/providers.py.
• Factory methods: Use OpaPremodit.from_snapshot() for snapshots or OpaPremodit.from_mdb() for legacy objects, both defined in src/exojax/opacity/premodit/api.py.
• Batch management: MultiMol in src/exojax/database/multimol.py automates opacity construction for multiple species while supporting mixed custom and standard databases.
• Backward compatibility: Custom MDB classes work immediately by implementing to_snapshot(), requiring no changes to existing file parsers.

Frequently Asked Questions

Can I use laboratory measurement line lists with ExoJAX?

Yes. Any line list—whether from laboratory spectroscopy, theoretical calculations, or custom synthetic spectra—can be used by wrapping the wavenumbers, lower state energies, and reference strengths in an MDBSnapshot as shown in src/exojax/database/contracts.py. You must also provide the molecular mass and partition function tables in MDBMeta.

How do I validate my custom opacity calculations?

Compare the output of xsvector() or xsmatrix() against reference calculations using the same line data. The test file tests/unittests/opacity/premodit/test_premodit_providers.py demonstrates the minimal setup required for unit testing custom providers. Verify that your partition function ratios (qr_single) return the correct Q(T)/Q(Tref) values for your specific molecule.

Does using custom physics providers impact performance?

No. The provider interface uses pure NumPy or JAX-compatible operations. Once the OpaPremodit object is initialized with manual_setting() or auto_setting(), the PreMODIT interpolation grids are pre-computed, and subsequent calls to xsvector() execute with the same performance as standard ExoMol or HITRAN opacities.

Can I mix standard ExoMol data with custom snapshots in the same retrieval?

Yes. MultiMol handles heterogeneous databases automatically. When store_single_opa encounters a standard MdbExomol object, it uses default providers; when it encounters your custom snapshot or MDB with to_snapshot(), it applies your injected pf_provider and broadening_strategy. The resulting opa_list contains homogeneous OpaPremodit objects ready for radiative transfer calculations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}