skyscapes.physical_model ======================== .. py:module:: skyscapes.physical_model .. autoapi-nested-parse:: Physical-model hierarchy: phase/wavelength-dependent planet-to-star contrast. Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/skyscapes/physical_model/base/index /autoapi/skyscapes/physical_model/cached/index /autoapi/skyscapes/physical_model/exojax/index /autoapi/skyscapes/physical_model/grid/index /autoapi/skyscapes/physical_model/lambertian/index Classes ------- .. autoapisummary:: skyscapes.physical_model.AbstractPhysicalModel skyscapes.physical_model.PrecomputedPhysicalModel skyscapes.physical_model.ExoJaxPhysicalModel skyscapes.physical_model.GridPhysicalModel skyscapes.physical_model.LambertianPhysicalModel Package Contents ---------------- .. py:class:: AbstractPhysicalModel Bases: :py:obj:`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=...)``. .. py:method:: contrast(phase_angle_rad, dist_AU, wavelength_nm, Rp_Rearth) :abstractmethod: 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)``. .. py:class:: PrecomputedPhysicalModel Bases: :py:obj:`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). .. py:attribute:: reflectivity :type: jaxtyping.Array .. py:attribute:: nu_grid :type: jaxtyping.Array .. py:attribute:: n_nu :type: int .. py:method:: from_physical_model(model) :classmethod: 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 :class:`ExoJaxPhysicalModel` is the supported source. .. py:method:: load(path) :classmethod: Load a previously-saved cache file. Args: path: Path to a ``.npz`` file produced by :meth:`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). .. py:method:: save(path) 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). .. 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, 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 :class:`ExoJaxPhysicalModel.contrast`. .. py:method:: contrast_cube(phase_angle_rad, dist_AU, wavelengths_nm, Rp_Rearth) 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 :meth:`contrast`. .. py:method:: __repr__() Compact summary of the cached spectrum. .. 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:class:: GridPhysicalModel Bases: :py:obj:`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)``. .. py:attribute:: wavelengths_nm :type: jaxtyping.Array .. py:attribute:: phase_angle_deg :type: jaxtyping.Array .. py:attribute:: contrast_grid :type: jaxtyping.Array .. py:method:: contrast(phase_angle_rad, dist_AU, wavelength_nm, Rp_Rearth) 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)``. .. py:method:: __repr__() One-line summary of grid shape (K, n_wl, n_phase) and wl range. .. py:class:: LambertianPhysicalModel Bases: :py:obj:`skyscapes.physical_model.base.AbstractPhysicalModel` Lambertian reflector. Attributes: Ag: Geometric albedo, shape ``(K,)``. .. py:attribute:: Ag :type: jaxtyping.Array .. py:method:: contrast(phase_angle_rad, dist_AU, wavelength_nm, Rp_Rearth) 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``. .. py:method:: __repr__() One-line summary of geometric albedo.