skyscapes.physical_model

skyscapes.physical_model#

Physical-model hierarchy: phase/wavelength-dependent planet-to-star contrast.

Submodules#

Classes#

`AbstractPhysicalModel`	Spectral-physics layer for a planet.
`PrecomputedPhysicalModel`	Physical model whose reflectivity spectrum has been pre-computed.
`ExoJaxPhysicalModel`	Composition-based reflected-light planet model over ExoJAX's 2-stream RT.
`GridPhysicalModel`	Per-planet 2D interpolated contrast over (wavelength, phase-angle).
`LambertianPhysicalModel`	Lambertian reflector.

Package Contents#

class skyscapes.physical_model.AbstractPhysicalModel[source]#

Bases: equinox.Module

Spectral-physics layer for a planet.

Concrete subclasses own per-planet physical-model parameters (albedo, grid contrast cube, ExoJax composition, etc.). The host planet’s radius is supplied at call time via contrast(..., Rp_Rearth=...).

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

Return planet-to-star contrast.

Args:

phase_angle_rad: Star-planet-observer phase angle beta,: shape (K, T) (matching AbstractOrbit.propagate).

dist_AU: Star-planet distance, shape (K, T). wavelength_nm: Wavelength, scalar or shape (W,). Rp_Rearth: Planet radius [Earth radii], shape (K,),

supplied by the host Planet.

Returns:

Flux-ratio contrast, shape (K, T).

Parameters:

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

Return type:

jaxtyping.Array

class skyscapes.physical_model.PrecomputedPhysicalModel[source]#

Bases: skyscapes.physical_model.base.AbstractPhysicalModel

Physical model whose reflectivity spectrum has been pre-computed.

Attributes:

reflectivity: Per-planet plane-parallel reflectivity R(nu): on the wavenumber grid, shape (K, n_nu). Already includes all the physics of the original physical model (absorption, Rayleigh, clouds, surface).

nu_grid: Wavenumber grid [cm^-1], shape (n_nu,). n_nu: Length of nu_grid (static for JIT).

reflectivity: jaxtyping.Array#

nu_grid: jaxtyping.Array#

n_nu: int#

classmethod from_physical_model(model)[source]#

Pre-compute the reflectivity from an existing physical model.

Calls the model’s internal _reflectivity_all_planets once and packages the result. Requires the model to expose _reflectivity_all_planets, nu_grid, and n_nu – presently ExoJaxPhysicalModel is the supported source.

Parameters:: model (skyscapes.physical_model.base.AbstractPhysicalModel)
Return type:: PrecomputedPhysicalModel

classmethod load(path)[source]#

Load a previously-saved cache file.

Args:

path: Path to a .npz file produced by save().

Returns:

A PrecomputedPhysicalModel ready to evaluate.

Raises:

ValueError: if the file’s cache_format_version differs: from the current code’s version (would silently produce wrong spectra otherwise).

Parameters:: path (str | pathlib.Path)
Return type:: PrecomputedPhysicalModel

save(path)[source]#

Save the cached reflectivity to a .npz file.

Idempotent and safe to call from JAX-traced contexts (the save itself is plain NumPy – callers should not call this inside a JIT region, but the data is just regular arrays).

Parameters:: path (str | pathlib.Path)
Return type:: None

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, shape (K, T). dist_AU: Star-planet distance [AU], shape (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 cached array stores the underlying model’s plane-parallel (spherical) reflectivity; we convert to geometric albedo via the Lambertian-sphere factor 2/3 (Seager 2010, eq 3.36) at call time – same convention as ExoJaxPhysicalModel.contrast.

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 geometric-albedo contrast across wavelengths.

Returns shape (W, K, T). Avoids per-wavelength recomputation by vectorising the interpolation; applies the Lambertian-sphere spherical-to-geometric conversion 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]#

Compact summary of the cached spectrum.

Return type:: str

class skyscapes.physical_model.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

class skyscapes.physical_model.GridPhysicalModel[source]#

Bases: skyscapes.physical_model.base.AbstractPhysicalModel

Per-planet 2D interpolated contrast over (wavelength, phase-angle).

Distance and planet radius are ignored – the grid already encodes a flux ratio.

Attributes:: wavelengths_nm: 1-D wavelength grid [nm], shape (n_wl,). phase_angle_deg: 1-D phase-angle grid [deg], shape (n_phase,). contrast_grid: Contrast cube, shape (K, n_wl, n_phase).

wavelengths_nm: jaxtyping.Array#

phase_angle_deg: jaxtyping.Array#

contrast_grid: jaxtyping.Array#

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

Per-planet contrast at (wavelength, phase).

Args:: phase_angle_rad: Phase angle per planet [rad], shape (K, T). dist_AU: Shape (K, T); ignored (grid encodes flux ratio). wavelength_nm: Scalar wavelength. Rp_Rearth: Shape (K,); ignored (grid encodes flux ratio).
Returns:: Contrast, shape (K, T).

Parameters:

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

Return type:

jaxtyping.Array

__repr__()[source]#

One-line summary of grid shape (K, n_wl, n_phase) and wl range.

Return type:: str

class skyscapes.physical_model.LambertianPhysicalModel[source]#

Bases: skyscapes.physical_model.base.AbstractPhysicalModel

Lambertian reflector.

Attributes:: Ag: Geometric albedo, shape (K,).

Ag: jaxtyping.Array#

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

Contrast = Ag * Phi(beta) * (Rp/r)**2. Wavelength-independent.

phase_angle_rad and dist_AU are shape (K, T). wavelength_nm is part of the interface but ignored in the grey case. Rp_Rearth is shape (K,), supplied by the host Planet.

Parameters:

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

Return type:

jaxtyping.Array

__repr__()[source]#

One-line summary of geometric albedo.

Return type:: str