skyscapes.physical_model.exojax =============================== .. py:module:: skyscapes.physical_model.exojax .. autoapi-nested-parse:: ExoJAX-backed physical-model components. Self-contained subpackage holding everything specific to the ExoJAX radiative-transfer backend: the :class:`ExoJaxPhysicalModel` orchestrator, swappable physics components (T-P, absorption, scattering, clouds, surface), PSG cross-section helpers, and Earth-epoch archetypes. The backend-agnostic :class:`AbstractPhysicalModel` interface and the simple non-RT physical models (Lambertian, grid) live one level up in :mod:`skyscapes.physical_model`. Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/skyscapes/physical_model/exojax/components/index /autoapi/skyscapes/physical_model/exojax/o3_chappuis/index /autoapi/skyscapes/physical_model/exojax/physical_model/index /autoapi/skyscapes/physical_model/exojax/presets/index /autoapi/skyscapes/physical_model/exojax/psg_xs/index Attributes ---------- .. autoapisummary:: skyscapes.physical_model.exojax.MOLECULE_RECIPES skyscapes.physical_model.exojax.O3_MOLMASS skyscapes.physical_model.exojax.DEFAULT_MOLECULES skyscapes.physical_model.exojax.EARTH_ARCHEAN_VMRS skyscapes.physical_model.exojax.EARTH_EARLY_PROTEROZOIC_VMRS skyscapes.physical_model.exojax.EARTH_LATE_PROTEROZOIC_VMRS skyscapes.physical_model.exojax.EARTH_MODERN_VMRS Classes ------- .. autoapisummary:: skyscapes.physical_model.exojax.Absorption skyscapes.physical_model.exojax.AbstractAbsorption skyscapes.physical_model.exojax.AbstractClouds skyscapes.physical_model.exojax.AbstractMmrProfile skyscapes.physical_model.exojax.AbstractScattering skyscapes.physical_model.exojax.AbstractSurface skyscapes.physical_model.exojax.AbstractTPProfile skyscapes.physical_model.exojax.BulkGasRecipe skyscapes.physical_model.exojax.BulkGasResidual skyscapes.physical_model.exojax.ConstantMmr skyscapes.physical_model.exojax.Contribution skyscapes.physical_model.exojax.FlatSurface skyscapes.physical_model.exojax.GrayCloud skyscapes.physical_model.exojax.MieCloud skyscapes.physical_model.exojax.MolecularSpecies skyscapes.physical_model.exojax.MoleculeRecipe skyscapes.physical_model.exojax.NoCloud skyscapes.physical_model.exojax.NullScattering skyscapes.physical_model.exojax.PowerLawTPProfile skyscapes.physical_model.exojax.RayleighScattering skyscapes.physical_model.exojax.StratosphericPeakMmr skyscapes.physical_model.exojax.TroposphericMmr skyscapes.physical_model.exojax.WavelengthDependentSurface skyscapes.physical_model.exojax.O3ChappuisOpacity skyscapes.physical_model.exojax.ExoJaxPhysicalModel skyscapes.physical_model.exojax.PsgCrossSectionOpacity Functions --------- .. autoapisummary:: skyscapes.physical_model.exojax.build_bulk_prebuilt skyscapes.physical_model.exojax.build_mie_cloud skyscapes.physical_model.exojax.build_species_prebuilt skyscapes.physical_model.exojax.build_exojax_engines skyscapes.physical_model.exojax.vmr_dict_to_earth_profile_dict skyscapes.physical_model.exojax.vmr_dict_to_log_mmr_dict skyscapes.physical_model.exojax.vmr_dict_to_mmr_dict Package Contents ---------------- .. py:data:: MOLECULE_RECIPES :type: dict[str, MoleculeRecipe] .. py:class:: Absorption Bases: :py:obj:`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). .. py:method:: compute(species, Tarr, pressure, gravity, rt_engine) Sum per-species absorption optical depth. .. py:class:: AbstractAbsorption Bases: :py:obj:`equinox.Module` Per-layer absorption opacity from line lists / cross-sections. Concrete implementations iterate over the atmosphere's :class:`MolecularSpecies` tuple, calling each species' ``opa`` engine and summing the contributions. .. py:method:: compute(species, Tarr, pressure, gravity, rt_engine) :abstractmethod: 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. .. py:class:: AbstractClouds Bases: :py:obj:`equinox.Module` Per-layer cloud opacity (mixed scattering + absorption). .. py:method:: compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu) :abstractmethod: 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``. .. py:class:: AbstractMmrProfile Bases: :py:obj:`equinox.Module` Per-species mass-mixing-ratio profile evaluated at layer pressures. .. py:method:: evaluate(pressure) :abstractmethod: Return shape ``(n_layers,)`` mmr values at the layer pressures. .. py:class:: AbstractScattering Bases: :py:obj:`equinox.Module` Per-layer scattering opacity (e.g. Rayleigh). .. py:method:: compute(species, bulk, gravity, rt_engine, n_layers, n_nu) :abstractmethod: Scattering contribution from tracked species + optional bulk gas. For a pure scatterer ``dtau_total == dtau_scatter``; ``g_weighted_num`` is zero for isotropic Rayleigh. .. py:class:: AbstractSurface Bases: :py:obj:`equinox.Module` Bottom-of-atmosphere reflectivity. .. py:method:: compute_refl(log_albedo_scalar, n_nu) :abstractmethod: Return surface reflectivity, shape ``(n_nu,)``. .. py:class:: AbstractTPProfile Bases: :py:obj:`equinox.Module` Temperature-pressure profile. .. py:method:: compute_Tarr(rt_engine, T_eq_K_scalar, T_alpha_scalar) :abstractmethod: Return layer temperatures, shape ``(n_layers,)``. .. py:class:: BulkGasRecipe 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. .. py:attribute:: name :type: str .. py:attribute:: molmass :type: float .. py:attribute:: polarizability_override :type: float | None :value: None .. py:class:: BulkGasResidual Bases: :py:obj:`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,)``. .. py:attribute:: name :type: str .. py:attribute:: molmass :type: float .. py:attribute:: rayleigh_xs :type: jaxtyping.Array .. py:class:: ConstantMmr Bases: :py:obj:`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). .. py:attribute:: log_mmr :type: jaxtyping.Array .. py:method:: evaluate(pressure) Return ``10**log_mmr`` broadcast to ``pressure.shape``. .. py:class:: Contribution Bases: :py:obj:`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``. .. py:attribute:: dtau_total :type: jaxtyping.Array .. py:attribute:: dtau_scatter :type: jaxtyping.Array .. py:attribute:: g_weighted_num :type: jaxtyping.Array .. py:class:: FlatSurface Bases: :py:obj:`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,)``. .. py:attribute:: log_albedo :type: jaxtyping.Array .. py:method:: compute_refl(log_albedo_scalar, n_nu) Return scalar albedo broadcast across the wavenumber grid. .. py:class:: GrayCloud Bases: :py:obj:`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]. .. py:attribute:: log_pressure_bar :type: jaxtyping.Array .. py:attribute:: log_opt_depth :type: jaxtyping.Array .. py:attribute:: ssa :type: float .. py:attribute:: g :type: float .. py:attribute:: log_sigma :type: float .. py:method:: compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu) Gray cloud contribution at a single pressure level. .. py:class:: MieCloud Bases: :py:obj:`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 :class:`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]. .. py:attribute:: log_pressure_bar :type: jaxtyping.Array .. py:attribute:: log_opt_depth :type: jaxtyping.Array .. py:attribute:: ssa_grid :type: jaxtyping.Array .. py:attribute:: g_grid :type: jaxtyping.Array .. py:attribute:: condensate :type: str .. py:attribute:: rg_um :type: float .. py:attribute:: sigmag :type: float .. py:attribute:: log_sigma :type: float .. py:method:: compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu) Mie-cloud contribution: gray-depth distribution, spectral ssa/g. .. py:class:: MolecularSpecies Bases: :py:obj:`equinox.Module` One atmospheric molecule with an altitude-resolved mixing ratio. The mixing ratio is encoded as a profile component (e.g. :class:`~skyscapes.physical_model.exojax.components.mmr_profile.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 :class:`~skyscapes.physical_model.exojax.o3_chappuis.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. .. py:attribute:: profile :type: skyscapes.physical_model.exojax.components.mmr_profile.AbstractMmrProfile .. py:attribute:: name :type: str .. py:attribute:: molmass :type: float .. py:attribute:: opa :type: Any .. py:attribute:: rayleigh_xs :type: jaxtyping.Array .. py:class:: MoleculeRecipe 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 (:class:`~skyscapes.physical_model.exojax.psg_xs.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. .. py:attribute:: name :type: str .. py:attribute:: psg_xs_url :type: str | None :value: None .. py:attribute:: psg_xs_molmass :type: float | None :value: None .. py:attribute:: polarizability_override :type: float | None :value: None .. py:class:: NoCloud Bases: :py:obj:`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. .. py:method:: compute(log_pressure_bar_scalar, log_opt_depth_scalar, pressure, n_nu) Return zero-everywhere cloud contribution. ``log_pressure_bar_scalar`` and ``log_opt_depth_scalar`` are accepted for parity with :class:`GrayCloud` but unused. .. py:class:: NullScattering Bases: :py:obj:`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. .. py:method:: compute(species, bulk, gravity, rt_engine, n_layers, n_nu) Return zero-everywhere contribution. .. py:class:: PowerLawTPProfile Bases: :py:obj:`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 :meth:`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,)``. .. py:attribute:: T_eq_K :type: jaxtyping.Array .. py:attribute:: T_alpha :type: jaxtyping.Array .. py:method:: compute_Tarr(rt_engine, T_eq_K_scalar, T_alpha_scalar) ``T(P) = T_eq * P^alpha`` on the rt_engine's layer pressures. .. py:class:: RayleighScattering Bases: :py:obj:`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 :class:`NullScattering`. .. py:method:: compute(species, bulk, gravity, rt_engine, n_layers, n_nu) 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. .. py:class:: StratosphericPeakMmr Bases: :py:obj:`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,)``. .. py:attribute:: log_peak_mmr :type: jaxtyping.Array .. py:attribute:: log_peak_pressure_bar :type: jaxtyping.Array .. py:attribute:: log_sigma_decades :type: jaxtyping.Array .. py:method:: evaluate(pressure) Gaussian peak in log-pressure. .. py:class:: TroposphericMmr Bases: :py:obj:`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. .. py:attribute:: log_mmr_below :type: jaxtyping.Array .. py:attribute:: log_mmr_above :type: jaxtyping.Array .. py:attribute:: log_pressure_step_bar :type: jaxtyping.Array .. py:attribute:: log_transition_width_decades :type: jaxtyping.Array .. py:method:: evaluate(pressure) Sigmoid step in log-pressure. .. py:class:: WavelengthDependentSurface Bases: :py:obj:`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". .. py:attribute:: log_albedo :type: jaxtyping.Array .. py:attribute:: spectrum :type: jaxtyping.Array .. py:method:: compute_refl(log_albedo_scalar, n_nu) Return wavelength-dependent surface reflectivity. ``n_nu`` is accepted for parity with :class:`FlatSurface` but unused -- the wavelength axis comes from ``self.spectrum``. .. py:function:: build_bulk_prebuilt(*, name, nu_grid) Construct (molmass, rayleigh_xs) for the implicit residual gas. .. py:function:: 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) 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 :func:`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 :class:`MieCloud` instance ready to drop into :class:`ExoJaxPhysicalModel`. .. py:function:: build_species_prebuilt(*, name, nu_grid, nu_min, nu_max, databases_dir, crit) 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 :class:`MolecularSpecies` (everything except ``log_mmr``). .. py:data:: O3_MOLMASS :value: 47.9982 .. py:class:: O3ChappuisOpacity(nu_grid, cache_dir = None, xs_table_path = None) Bases: :py:obj:`skyscapes.physical_model.exojax.psg_xs.PsgCrossSectionOpacity` PSG cross-section opacity for O3 (Serdyuchenko et al. 2014). .. py:data:: DEFAULT_MOLECULES :type: tuple[str, Ellipsis] :value: ('H2O', 'CO2', 'CH4', 'O2', 'O3') .. py:class:: ExoJaxPhysicalModel Bases: :py:obj:`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 :class:`MolecularSpecies`. Each species owns its own ``log_mmr`` (per-planet, shape ``(K,)``). The number and identity of species is configurable via :func:`build_exojax_engines`. bulk: Optional :class:`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). .. py:attribute:: log_gravity_cgs :type: jaxtyping.Array .. py:attribute:: species :type: tuple[skyscapes.physical_model.exojax.components.MolecularSpecies, Ellipsis] .. py:attribute:: bulk :type: skyscapes.physical_model.exojax.components.BulkGasResidual | None .. py:attribute:: tp_profile :type: skyscapes.physical_model.exojax.components.AbstractTPProfile .. py:attribute:: absorption :type: skyscapes.physical_model.exojax.components.AbstractAbsorption .. py:attribute:: scattering :type: skyscapes.physical_model.exojax.components.AbstractScattering .. py:attribute:: clouds :type: skyscapes.physical_model.exojax.components.AbstractClouds .. py:attribute:: surface :type: skyscapes.physical_model.exojax.components.AbstractSurface .. py:attribute:: rt_engine :type: Any .. py:attribute:: nu_grid :type: jaxtyping.Array .. py:attribute:: n_nu :type: int .. py:method:: 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) :classmethod: 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 :func:`build_exojax_engines`. wavelength_max_nm: See :func:`build_exojax_engines`. n_wavenumbers: See :func:`build_exojax_engines`. n_layers: See :func:`build_exojax_engines`. pressure_top_bar: See :func:`build_exojax_engines`. pressure_btm_bar: See :func:`build_exojax_engines`. databases_dir: See :func:`build_exojax_engines`. crit: See :func:`build_exojax_engines`. .. py:method:: 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) :classmethod: Cached one-shot factory: returns a fast :class:`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 :class:`ExoJaxPhysicalModel` via :meth:`from_default_setup`, runs the 2-stream RT once, precomputes the reflectivity, saves to disk, and returns the :class:`PrecomputedPhysicalModel`. Use this when the physical-model parameters are fixed across many evaluations (coronagraphoto sims, ETC studies). For HMC retrievals where parameters vary, use :meth:`from_default_setup` directly. Args: cache_dir: Override the cache directory. Defaults to ``~/.cache/skyscapes/physical_models/``. log_mmrs: See :meth:`from_default_setup`. T_eq_K: See :meth:`from_default_setup`. T_alpha: See :meth:`from_default_setup`. log_surface_albedo: See :meth:`from_default_setup`. log_gravity_cgs: See :meth:`from_default_setup`. log_cloud_pressure_bar: See :meth:`from_default_setup`. log_cloud_opt_depth: See :meth:`from_default_setup`. surface_albedo_spectrum: See :meth:`from_default_setup`. molecules: See :meth:`from_default_setup`. bulk_gas: See :meth:`from_default_setup`. wavelength_min_nm: See :meth:`from_default_setup`. wavelength_max_nm: See :meth:`from_default_setup`. n_wavenumbers: See :meth:`from_default_setup`. n_layers: See :meth:`from_default_setup`. pressure_top_bar: See :meth:`from_default_setup`. pressure_btm_bar: See :meth:`from_default_setup`. databases_dir: See :meth:`from_default_setup`. crit: See :meth:`from_default_setup`. Returns: A :class:`PrecomputedPhysicalModel` with no RT cost at evaluation time. .. py:method:: _reflectivity_one_planet(k) 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). .. py:method:: _reflectivity_all_planets() 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. .. py:method:: contrast(phase_angle_rad, dist_AU, wavelength_nm, Rp_Rearth) 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 :class:`LambertianPhysicalModel`. .. py:method:: contrast_cube(phase_angle_rad, dist_AU, wavelengths_nm, Rp_Rearth) 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 :meth:`contrast`. .. py:method:: __repr__() Human-readable summary of components + species + per-planet state. .. py:function:: 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) 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 :data:`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 :class:`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``). .. py:data:: EARTH_ARCHEAN_VMRS :type: dict[str, float] .. py:data:: EARTH_EARLY_PROTEROZOIC_VMRS :type: dict[str, float] .. py:data:: EARTH_LATE_PROTEROZOIC_VMRS :type: dict[str, float] .. py:data:: EARTH_MODERN_VMRS :type: dict[str, float] .. py:function:: vmr_dict_to_earth_profile_dict(vmrs, K = 1, bulk_gas = 'N2') VMRs -> dict of Earth-realistic profiles per species. Most species get a :class:`ConstantMmr` (well-mixed assumption). For H2O and O3, where Earth's altitude profile differs sharply from uniform, the function builds: - **H2O**: :class:`TroposphericMmr` with the supplied VMR below the cold trap and a 3-decade drop above (parametrized at ``P = 0.1 bar``). - **O3**: :class:`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 :meth:`ExoJaxPhysicalModel.from_default_setup`. .. py:function:: vmr_dict_to_log_mmr_dict(vmrs, K = 1, bulk_gas = 'N2') 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 :meth:`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 :class:`ConstantMmr` profile by default (uniform mmr at every altitude). For Earth-realistic profiles (O3 stratospheric peak, H2O cold-trap drop) use :func:`vmr_dict_to_earth_profile_dict` instead. .. py:function:: vmr_dict_to_mmr_dict(vmrs, bulk_gas = 'N2') 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 :func:`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 :data:`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. .. py:class:: PsgCrossSectionOpacity(name, url, molmass, nu_grid, cache_dir = None, xs_table_path = None, cache_filename = None) 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. .. py:attribute:: name .. py:attribute:: molmass .. py:attribute:: nu_grid .. py:attribute:: _T_grid .. py:attribute:: _sigma_at_nu .. py:method:: xsmatrix(Tarr, pressure) 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.