skyscapes.physical_model.exojax

skyscapes.physical_model.exojax#

ExoJAX-backed physical-model components.

Self-contained subpackage holding everything specific to the ExoJAX radiative-transfer backend: the ExoJaxPhysicalModel orchestrator, swappable physics components (T-P, absorption, scattering, clouds, surface), PSG cross-section helpers, and Earth-epoch archetypes.

The backend-agnostic AbstractPhysicalModel interface and the simple non-RT physical models (Lambertian, grid) live one level up in skyscapes.physical_model.

Submodules#

Attributes#

`MOLECULE_RECIPES`
`O3_MOLMASS`
`DEFAULT_MOLECULES`
`EARTH_ARCHEAN_VMRS`
`EARTH_EARLY_PROTEROZOIC_VMRS`
`EARTH_LATE_PROTEROZOIC_VMRS`
`EARTH_MODERN_VMRS`

Classes#

`Absorption`	Sum of per-species line-list / cross-section absorption.
`AbstractAbsorption`	Per-layer absorption opacity from line lists / cross-sections.
`AbstractClouds`	Per-layer cloud opacity (mixed scattering + absorption).
`AbstractMmrProfile`	Per-species mass-mixing-ratio profile evaluated at layer pressures.
`AbstractScattering`	Per-layer scattering opacity (e.g. Rayleigh).
`AbstractSurface`	Bottom-of-atmosphere reflectivity.
`AbstractTPProfile`	Temperature-pressure profile.
`BulkGasRecipe`	Build instructions for the implicit residual gas.
`BulkGasResidual`	Implicit residual gas filling the unallocated mass fraction.
`ConstantMmr`	Well-mixed gas: constant mmr at every layer.
`Contribution`	Per-layer optical-property contribution from one component.
`FlatSurface`	Wavelength-independent (gray) Lambertian surface.
`GrayCloud`	Single-layer gray-scattering cloud.
`MieCloud`	Single-layer cloud with Mie-scattering optical properties.
`MolecularSpecies`	One atmospheric molecule with an altitude-resolved mixing ratio.
`MoleculeRecipe`	Build instructions for one molecule.
`NoCloud`	Disable cloud opacity: zero contribution.
`NullScattering`	Disable scattering entirely: contributes zero opacity.
`PowerLawTPProfile`	Power-law T-P profile: `T(P) = T_eq * P^alpha`.
`RayleighScattering`	Rayleigh scattering from tracked species + optional bulk residual.
`StratosphericPeakMmr`	Gaussian-in-log-pressure peak; canonical for O3.
`TroposphericMmr`	Constant below a step pressure, drops sharply above; canonical for H2O.
`WavelengthDependentSurface`	Wavelength-dependent surface reflectivity.
`O3ChappuisOpacity`	PSG cross-section opacity for O3 (Serdyuchenko et al. 2014).
`ExoJaxPhysicalModel`	Composition-based reflected-light planet model over ExoJAX's 2-stream RT.
`PsgCrossSectionOpacity`	Cross-section-backed opacity from any PSG xs file.

Functions#

`build_bulk_prebuilt`(*, name, nu_grid)	Construct (molmass, rayleigh_xs) for the implicit residual gas.
`build_mie_cloud`(*, nu_grid, log_pressure_bar, ...[, ...])	Build a Mie cloud component by pre-computing Mie params.
`build_species_prebuilt`(*, name, nu_grid, nu_min, ...)	Construct (molmass, opa, rayleigh_xs) for one molecule.
`build_exojax_engines`(*[, molecules, bulk_gas, ...])	Build the heavy shared engines and pre-built per-species data.
`vmr_dict_to_earth_profile_dict`(vmrs[, K, bulk_gas])	VMRs -> dict of Earth-realistic profiles per species.
`vmr_dict_to_log_mmr_dict`(vmrs[, K, bulk_gas])	Convert VMR dict to log10 MMR dict with K-shaped arrays.
`vmr_dict_to_mmr_dict`(vmrs[, bulk_gas])	Convert volume-mixing-ratio dict to mass-mixing-ratio dict.

Package Contents#

skyscapes.physical_model.exojax.MOLECULE_RECIPES: dict[str, MoleculeRecipe]#

class skyscapes.physical_model.exojax.Absorption[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractAbsorption

Sum of per-species line-list / cross-section absorption.

Iterates over the species tuple, skipping any with opa is None (e.g. a species included purely for its Rayleigh contribution).

compute(species, Tarr, pressure, gravity, rt_engine)[source]#

Sum per-species absorption optical depth.

Parameters:

species (tuple[skyscapes.physical_model.exojax.components.species.MolecularSpecies, Ellipsis])
Tarr (jaxtyping.Array)
pressure (jaxtyping.Array)
gravity (jaxtyping.Array)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.AbstractAbsorption[source]#

Bases: equinox.Module

Per-layer absorption opacity from line lists / cross-sections.

Concrete implementations iterate over the atmosphere’s MolecularSpecies tuple, calling each species’ opa engine and summing the contributions.

abstractmethod compute(species, Tarr, pressure, gravity, rt_engine)[source]#

Absorption contribution.

dtau_scatter and g_weighted_num are zero for pure absorbers, but the return shape is the same as for scattering components so the atmosphere can combine them uniformly.

Parameters:

species (tuple[skyscapes.physical_model.exojax.components.species.MolecularSpecies, Ellipsis])
Tarr (jaxtyping.Array)
pressure (jaxtyping.Array)
gravity (jaxtyping.Array)

Return type:

Contribution

class skyscapes.physical_model.exojax.AbstractClouds[source]#

Bases: equinox.Module

Per-layer cloud opacity (mixed scattering + absorption).

abstractmethod compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu)[source]#

Cloud contribution.

For a single-scattering-albedo cloud the scattering and absorption parts are both nonzero; for a purely scattering cloud dtau_total == dtau_scatter.

Parameters:

log_pressure_bar_scalar (jaxtyping.Array)
log_opt_depth_scalar (jaxtyping.Array)
pressure (jaxtyping.Array)
n_nu (int)

Return type:

Contribution

class skyscapes.physical_model.exojax.AbstractMmrProfile[source]#

Bases: equinox.Module

Per-species mass-mixing-ratio profile evaluated at layer pressures.

abstractmethod evaluate(pressure)[source]#

Return shape (n_layers,) mmr values at the layer pressures.

Parameters:: pressure (jaxtyping.Array)
Return type:: jaxtyping.Array

class skyscapes.physical_model.exojax.AbstractScattering[source]#

Bases: equinox.Module

Per-layer scattering opacity (e.g. Rayleigh).

abstractmethod compute(species, bulk, gravity, rt_engine, n_layers, n_nu)[source]#

Scattering contribution from tracked species + optional bulk gas.

For a pure scatterer dtau_total == dtau_scatter; g_weighted_num is zero for isotropic Rayleigh.

Parameters:

species (tuple[skyscapes.physical_model.exojax.components.species.MolecularSpecies, Ellipsis])
bulk (skyscapes.physical_model.exojax.components.species.BulkGasResidual | None)
gravity (jaxtyping.Array)
n_layers (int)
n_nu (int)

Return type:

Contribution

class skyscapes.physical_model.exojax.AbstractSurface[source]#

Bases: equinox.Module

Bottom-of-atmosphere reflectivity.

abstractmethod compute_refl(log_albedo_scalar, n_nu)[source]#

Return surface reflectivity, shape (n_nu,).

Parameters:

log_albedo_scalar (jaxtyping.Array)
n_nu (int)

Return type:

jaxtyping.Array

class skyscapes.physical_model.exojax.AbstractTPProfile[source]#

Bases: equinox.Module

Temperature-pressure profile.

abstractmethod compute_Tarr(rt_engine, T_eq_K_scalar, T_alpha_scalar)[source]#

Return layer temperatures, shape (n_layers,).

Parameters:

T_eq_K_scalar (jaxtyping.Array)
T_alpha_scalar (jaxtyping.Array)

Return type:

jaxtyping.Array

class skyscapes.physical_model.exojax.BulkGasRecipe[source]#

Build instructions for the implicit residual gas.

Attributes:

name: Gas name; must match an ExoJAX polarizability key (or: provide polarizability_override).

molmass: Molar mass [g/mol]. polarizability_override: Override polarizability if missing

from ExoJAX’s table.

name: str#

molmass: float#

polarizability_override: float | None = None#

class skyscapes.physical_model.exojax.BulkGasResidual[source]#

Bases: equinox.Module

Implicit residual gas filling the unallocated mass fraction.

The mass-mixing ratio is computed dynamically as max(0, 1 - sum(tracked species mmrs)). Contributes only to Rayleigh scattering (no line-list absorption is associated with the bulk gas in this model – N2 is essentially transparent across 300–1100 nm, H2/He likewise).

Attributes:: name: Gas name, e.g. "N2" for Earth or "H2" for gas giants. molmass: Molar mass [g/mol]. rayleigh_xs: Rayleigh cross-section [cm^2/molecule] on the

atmosphere’s wavenumber grid, shape (n_nu,).

name: str#

molmass: float#

rayleigh_xs: jaxtyping.Array#

class skyscapes.physical_model.exojax.ConstantMmr[source]#

Bases: AbstractMmrProfile

Well-mixed gas: constant mmr at every layer.

Use for gases with chemical lifetimes longer than mixing timescales in the modelled pressure range (CO2, CH4, O2, N2, etc. in troposphere + lower stratosphere).

Attributes:: log_mmr: Log10 mass-mixing ratio, shape (K,) (per planet).

log_mmr: jaxtyping.Array#

evaluate(pressure)[source]#

Return 10**log_mmr broadcast to pressure.shape.

Parameters:: pressure (jaxtyping.Array)
Return type:: jaxtyping.Array

class skyscapes.physical_model.exojax.Contribution[source]#

Bases: NamedTuple

Per-layer optical-property contribution from one component.

Fields:

dtau_total: Layer optical depth this component adds to the total: opacity budget, shape (n_layers, n_nu).
dtau_scatter: Subset of dtau_total that is scattering: (non-absorbing), shape (n_layers, n_nu). For a pure absorber this is zero. For Rayleigh ssa=1 so it equals dtau_total. For a partly-absorbing cloud with single- scattering albedo ssa_c, this is ssa_c * dtau_cloud.
g_weighted_num: g * dtau_scatter numerator for the weighted-: average asymmetry parameter, shape (n_layers, n_nu). Rayleigh contributes zero (isotropic, g=0). A cloud with asymmetry g_c contributes g_c * ssa_c * dtau_cloud.

dtau_total: jaxtyping.Array#

dtau_scatter: jaxtyping.Array#

g_weighted_num: jaxtyping.Array#

class skyscapes.physical_model.exojax.FlatSurface[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractSurface

Wavelength-independent (gray) Lambertian surface.

Equivalent to WavelengthDependentSurface with a flat spectrum but structurally explicit; useful when callers want to document “I am intentionally using a featureless surface”.

Attributes:: log_albedo: Log10 surface albedo per planet, shape (K,).

log_albedo: jaxtyping.Array#

compute_refl(log_albedo_scalar, n_nu)[source]#

Return scalar albedo broadcast across the wavenumber grid.

Parameters:

log_albedo_scalar (jaxtyping.Array)
n_nu (int)

Return type:

jaxtyping.Array

class skyscapes.physical_model.exojax.GrayCloud[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractClouds

Single-layer gray-scattering cloud.

The total cloud scattering optical depth is distributed across layers via a Gaussian in log-pressure (softmax-normalised so the weights are well-defined even when the cloud pressure is far outside the layer grid – useful for ablation runs with log_opt_depth = -inf).

Wavelength-grey: the cloud’s cross-section is constant across the spectrum. A Mie cloud with composition-dependent scattering would be a separate component implementing the same contract.

Attributes (PyTree leaves, fittable):: log_pressure_bar: Log10 cloud-deck pressure [bar], shape (K,). log_opt_depth: Log10 of the vertically-integrated cloud

scattering optical depth, shape (K,).
Static attributes (configuration):: ssa: Single-scattering albedo (default 1.0, pure scattering). g: Asymmetry parameter (default 0.85, forward-peaked). log_sigma: Vertical-distribution width in log10(P) [dex].

log_pressure_bar: jaxtyping.Array#

log_opt_depth: jaxtyping.Array#

ssa: float#

g: float#

log_sigma: float#

compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu)[source]#

Gray cloud contribution at a single pressure level.

Parameters:

log_pressure_bar_scalar (jaxtyping.Array)
log_opt_depth_scalar (jaxtyping.Array)
pressure (jaxtyping.Array)
n_nu (int)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.MieCloud[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractClouds

Single-layer cloud with Mie-scattering optical properties.

The total cloud scattering optical depth is distributed vertically via a softmax-Gaussian in log-pressure (same as GrayCloud). What’s different is that the single-scattering albedo and asymmetry parameter come from pre-computed Mie cross-sections rather than being scalar constants – they vary with wavelength according to the condensate’s refractive index n(lambda) + k(lambda).

Attributes (PyTree leaves, fittable):: log_pressure_bar: Log10 cloud-deck pressure [bar], shape (K,). log_opt_depth: Log10 of the vertically-integrated cloud

extinction optical depth, shape (K,).

Pre-computed Mie quantities (built at engine time, shared across planets):

ssa_grid: Single-scattering albedo on the wavenumber grid,
shape (n_nu,). sigma_scattering / sigma_extinction.

g_grid: Asymmetry parameter on the wavenumber grid,
shape (n_nu,).

Static config:

condensate: Condensate name ("H2O", "H2O_ice",: "MgSiO3", etc.). Matters for the repr; the actual Mie params are baked into ssa_grid and g_grid.

rg_um: Geometric mean particle radius [um]. sigmag: Geometric standard deviation of the lognormal size

distribution.

log_sigma: Cloud-deck vertical Gaussian half-width in: log10(pressure) [dex].

log_pressure_bar: jaxtyping.Array#

log_opt_depth: jaxtyping.Array#

ssa_grid: jaxtyping.Array#

g_grid: jaxtyping.Array#

condensate: str#

rg_um: float#

sigmag: float#

log_sigma: float#

compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu)[source]#

Mie-cloud contribution: gray-depth distribution, spectral ssa/g.

Parameters:

log_pressure_bar_scalar (jaxtyping.Array)
log_opt_depth_scalar (jaxtyping.Array)
pressure (jaxtyping.Array)
n_nu (int)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.MolecularSpecies[source]#

Bases: equinox.Module

One atmospheric molecule with an altitude-resolved mixing ratio.

The mixing ratio is encoded as a profile component (e.g. ConstantMmr for well-mixed gases, StratosphericPeakMmr for O3, TroposphericMmr for H2O) rather than a single scalar, so altitude variation is represented explicitly.

Attributes:

profile: Per-species mmr profile. Owns the fittable: log-mixing-ratio leaves.

name: Molecule name, e.g. "H2O". molmass: Molar mass [g/mol]. opa: Opacity engine. Either an ExoJAX OpaPremodit (for

HITRAN-backed line-list molecules) or O3ChappuisOpacity (for visible cross-section absorbers like O3). May be None if the species contributes only Rayleigh scattering.

rayleigh_xs: Pre-computed Rayleigh cross-section: [cm^2/molecule] on the atmosphere’s wavenumber grid, shape (n_nu,). Set to all-zeros to disable per-species Rayleigh contribution.

profile: skyscapes.physical_model.exojax.components.mmr_profile.AbstractMmrProfile#

name: str#

molmass: float#

opa: Any#

rayleigh_xs: jaxtyping.Array#

class skyscapes.physical_model.exojax.MoleculeRecipe[source]#

Build instructions for one molecule.

Attributes:

name: Molecule name (must match ExoJAX’s HITRAN / polarizability: keys when those tables apply).
psg_xs_url: If set, use a PSG cross-section file at this URL: (PsgCrossSectionOpacity) instead of an ExoJAX line-list opa. Required for species whose absorption is dominated by electronic transitions in the visible/NIR (e.g. O3 Chappuis, SO2 UV).
psg_xs_molmass: Molar mass [g/mol]; required when psg_xs_url: is set since we can’t extract it from HITRAN in that case.
polarizability_override: Explicit polarizability [cm^3] for the: Rayleigh cross-section, used when ExoJAX’s table doesn’t have an entry. None lets OpaRayleigh look it up.

name: str#

psg_xs_url: str | None = None#

psg_xs_molmass: float | None = None#

polarizability_override: float | None = None#

class skyscapes.physical_model.exojax.NoCloud[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractClouds

Disable cloud opacity: zero contribution.

Equivalent to a GrayCloud with log_opt_depth = -inf but structurally explicit, so that isinstance(atm.clouds, NoCloud) documents intent in code that builds cloud-free atmospheres.

compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu)[source]#

Return zero-everywhere cloud contribution.

log_pressure_bar_scalar and log_opt_depth_scalar are accepted for parity with GrayCloud but unused.

Parameters:

log_pressure_bar_scalar (jaxtyping.Array)
log_opt_depth_scalar (jaxtyping.Array)
pressure (jaxtyping.Array)
n_nu (int)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.NullScattering[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractScattering

Disable scattering entirely: contributes zero opacity.

Useful for ablation studies (e.g. quantifying how much Rayleigh affects a retrieval) and for atmospheres where scattering is negligible.

compute(species, bulk, gravity, rt_engine, n_layers, n_nu)[source]#

Return zero-everywhere contribution.

Parameters:

species (tuple[skyscapes.physical_model.exojax.components.species.MolecularSpecies, Ellipsis])
bulk (skyscapes.physical_model.exojax.components.species.BulkGasResidual | None)
gravity (jaxtyping.Array)
n_layers (int)
n_nu (int)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.PowerLawTPProfile[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractTPProfile

Power-law T-P profile: T(P) = T_eq * P^alpha.

The two parameters are PyTree leaves but they are passed to compute_Tarr() as scalar args rather than being read from self so the atmosphere can vmap over per-planet axes without extra plumbing.

Attributes:: T_eq_K: Reference temperature at 1 bar [K], shape (K,). T_alpha: Power-law exponent, shape (K,).

T_eq_K: jaxtyping.Array#

T_alpha: jaxtyping.Array#

compute_Tarr(rt_engine, T_eq_K_scalar, T_alpha_scalar)[source]#

T(P) = T_eq * P^alpha on the rt_engine’s layer pressures.

Parameters:

T_eq_K_scalar (jaxtyping.Array)
T_alpha_scalar (jaxtyping.Array)

Return type:

jaxtyping.Array

class skyscapes.physical_model.exojax.RayleighScattering[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractScattering

Rayleigh scattering from tracked species + optional bulk residual.

Iterates over the atmosphere’s species tuple plus the bulk gas (if present), computing each gas’s contribution from its own rayleigh_xs and molmass. The bulk gas’s mass-mixing ratio is computed dynamically as max(0, 1 - sum(tracked mmrs)).

To disable per-species Rayleigh while keeping the bulk: set the species’ rayleigh_xs to zeros. To disable scattering entirely: use NullScattering.

compute(species, bulk, gravity, rt_engine, n_layers, n_nu)[source]#

Sum tracked-species + bulk-gas Rayleigh contributions.

Bulk gas mmr is computed per-layer as max(0, 1 - sum(species profiles at this pressure)), so altitude variation of the tracked species translates into altitude variation of the bulk-gas residual.

Parameters:

species (tuple[skyscapes.physical_model.exojax.components.species.MolecularSpecies, Ellipsis])
bulk (skyscapes.physical_model.exojax.components.species.BulkGasResidual | None)
gravity (jaxtyping.Array)
n_layers (int)
n_nu (int)

Return type:

skyscapes.physical_model.exojax.components.base.Contribution

class skyscapes.physical_model.exojax.StratosphericPeakMmr[source]#

Bases: AbstractMmrProfile

Gaussian-in-log-pressure peak; canonical for O3.

mmr(P) = 10**log_peak_mmr * exp(-0.5 * ((log10(P) - log_peak_pressure_bar) / log_sigma_decades)**2).

For Earth’s O3: log_peak_pressure_bar = log10(0.01) = -2 (10 mbar, ~30 km altitude), log_sigma_decades = 0.5.

Attributes:: log_peak_mmr: Log10 of the peak mmr, shape (K,). log_peak_pressure_bar: Log10 pressure [bar] at peak, (K,). log_sigma_decades: Gaussian width in log10-pressure, (K,).

log_peak_mmr: jaxtyping.Array#

log_peak_pressure_bar: jaxtyping.Array#

log_sigma_decades: jaxtyping.Array#

evaluate(pressure)[source]#

Gaussian peak in log-pressure.

Parameters:: pressure (jaxtyping.Array)
Return type:: jaxtyping.Array

class skyscapes.physical_model.exojax.TroposphericMmr[source]#

Bases: AbstractMmrProfile

Constant below a step pressure, drops sharply above; canonical for H2O.

Uses a sigmoid in log-pressure for smooth transition (differentiable for HMC retrievals). For Earth’s H2O: tropospheric ~3e-3 below 0.1 bar, ~3e-6 above.

Attributes:: log_mmr_below: Log10 mmr at high pressure (below cold trap), (K,). log_mmr_above: Log10 mmr at low pressure (above cold trap), (K,). log_pressure_step_bar: Log10 transition pressure [bar], (K,). log_transition_width_decades: Width of the sigmoid transition

in log10-pressure, (K,). Smaller = sharper step.

log_mmr_below: jaxtyping.Array#

log_mmr_above: jaxtyping.Array#

log_pressure_step_bar: jaxtyping.Array#

log_transition_width_decades: jaxtyping.Array#

evaluate(pressure)[source]#

Sigmoid step in log-pressure.

Parameters:: pressure (jaxtyping.Array)
Return type:: jaxtyping.Array

class skyscapes.physical_model.exojax.WavelengthDependentSurface[source]#

Bases: skyscapes.physical_model.exojax.components.base.AbstractSurface

Wavelength-dependent surface reflectivity.

The bottom-of-atmosphere reflectivity is 10**log_albedo * spectrum: the log_albedo PyTree leaf is a per-planet fittable scaling, and spectrum is a fixed wavelength-dependent profile (vegetation red-edge, water absorption, snow/ice, …) shared across planets.

Attributes:: log_albedo: Log10 surface albedo scaling per planet, shape (K,). spectrum: Wavelength-dependent reflectivity profile,

shape (n_nu,). Defaults to flat ones for “no spectral shape; albedo is constant across the band”.

log_albedo: jaxtyping.Array#

spectrum: jaxtyping.Array#

compute_refl(log_albedo_scalar, n_nu)[source]#

Return wavelength-dependent surface reflectivity.

n_nu is accepted for parity with FlatSurface but unused – the wavelength axis comes from self.spectrum.

Parameters:

log_albedo_scalar (jaxtyping.Array)
n_nu (int)

Return type:

jaxtyping.Array

skyscapes.physical_model.exojax.build_bulk_prebuilt(*, name, nu_grid)[source]#

Construct (molmass, rayleigh_xs) for the implicit residual gas.

Parameters:

name (str)
nu_grid (jaxtyping.Array)

Return type:

tuple[float, jaxtyping.Array]

skyscapes.physical_model.exojax.build_mie_cloud(*, nu_grid, log_pressure_bar, log_opt_depth, condensate='H2O', rg_um=10.0, sigmag=2.0, log_sigma=DEFAULT_CLOUD_LOG_SIGMA)[source]#

Build a Mie cloud component by pre-computing Mie params.

On first use for a given condensate this triggers two downloads / computations cached under ./.database/particulates/virga/:

Refractive-index file (small, fetched from Zenodo).

Mie-grid lookup table (built locally via PyMieScatt; takes a couple of minutes per condensate). Subsequent calls reuse the cache.

Args:

nu_grid: Wavenumber grid [cm^-1] from: build_exojax_engines().

log_pressure_bar: Per-planet log10 cloud pressure [bar]. log_opt_depth: Per-planet log10 total cloud extinction tau. condensate: Condensate name. ExoJAX/Virga ships e.g.

"H2O", "H2O_ice", "NH3", "MgSiO3", "Mg2SiO4", "Fe", "KCl", "Na2S", "ZnS", "MnS", "Cr", "Al2O3", "TiO2".

rg_um: Mean particle radius [um] of the lognormal size: distribution.

sigmag: Geometric standard deviation of the size distribution. log_sigma: Vertical-distribution Gaussian half-width [dex].

Returns:

A MieCloud instance ready to drop into ExoJaxPhysicalModel.

Parameters:

nu_grid (jaxtyping.Array)
log_pressure_bar (jaxtyping.Array)
log_opt_depth (jaxtyping.Array)
condensate (str)
rg_um (float)
sigmag (float)
log_sigma (float)

Return type:

MieCloud

skyscapes.physical_model.exojax.build_species_prebuilt(*, name, nu_grid, nu_min, nu_max, databases_dir, crit)[source]#

Construct (molmass, opa, rayleigh_xs) for one molecule.

Used by build_exojax_engines; not typically called by users directly. Returns the static parts of a MolecularSpecies (everything except log_mmr).

Parameters:

name (str)
nu_grid (jaxtyping.Array)
nu_min (float)
nu_max (float)
databases_dir (str)
crit (float)

Return type:

tuple[float, Any, jaxtyping.Array]

skyscapes.physical_model.exojax.O3_MOLMASS = 47.9982#

class skyscapes.physical_model.exojax.O3ChappuisOpacity(nu_grid, cache_dir=None, xs_table_path=None)[source]#

Bases: skyscapes.physical_model.exojax.psg_xs.PsgCrossSectionOpacity

PSG cross-section opacity for O3 (Serdyuchenko et al. 2014).

Parameters:

nu_grid (jaxtyping.Array | numpy.ndarray)
cache_dir (str | pathlib.Path | None)
xs_table_path (str | pathlib.Path | None)

skyscapes.physical_model.exojax.DEFAULT_MOLECULES: tuple[str, Ellipsis] = ('H2O', 'CO2', 'CH4', 'O2', 'O3')#

class skyscapes.physical_model.exojax.ExoJaxPhysicalModel[source]#

Bases: skyscapes.physical_model.base.AbstractPhysicalModel

Composition-based reflected-light planet model over ExoJAX’s 2-stream RT.

Per-planet state (PyTree leaves, fittable):

log_gravity_cgs: Log10 surface gravity [cm/s^2], shape (K,). species: Tuple of MolecularSpecies. Each species owns

its own log_mmr (per-planet, shape (K,)). The number and identity of species is configurable via build_exojax_engines().

bulk: Optional BulkGasResidual (implicit residual gas).

Components (each owns its own per-planet PyTree leaves where applicable):

tp_profile: T-P profile component (e.g. PowerLawTPProfile). absorption: Absorption orchestrator (e.g. Absorption). scattering: Scattering component (e.g. RayleighScattering,

NullScattering).

clouds: Cloud component (e.g. GrayCloud, NoCloud). surface: Surface component (e.g. WavelengthDependentSurface).

Shared / configuration attributes:: rt_engine: ExoJAX ArtReflectPure instance. nu_grid: Wavenumber grid [cm^-1]. n_nu: Length of nu_grid (static for JIT).

log_gravity_cgs: jaxtyping.Array#

species: tuple[skyscapes.physical_model.exojax.components.MolecularSpecies, Ellipsis]#

bulk: skyscapes.physical_model.exojax.components.BulkGasResidual | None#

tp_profile: skyscapes.physical_model.exojax.components.AbstractTPProfile#

absorption: skyscapes.physical_model.exojax.components.AbstractAbsorption#

scattering: skyscapes.physical_model.exojax.components.AbstractScattering#

clouds: skyscapes.physical_model.exojax.components.AbstractClouds#

surface: skyscapes.physical_model.exojax.components.AbstractSurface#

rt_engine: Any#

nu_grid: jaxtyping.Array#

n_nu: int#

classmethod from_default_setup(*, log_mmrs, T_eq_K, T_alpha, log_surface_albedo, log_gravity_cgs, log_cloud_pressure_bar=None, log_cloud_opt_depth=None, surface_albedo_spectrum=None, molecules=None, bulk_gas='N2', wavelength_min_nm=400.0, wavelength_max_nm=1000.0, n_wavenumbers=2000, n_layers=100, pressure_top_bar=1e-05, pressure_btm_bar=1.0, databases_dir=None, crit=0.0)[source]#

One-shot convenience: build engines + default components in one call.

Defaults to Earth-like physics: PowerLawTPProfile, Absorption, RayleighScattering with N2 bulk, GrayCloud (Earth water clouds), and a flat WavelengthDependentSurface.

Args:

log_mmrs: Dict mapping molecule name to per-planet log10: mass-mixing ratio, shape (K,) each. The dict’s molecules determine which species are built.

T_eq_K: Per-planet equatorial T [K], (K,). T_alpha: Per-planet T-P power-law exponent, (K,). log_surface_albedo: Per-planet surface scaling, (K,). log_gravity_cgs: Per-planet log10 gravity, (K,). log_cloud_pressure_bar: Per-planet cloud-deck pressure, (K,). log_cloud_opt_depth: Per-planet cloud optical depth, (K,). surface_albedo_spectrum: (n_nu,) spectral profile.

Defaults to flat ones.

molecules: Override molecule list (default: derived from: log_mmrs.keys()).

bulk_gas: Implicit residual gas (default "N2"). wavelength_min_nm: See build_exojax_engines(). wavelength_max_nm: See build_exojax_engines(). n_wavenumbers: See build_exojax_engines(). n_layers: See build_exojax_engines(). pressure_top_bar: See build_exojax_engines(). pressure_btm_bar: See build_exojax_engines(). databases_dir: See build_exojax_engines(). crit: See build_exojax_engines().

Parameters:

log_mmrs (dict[str, jaxtyping.Array])
T_eq_K (jaxtyping.Array)
T_alpha (jaxtyping.Array)
log_surface_albedo (jaxtyping.Array)
log_gravity_cgs (jaxtyping.Array)
log_cloud_pressure_bar (jaxtyping.Array | None)
log_cloud_opt_depth (jaxtyping.Array | None)
surface_albedo_spectrum (jaxtyping.Array | None)
molecules (tuple[str, Ellipsis] | None)
bulk_gas (str | None)
wavelength_min_nm (float)
wavelength_max_nm (float)
n_wavenumbers (int)
n_layers (int)
pressure_top_bar (float)
pressure_btm_bar (float)
databases_dir (str | None)
crit (float)

Return type:

ExoJaxPhysicalModel

classmethod from_default_setup_cached(*, log_mmrs, T_eq_K, T_alpha, log_surface_albedo, log_gravity_cgs, log_cloud_pressure_bar=None, log_cloud_opt_depth=None, surface_albedo_spectrum=None, molecules=None, bulk_gas='N2', wavelength_min_nm=400.0, wavelength_max_nm=1000.0, n_wavenumbers=2000, n_layers=100, pressure_top_bar=1e-05, pressure_btm_bar=1.0, databases_dir=None, crit=0.0, cache_dir=None)[source]#

Cached one-shot factory: returns a fast PrecomputedPhysicalModel.

Hashes every input and looks up a cached spectrum on disk. On cache hit, returns the cached spectrum in ~10 ms. On cache miss, builds the full ExoJaxPhysicalModel via from_default_setup(), runs the 2-stream RT once, precomputes the reflectivity, saves to disk, and returns the PrecomputedPhysicalModel.

Use this when the physical-model parameters are fixed across many evaluations (coronagraphoto sims, ETC studies). For HMC retrievals where parameters vary, use from_default_setup() directly.

Args:

cache_dir: Override the cache directory. Defaults to: ~/.cache/skyscapes/physical_models/.

Returns:

A PrecomputedPhysicalModel with no RT cost at evaluation time.

Parameters:

log_mmrs (dict[str, jaxtyping.Array | skyscapes.physical_model.exojax.components.AbstractMmrProfile])
T_eq_K (jaxtyping.Array)
T_alpha (jaxtyping.Array)
log_surface_albedo (jaxtyping.Array)
log_gravity_cgs (jaxtyping.Array)
log_cloud_pressure_bar (jaxtyping.Array | None)
log_cloud_opt_depth (jaxtyping.Array | None)
surface_albedo_spectrum (jaxtyping.Array | None)
molecules (tuple[str, Ellipsis] | None)
bulk_gas (str | None)
wavelength_min_nm (float)
wavelength_max_nm (float)
n_wavenumbers (int)
n_layers (int)
pressure_top_bar (float)
pressure_btm_bar (float)
databases_dir (str | None)
crit (float)
cache_dir (str | pathlib.Path | None)

Return type:

skyscapes.physical_model.cached.PrecomputedPhysicalModel

_reflectivity_one_planet(k)[source]#

Plane-parallel reflectivity for the k-th planet, shape (n_nu,).

Indexes every K-shaped leaf at position k (cheap; JIT traces this as plain array indexing). With per-species mmr profiles the previous “vmap over an extracted log_mmrs array” pattern broke because different profile types have different K-shaped fields – a Python loop over K avoids that issue without losing JIT efficiency at K=1 (the common case).

Parameters:: k (int)
Return type:: jaxtyping.Array

_reflectivity_all_planets()[source]#

Plane-parallel reflectivity for every planet, shape (K, n_nu).

Loops over K planets in Python. JIT unrolls the loop; for the common K=1 case this is a single iteration with no overhead. For HMC retrievals (K=1 per sample), the JIT cache is reused across MCMC steps.

Return type:: jaxtyping.Array

contrast(phase_angle_rad, dist_AU, wavelength_nm, Rp_Rearth)[source]#

Per-planet, per-time geometric-albedo contrast at one wavelength.

Args:: phase_angle_rad: Star-planet-observer phase angle, (K, T). dist_AU: Star-planet distance [AU], (K, T). wavelength_nm: Scalar wavelength [nm]. Rp_Rearth: Planet radius [Earth radii], shape (K,).
Returns:: Contrast = A_g(lambda) * Lambert_phase(beta) * (Rp/d)^2, shape (K, T). The 2-stream RT produces the plane- parallel (spherical) reflectivity; we convert to geometric albedo via the Lambertian-sphere factor 2/3 (Seager 2010, eq 3.36) so the output is a geometric-albedo contrast – the same convention as LambertianPhysicalModel.

Parameters:

phase_angle_rad (jaxtyping.Array)
dist_AU (jaxtyping.Array)
wavelength_nm (jaxtyping.Array)
Rp_Rearth (jaxtyping.Array)

Return type:

jaxtyping.Array

contrast_cube(phase_angle_rad, dist_AU, wavelengths_nm, Rp_Rearth)[source]#

Per-planet, per-time contrast across many wavelengths.

Returns (W, K, T) geometric-albedo contrast cube. Computes the underlying 2-stream RT exactly once per planet rather than once per wavelength; applies the spherical-to-geometric Lambertian-sphere conversion (Seager 2010, eq 3.36) the same way as contrast().

Parameters:

phase_angle_rad (jaxtyping.Array)
dist_AU (jaxtyping.Array)
wavelengths_nm (jaxtyping.Array)
Rp_Rearth (jaxtyping.Array)

Return type:

jaxtyping.Array

__repr__()[source]#

Human-readable summary of components + species + per-planet state.

Return type:: str

skyscapes.physical_model.exojax.build_exojax_engines(*, molecules=DEFAULT_MOLECULES, bulk_gas='N2', wavelength_min_nm=400.0, wavelength_max_nm=1000.0, n_wavenumbers=2000, n_layers=100, pressure_top_bar=1e-05, pressure_btm_bar=1.0, databases_dir=None, crit=0.0)[source]#

Build the heavy shared engines and pre-built per-species data.

Each molecule’s opacity engine (line-list LSD tables or cross- section interpolant) is built once here and reused across many atmosphere instantiations that vary only the per-planet log_mmrs. The retrieval-friendly construction pattern is:

engines = build_exojax_engines(molecules=("H2O", "CO", "O2"))

def make_atmosphere(per_planet_log_mmrs):
    return ExoJaxPhysicalModel.from_default_setup(
        log_mmrs=per_planet_log_mmrs, ..., **engines
    )

Args:

molecules: Tuple of molecule names to include. Must all appear: in MOLECULE_RECIPES. Default is the 5-molecule biosignature set.
bulk_gas: Implicit residual gas name ("N2", "H2",: "He"), or None to disable the bulk-gas Rayleigh contribution. Default "N2" for terrestrial atmospheres.

wavelength_min_nm: Short-wavelength end of the spectral range [nm]. wavelength_max_nm: Long-wavelength end of the spectral range [nm]. n_wavenumbers: Number of points in the wavenumber grid. n_layers: Number of layers in the plane-parallel RT solver. pressure_top_bar: Pressure at the top of the model [bar]. pressure_btm_bar: Pressure at the bottom of the model [bar]. databases_dir: ExoJAX line-list cache. crit: Line-strength cutoff [cm/(molecule.cm^-2)] passed to

MdbHitran (default 0.0 = no filtering). Set crit=1e-26 for ~25% speedup at HWO contrast levels.

Returns:

Dict ready to be **-unpacked into ExoJaxPhysicalModel. Keys: rt_engine, nu_grid, n_nu, species_prebuilt (dict mapping name to {molmass, opa, rayleigh_xs}), bulk_prebuilt (dict with {name, molmass, rayleigh_xs} or None).

Parameters:

molecules (tuple[str, Ellipsis])
bulk_gas (str | None)
wavelength_min_nm (float)
wavelength_max_nm (float)
n_wavenumbers (int)
n_layers (int)
pressure_top_bar (float)
pressure_btm_bar (float)
databases_dir (str | None)
crit (float)

Return type:

dict

skyscapes.physical_model.exojax.EARTH_ARCHEAN_VMRS: dict[str, float]#

skyscapes.physical_model.exojax.EARTH_EARLY_PROTEROZOIC_VMRS: dict[str, float]#

skyscapes.physical_model.exojax.EARTH_LATE_PROTEROZOIC_VMRS: dict[str, float]#

skyscapes.physical_model.exojax.EARTH_MODERN_VMRS: dict[str, float]#

skyscapes.physical_model.exojax.vmr_dict_to_earth_profile_dict(vmrs, K=1, bulk_gas='N2')[source]#

VMRs -> dict of Earth-realistic profiles per species.

Most species get a ConstantMmr (well-mixed assumption). For H2O and O3, where Earth’s altitude profile differs sharply from uniform, the function builds:

H2O: TroposphericMmr with the supplied VMR below the cold trap and a 3-decade drop above (parametrized at P = 0.1 bar).

O3: StratosphericPeakMmr with the supplied VMR as the peak value at P = 10 mbar (Gaussian width 0.5 dex).

All values are taken from canonical Earth-atmosphere structure – customize by constructing your own profile dict if you need to fit per-planet profile parameters in an HMC retrieval.

Args:

vmrs: {name: VMR}; H2O VMR is interpreted as the: tropospheric value, O3 VMR as the stratospheric peak.

K: Number of planets (per-planet broadcasting). bulk_gas: Implicit residual gas name.

Returns:

{name: AbstractMmrProfile} ready to pass as log_mmrs= to ExoJaxPhysicalModel.from_default_setup().

Parameters:

vmrs (dict[str, float])
K (int)
bulk_gas (str)

Return type:

dict[str, skyscapes.physical_model.exojax.components.AbstractMmrProfile]

skyscapes.physical_model.exojax.vmr_dict_to_log_mmr_dict(vmrs, K=1, bulk_gas='N2')[source]#

Convert VMR dict to log10 MMR dict with K-shaped arrays.

Each value in the returned dict is shape (K,) so it can be passed directly as log_mmrs= to ExoJaxPhysicalModel.from_default_setup(). The same VMR is used for every planet – to give per-planet variation, modify the returned arrays before constructing the atmosphere.

Every species gets a ConstantMmr profile by default (uniform mmr at every altitude). For Earth-realistic profiles (O3 stratospheric peak, H2O cold-trap drop) use vmr_dict_to_earth_profile_dict() instead.

Parameters:

vmrs (dict[str, float])
K (int)
bulk_gas (str)

Return type:

dict[str, jaxtyping.Array]

skyscapes.physical_model.exojax.vmr_dict_to_mmr_dict(vmrs, bulk_gas='N2')[source]#

Convert volume-mixing-ratio dict to mass-mixing-ratio dict.

Computes mean molecular weight from the supplied VMRs plus the bulk-gas residual (mmr_bulk = 1 - sum(vmrs_tracked), in VMR space; then convert all to mass). Used internally by vmr_dict_to_log_mmr_dict(); exposed here for inspection.

Args:

vmrs: {name: VMR} for tracked molecules. Sum must be < 1: (the residual is the bulk gas).
bulk_gas: Name of the implicit residual gas ("N2" for: terrestrial atmospheres, "H2" for gas giants, …). Must appear in BULK_GAS_RECIPES.

Returns:

{name: MMR} for the tracked molecules. The bulk gas’s MMR is not returned – ExoJaxPhysicalModel computes it dynamically as the residual.

Parameters:

vmrs (dict[str, float])
bulk_gas (str)

Return type:

dict[str, float]

class skyscapes.physical_model.exojax.PsgCrossSectionOpacity(name, url, molmass, nu_grid, cache_dir=None, xs_table_path=None, cache_filename=None)[source]#

Cross-section-backed opacity from any PSG xs file.

Quacks like ExoJAX’s OpaPremodit: exposes xsmatrix(Tarr, pressure) -> (n_layers, n_nu) and a molmass attribute, so it slots into the same opa_engines tuple as the line-list-backed opacity engines.

Pressure broadening is ignored. The PSG cross-section files describe electronic-continuum / pre-broadened absorption; the pressure-dependence at typical atmospheric pressures is below the photon-noise floor for HWO contrast levels. Temperature interpolation is linear against the file’s temperature axis, with out-of-range temperatures clamped to the nearest endpoint.

For files with a single temperature column the cross-section is returned independent of layer temperature.

Attributes:

name: Molecule name (e.g. "O3", "SO2"). molmass: Molar mass [g/mol]. Caller supplies; used by

opacity_profile_xs to convert mmr to column density.

nu_grid: Wavenumber grid [cm^-1] this adapter was built for.

Parameters:

name (str)
url (str)
molmass (float)
nu_grid (jaxtyping.Array | numpy.ndarray)
cache_dir (str | pathlib.Path | None)
xs_table_path (str | pathlib.Path | None)
cache_filename (str | None)

name#

molmass#

nu_grid#

_T_grid#

_sigma_at_nu#

xsmatrix(Tarr, pressure)[source]#

Cross-section matrix at the requested layer temperatures.

pressure is accepted for interface parity with OpaPremodit but is unused.

Args:: Tarr: Layer temperatures [K], shape (n_layers,). pressure: Layer pressures [bar], shape (n_layers,);

ignored.
Returns:: Cross-section matrix [cm^2/molecule], shape (n_layers, n_nu). Temperatures outside the PSG file’s range clamp to the nearest endpoint cross-section. For single-temperature files the returned matrix is constant in T.

Parameters:

Tarr (jaxtyping.Array)
pressure (jaxtyping.Array)

Return type:

jaxtyping.Array