skyscapes.physical_model.exojax.components ========================================== .. py:module:: skyscapes.physical_model.exojax.components .. autoapi-nested-parse:: Composable physics components for :class:`ExoJaxPhysicalModel`. Each component encapsulates one piece of atmospheric physics (temperature-pressure profile, line absorption, Rayleigh scattering, clouds, surface) and exposes a uniform contract: - ``compute_*(per_planet_args, ...) -> Contribution`` where :class:`Contribution` carries the per-layer optical-property arrays (``dtau_total``, ``dtau_scatter``, ``g_weighted_num``) that the atmosphere combines before passing to the 2-stream RT solver. Surface and T-P components have analogous compute methods that return their respective quantities. To swap physics: replace one component instance with another that implements the same contract (e.g. ``GrayCloud`` -> ``NoCloud``, ``RayleighScattering`` -> ``NullScattering``). Per-planet PyTree leaves live on the components themselves, so HMC over their parameters "just works" with ``eqx.filter_vmap`` or ``jax.vmap``. Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/skyscapes/physical_model/exojax/components/absorption/index /autoapi/skyscapes/physical_model/exojax/components/base/index /autoapi/skyscapes/physical_model/exojax/components/clouds/index /autoapi/skyscapes/physical_model/exojax/components/mie_cloud/index /autoapi/skyscapes/physical_model/exojax/components/mmr_profile/index /autoapi/skyscapes/physical_model/exojax/components/scattering/index /autoapi/skyscapes/physical_model/exojax/components/species/index /autoapi/skyscapes/physical_model/exojax/components/surface/index /autoapi/skyscapes/physical_model/exojax/components/tp/index Attributes ---------- .. autoapisummary:: skyscapes.physical_model.exojax.components.BULK_GAS_RECIPES skyscapes.physical_model.exojax.components.MOLECULE_RECIPES Classes ------- .. autoapisummary:: skyscapes.physical_model.exojax.components.Absorption skyscapes.physical_model.exojax.components.AbstractAbsorption skyscapes.physical_model.exojax.components.AbstractClouds skyscapes.physical_model.exojax.components.AbstractScattering skyscapes.physical_model.exojax.components.AbstractSurface skyscapes.physical_model.exojax.components.AbstractTPProfile skyscapes.physical_model.exojax.components.Contribution skyscapes.physical_model.exojax.components.GrayCloud skyscapes.physical_model.exojax.components.NoCloud skyscapes.physical_model.exojax.components.MieCloud skyscapes.physical_model.exojax.components.AbstractMmrProfile skyscapes.physical_model.exojax.components.ConstantMmr skyscapes.physical_model.exojax.components.StratosphericPeakMmr skyscapes.physical_model.exojax.components.TroposphericMmr skyscapes.physical_model.exojax.components.NullScattering skyscapes.physical_model.exojax.components.RayleighScattering skyscapes.physical_model.exojax.components.BulkGasRecipe skyscapes.physical_model.exojax.components.BulkGasResidual skyscapes.physical_model.exojax.components.MolecularSpecies skyscapes.physical_model.exojax.components.MoleculeRecipe skyscapes.physical_model.exojax.components.FlatSurface skyscapes.physical_model.exojax.components.WavelengthDependentSurface skyscapes.physical_model.exojax.components.PowerLawTPProfile Functions --------- .. autoapisummary:: skyscapes.physical_model.exojax.components.build_mie_cloud skyscapes.physical_model.exojax.components.build_bulk_prebuilt skyscapes.physical_model.exojax.components.build_species_prebuilt Package Contents ---------------- .. 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:: 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:: 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:: 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:: 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:: 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: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: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:: 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:: 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:: 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:: 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:data:: BULK_GAS_RECIPES :type: dict[str, BulkGasRecipe] .. py:data:: MOLECULE_RECIPES :type: dict[str, MoleculeRecipe] .. 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:: 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:function:: build_bulk_prebuilt(*, name, nu_grid) Construct (molmass, rayleigh_xs) for the implicit residual gas. .. 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: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:: 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: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.