CKD Transmission Tutorial: ArtTransPure with OpaCKD
Hajime Kawahara with Claude Code, July 2 (2025)
This tutorial demonstrates how to use the Correlated K-Distribution (CKD) method for atmospheric transmission calculations with ExoJAX. Transmission spectroscopy is a key technique for characterizing exoplanet atmospheres by observing starlight passing through the planetary atmosphere.
# Import required packages
import numpy as np
import matplotlib.pyplot as plt
from jax import config
# ExoJAX imports
from exojax.test.emulate_mdb import mock_mdbExomol, mock_wavenumber_grid
from exojax.opacity import OpaCKD, OpaPremodit
from exojax.rt import ArtTransPure
from exojax.test.data import get_testdata_filename, TESTDATA_CO_EXOMOL_PREMODIT_TRANSMISSION_REF
# Enable 64-bit precision for accurate calculations
config.update("jax_enable_x64", True)
print("ExoJAX CKD Tutorial: Transmission Spectroscopy")
print("=============================================")
ExoJAX CKD Tutorial: Transmission Spectroscopy
=============================================
1. Setup Atmospheric Model and Molecular Database
First, we’ll set up our atmospheric model for transmission spectroscopy calculations.
# Setup wavenumber grid and molecular database
nu_grid, wav, res = mock_wavenumber_grid()
print(f"Wavenumber grid: {len(nu_grid)} points from {nu_grid[0]:.1f} to {nu_grid[-1]:.1f} cm⁻¹")
print(f"Spectral resolution: {res:.1f}")
# Create mock H2O molecular database
mdb = mock_mdbExomol("H2O")
print(f"Molecular database: {mdb.nurange[0]:.1f} - {mdb.nurange[1]:.1f} cm⁻¹")
# Setup atmospheric radiative transfer for transmission
art = ArtTransPure(
pressure_top=1.0e-8,
pressure_btm=1.0e2,
nlayer=50, # Fewer layers for transmission calculations
integration="simpson" # Simpson integration for better accuracy
)
print(f"Atmospheric layers: {art.nlayer}")
print(f"Pressure range: {art.pressure_top:.1e} - {art.pressure_btm:.1e} bar")
print(f"Integration method: {art.integration}")
xsmode = modit xsmode assumes ESLOG in wavenumber space: xsmode=modit Your wavelength grid is in * ascending * order The wavenumber grid is in ascending order by definition. Please be careful when you use the wavelength grid. Wavenumber grid: 20000 points from 4329.0 to 4363.0 cm⁻¹ Spectral resolution: 2556525.8 xsmode = modit xsmode assumes ESLOG in wavenumber space: xsmode=modit Your wavelength grid is in * ascending * order The wavenumber grid is in ascending order by definition. Please be careful when you use the wavelength grid. radis== 0.15.2 HITRAN exact name= H2(16O) radis engine = vaex
/home/kawahara/exojax/src/exojax/utils/grids.py:85: UserWarning: Both input wavelength and output wavenumber are in ascending order.
warnings.warn(
/home/kawahara/exojax/src/exojax/utils/grids.py:85: UserWarning: Both input wavelength and output wavenumber are in ascending order.
warnings.warn(
/home/kawahara/exojax/src/exojax/utils/grids.py:85: UserWarning: Both input wavelength and output wavenumber are in ascending order.
warnings.warn(
/home/kawahara/exojax/src/exojax/utils/grids.py:85: UserWarning: Both input wavelength and output wavenumber are in ascending order.
warnings.warn(
/home/kawahara/exojax/src/exojax/database/api.py:134: UserWarning: The current version of radis does not support broadf_download (requires >=0.16).
warnings.warn(msg, UserWarning)
/home/kawahara/exojax/src/exojax/utils/molname.py:197: FutureWarning: e2s will be replaced to exact_molname_exomol_to_simple_molname.
warnings.warn(
/home/kawahara/exojax/src/exojax/utils/molname.py:91: FutureWarning: exojax.utils.molname.exact_molname_exomol_to_simple_molname will be replaced to radis.api.exomolapi.exact_molname_exomol_to_simple_molname.
warnings.warn(
/home/kawahara/exojax/src/exojax/utils/molname.py:91: FutureWarning: exojax.utils.molname.exact_molname_exomol_to_simple_molname will be replaced to radis.api.exomolapi.exact_molname_exomol_to_simple_molname.
warnings.warn(
Molecule: H2O
Isotopologue: 1H2-16O
Background atmosphere: H2
ExoMol database: None
Local folder: H2O/1H2-16O/SAMPLE
Transition files:
=> File 1H2-16O__SAMPLE__04300-04400.trans
Broadener: H2
Broadening code level: a1
DataFrame (self.df) available.
Molecular database: 4329.0 - 4363.0 cm⁻¹
integration: simpson
Simpson integration, uses the chord optical depth at the lower boundary and midppoint of the layers.
Atmospheric layers: 50
Pressure range: 1.0e-08 - 1.0e+02 bar
Integration method: simpson
/home/kawahara/exojax/src/exojax/rt/common.py:40: UserWarning: nu_grid is not given. specify nu_grid when using 'run'
warnings.warn(
2. Define Atmospheric and Planetary Parameters
We’ll create atmospheric profiles and define planetary parameters for transmission calculations.
# Create atmospheric profiles
Tarr = np.linspace(1000.0, 1500.0, 50) # Temperature profile
mmr_arr = np.full(50, 0.1) # Constant H2O mixing ratio
mean_molecular_weight = np.full(50, 2.33) # Mean molecular weight (H2-dominated)
# Planetary parameters (Jupiter-like)
radius_btm = 6.9e9 # Planet radius at bottom of atmosphere (cm)
gravity = 2478.57 # Surface gravity (cm/s²)
# Plot atmospheric profiles
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))
# Temperature profile
ax1.semilogy(Tarr, art.pressure)
ax1.set_xlabel('Temperature (K)')
ax1.set_ylabel('Pressure (bar)')
ax1.set_title('Temperature Profile')
ax1.grid(True, alpha=0.3)
ax1.invert_yaxis()
# Mixing ratio profile
ax2.semilogy(mmr_arr, art.pressure)
ax2.set_xlabel('H₂O Mixing Ratio')
ax2.set_ylabel('Pressure (bar)')
ax2.set_title('H₂O Mixing Ratio Profile')
ax2.grid(True, alpha=0.3)
ax2.invert_yaxis()
# Mean molecular weight profile
ax3.semilogy(mean_molecular_weight, art.pressure)
ax3.set_xlabel('Mean Molecular Weight (amu)')
ax3.set_ylabel('Pressure (bar)')
ax3.set_title('Mean Molecular Weight Profile')
ax3.grid(True, alpha=0.3)
ax3.invert_yaxis()
plt.tight_layout()
plt.show()
print(f"Temperature range: {np.min(Tarr):.0f} - {np.max(Tarr):.0f} K")
print(f"H2O mixing ratio: {mmr_arr[0]:.1f} (constant)")
print(f"Mean molecular weight: {mean_molecular_weight[0]:.2f} amu (constant)")
print(f"Planet radius: {radius_btm/6.9e9:.1f} R_Jupiter")
print(f"Surface gravity: {gravity:.0f} cm/s² ({gravity/2478.57:.1f} × Jupiter)")
Temperature range: 1000 - 1500 K
H2O mixing ratio: 0.1 (constant)
Mean molecular weight: 2.33 amu (constant)
Planet radius: 1.0 R_Jupiter
Surface gravity: 2479 cm/s² (1.0 × Jupiter)
3. Setup Standard Line-by-Line Opacity Calculator
First, we’ll compute the standard high-resolution transmission spectrum using line-by-line calculations.
# Initialize standard opacity calculator (Premodit)
base_opa = OpaPremodit(mdb, nu_grid, auto_trange=[800.0, 1600.0])
print(f"Base opacity calculator: {base_opa.__class__.__name__}")
# Compute line-by-line cross-sections and transmission spectrum
print("\nComputing line-by-line transmission spectrum...")
xsmatrix = base_opa.xsmatrix(Tarr, art.pressure)
dtau = art.opacity_profile_xs(xsmatrix, mmr_arr, base_opa.mdb.molmass, gravity)
transit_lbl = art.run(dtau, Tarr, mean_molecular_weight, radius_btm, gravity)
print(f"Line-by-line spectrum computed!")
print(f"Transit radius ratio range: [{np.min(transit_lbl):.6f}, {np.max(transit_lbl):.6f}]")
print(f"Transit depth variation: {(np.max(transit_lbl) - np.min(transit_lbl))*1e6:.0f} ppm")
OpaPremodit: params automatically set.
default elower grid trange (degt) file version: 2
Robust range: 771.9537482657882 - 1647.2060977798953 K
OpaPremodit: Tref_broadening is set to 1131.3708498984759 K
max value of ngamma_ref_grid : 21.825321843011604
min value of ngamma_ref_grid : 13.242701248020088
ngamma_ref_grid grid : [13.24270058 15.00453705 17.00077107 19.26258809 21.8253231 ]
max value of n_Texp_grid : 0.541
min value of n_Texp_grid : 0.216
n_Texp_grid grid : [0.21599999 0.54100007]
uniqidx: 0%| | 0/3 [00:00<?, ?it/s]uniqidx: 100%|██████████| 3/3 [00:00<00:00, 8473.34it/s]
Premodit: Twt= 1383.2165049575465 K Tref= 840.335329973883 K
Making LSD:|####################| 100%
Base opacity calculator: OpaPremodit
Computing line-by-line transmission spectrum...
Line-by-line spectrum computed!
Transit radius ratio range: [1.042101, 1.109748]
Transit depth variation: 67647 ppm
4. Setup CKD Opacity Calculator and Compute Transmission
Now we’ll initialize the CKD opacity calculator and compute the CKD transmission spectrum.
# Initialize CKD opacity calculator
opa_ckd = OpaCKD(
base_opa, # Base opacity calculator
Ng=16, # Number of g-ordinates for quadrature
band_width=0.5 # Spectral band width
)
print(f"CKD Opacity Calculator Setup:")
print(f" Number of g-ordinates (Ng): {opa_ckd.Ng}")
print(f" Band width: {opa_ckd.band_width}")
print(f" Number of spectral bands: {len(opa_ckd.nu_bands)}")
print(f" Spectral range: {opa_ckd.nu_bands[0]:.1f} - {opa_ckd.nu_bands[-1]:.1f} cm⁻¹")
# Pre-compute CKD tables on temperature-pressure grid
print("\nPre-computing CKD tables...")
T_grid = np.linspace(np.min(Tarr), np.max(Tarr), 10)
P_grid = np.logspace(np.log10(np.min(art.pressure)), np.log10(np.max(art.pressure)), 10)
opa_ckd.precompute_tables(T_grid, P_grid)
# Get CKD cross-section tensor and compute CKD spectrum
print("Computing CKD transmission spectrum...")
xs_ckd = opa_ckd.xstensor_ckd(Tarr, art.pressure)
dtau_ckd = art.opacity_profile_xs_ckd(xs_ckd, mmr_arr, base_opa.mdb.molmass, gravity)
transit_ckd = art.run_ckd(dtau_ckd, Tarr, mean_molecular_weight, radius_btm, gravity, opa_ckd.ckd_info.weights)
print(f"CKD spectrum computed!")
print(f"CKD transit range: [{np.min(transit_ckd):.6f}, {np.max(transit_ckd):.6f}]")
CKD Opacity Calculator Setup:
Number of g-ordinates (Ng): 16
Band width: 0.5
Number of spectral bands: 68
Spectral range: 4329.3 - 4362.8 cm⁻¹
Pre-computing CKD tables...
Generated g-grid: 16 points, range [0.0053, 0.9947]
Processing 68 spectral bands...
Band 1: [4329.0, 4329.5] cm⁻¹, 295 frequencies
Band 2: [4329.5, 4330.0] cm⁻¹, 294 frequencies
Band 3: [4330.0, 4330.5] cm⁻¹, 294 frequencies
Band 4: [4330.5, 4331.0] cm⁻¹, 294 frequencies
Band 5: [4331.0, 4331.5] cm⁻¹, 294 frequencies
Band 6: [4331.5, 4332.0] cm⁻¹, 294 frequencies
Band 7: [4332.0, 4332.5] cm⁻¹, 294 frequencies
Band 8: [4332.5, 4333.0] cm⁻¹, 294 frequencies
Band 9: [4333.0, 4333.5] cm⁻¹, 294 frequencies
Band 10: [4333.5, 4334.0] cm⁻¹, 295 frequencies
Band 11: [4334.0, 4334.5] cm⁻¹, 294 frequencies
Band 12: [4334.5, 4335.0] cm⁻¹, 294 frequencies
Band 13: [4335.0, 4335.5] cm⁻¹, 294 frequencies
Band 14: [4335.5, 4336.0] cm⁻¹, 294 frequencies
Band 15: [4336.0, 4336.5] cm⁻¹, 294 frequencies
Band 16: [4336.5, 4337.0] cm⁻¹, 294 frequencies
Band 17: [4337.0, 4337.5] cm⁻¹, 294 frequencies
Band 18: [4337.5, 4338.0] cm⁻¹, 294 frequencies
Band 19: [4338.0, 4338.5] cm⁻¹, 294 frequencies
Band 20: [4338.5, 4339.0] cm⁻¹, 295 frequencies
Band 21: [4339.0, 4339.5] cm⁻¹, 294 frequencies
Band 22: [4339.5, 4340.0] cm⁻¹, 294 frequencies
Band 23: [4340.0, 4340.5] cm⁻¹, 294 frequencies
Band 24: [4340.5, 4341.0] cm⁻¹, 294 frequencies
Band 25: [4341.0, 4341.5] cm⁻¹, 294 frequencies
Band 26: [4341.5, 4342.0] cm⁻¹, 294 frequencies
Band 27: [4342.0, 4342.5] cm⁻¹, 294 frequencies
Band 28: [4342.5, 4343.0] cm⁻¹, 294 frequencies
Band 29: [4343.0, 4343.5] cm⁻¹, 294 frequencies
Band 30: [4343.5, 4344.0] cm⁻¹, 295 frequencies
Band 31: [4344.0, 4344.5] cm⁻¹, 294 frequencies
Band 32: [4344.5, 4345.0] cm⁻¹, 294 frequencies
Band 33: [4345.0, 4345.5] cm⁻¹, 294 frequencies
Band 34: [4345.5, 4346.0] cm⁻¹, 294 frequencies
Band 35: [4346.0, 4346.5] cm⁻¹, 294 frequencies
Band 36: [4346.5, 4347.0] cm⁻¹, 294 frequencies
Band 37: [4347.0, 4347.5] cm⁻¹, 294 frequencies
Band 38: [4347.5, 4348.0] cm⁻¹, 294 frequencies
Band 39: [4348.0, 4348.5] cm⁻¹, 295 frequencies
Band 40: [4348.5, 4349.0] cm⁻¹, 294 frequencies
Band 41: [4349.0, 4349.5] cm⁻¹, 294 frequencies
Band 42: [4349.5, 4350.0] cm⁻¹, 294 frequencies
Band 43: [4350.0, 4350.5] cm⁻¹, 294 frequencies
Band 44: [4350.5, 4351.0] cm⁻¹, 294 frequencies
Band 45: [4351.0, 4351.5] cm⁻¹, 294 frequencies
Band 46: [4351.5, 4352.0] cm⁻¹, 294 frequencies
Band 47: [4352.0, 4352.5] cm⁻¹, 294 frequencies
Band 48: [4352.5, 4353.0] cm⁻¹, 294 frequencies
Band 49: [4353.0, 4353.5] cm⁻¹, 295 frequencies
Band 50: [4353.5, 4354.0] cm⁻¹, 294 frequencies
Band 51: [4354.0, 4354.5] cm⁻¹, 294 frequencies
Band 52: [4354.5, 4355.0] cm⁻¹, 294 frequencies
Band 53: [4355.0, 4355.5] cm⁻¹, 294 frequencies
Band 54: [4355.5, 4356.0] cm⁻¹, 294 frequencies
Band 55: [4356.0, 4356.5] cm⁻¹, 294 frequencies
Band 56: [4356.5, 4357.0] cm⁻¹, 294 frequencies
Band 57: [4357.0, 4357.5] cm⁻¹, 294 frequencies
Band 58: [4357.5, 4358.0] cm⁻¹, 294 frequencies
Band 59: [4358.0, 4358.5] cm⁻¹, 295 frequencies
Band 60: [4358.5, 4359.0] cm⁻¹, 294 frequencies
Band 61: [4359.0, 4359.5] cm⁻¹, 294 frequencies
Band 62: [4359.5, 4360.0] cm⁻¹, 294 frequencies
Band 63: [4360.0, 4360.5] cm⁻¹, 294 frequencies
Band 64: [4360.5, 4361.0] cm⁻¹, 294 frequencies
Band 65: [4361.0, 4361.5] cm⁻¹, 294 frequencies
Band 66: [4361.5, 4362.0] cm⁻¹, 294 frequencies
Band 67: [4362.0, 4362.5] cm⁻¹, 294 frequencies
Band 68: [4362.5, 4363.0] cm⁻¹, 295 frequencies
Creating CKD table info...
CKD precomputation complete! Ready for interpolation.
Table dimensions: T=10, P=10, g=16, bands=68
Computing CKD transmission spectrum...
CKD spectrum computed!
CKD transit range: [1.042468, 1.071653]
5. Compare Results and Visualize
Let’s compare the CKD results with the line-by-line spectrum and compute band averages for validation.
# Compute reference band averages by direct integration
print("Computing reference band averages...")
transit_avg = []
band_edges = opa_ckd.band_edges
for band_idx in range(len(opa_ckd.nu_bands)):
mask = (band_edges[band_idx, 0] <= nu_grid) & (nu_grid < band_edges[band_idx, 1])
transit_avg.append(np.mean(transit_lbl[mask]))
transit_avg = np.array(transit_avg)
# Calculate accuracy metrics
res = np.sqrt(np.sum((transit_ckd - transit_avg)**2)/len(transit_ckd))/np.mean(transit_avg)
max_relative_error = np.max(np.abs((transit_ckd - transit_avg) / transit_avg))
resolution = opa_ckd.nu_bands[0]/(band_edges[0, 1] - band_edges[0, 0])
transit_diff_ppm = np.abs((transit_ckd - transit_avg) * 1e6)
print(f"CKD Accuracy Assessment:")
print(f" RMS relative error: {res:.6f}")
print(f" Maximum relative error: {max_relative_error:.6f}")
print(f" Effective resolution: {resolution:.1f}")
print(f" Maximum transit depth difference: {np.max(transit_diff_ppm):.1f} ppm")
print(f" Mean transit depth difference: {np.mean(transit_diff_ppm):.1f} ppm")
Computing reference band averages...
CKD Accuracy Assessment:
RMS relative error: 0.000111
Maximum relative error: 0.000226
Effective resolution: 8692.6
Maximum transit depth difference: 240.4 ppm
Mean transit depth difference: 105.1 ppm
6. Visualize Transmission Spectra Comparison
# Create comparison plot
plt.figure(figsize=(14, 8))
# Plot line-by-line spectrum (high resolution)
plt.plot(nu_grid, transit_lbl,
label="Line-by-Line (Premodit)",
alpha=0.7, linewidth=0.8, color='lightblue')
# Plot CKD spectrum
plt.plot(opa_ckd.nu_bands, transit_ckd,
'o-', label="CKD Method",
markersize=4, linewidth=2, color='red')
# Plot reference band averages
plt.plot(opa_ckd.nu_bands, transit_avg,
's-', label="Reference Band Average",
markersize=3, linewidth=1.5, color='black', alpha=0.8)
plt.xlabel('Wavenumber (cm⁻¹)', fontsize=12)
plt.ylabel('(R_p/R_*)²', fontsize=12)
plt.title(f'CKD vs Line-by-Line Transmission Spectrum\\n'
f'Resolution: {resolution:.0f}, RMS Error: {res:.6f}', fontsize=14)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)
# Add text box with key parameters
textstr = f'Ng = {opa_ckd.Ng}\\nBands = {len(opa_ckd.nu_bands)}\\nLayers = {art.nlayer}\\nMax Δ = {np.max(transit_diff_ppm):.1f} ppm'
props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
plt.text(0.02, 0.98, textstr, transform=plt.gca().transAxes, fontsize=10,
verticalalignment='top', bbox=props)
plt.tight_layout()
plt.show()
# Save the figure
plt.savefig(f"ckd_transmission_comparison_res{resolution:.0f}.png",
dpi=300, bbox_inches='tight')
print(f"Figure saved as: ckd_transmission_comparison_res{resolution:.0f}.png")
Figure saved as: ckd_transmission_comparison_res8693.png
<Figure size 640x480 with 0 Axes>
Summary
This tutorial demonstrated how to use the CKD method with ExoJAX for transmission spectroscopy:
Key Steps:
Setup: Initialize atmospheric model and molecular database for transmission
Profiles: Define temperature, mixing ratio, and planetary parameters
Line-by-Line: Compute high-resolution transmission spectrum
CKD Setup: Initialize CKD calculator and pre-compute tables
CKD Calculation: Compute band-averaged transmission spectrum using CKD
Validation: Compare CKD results with line-by-line band averages
Visualization: Plot comparison and analyze accuracy in ppm