"""Transit forward model."""
import typing as t
import numpy as np
import numpy.typing as npt
from taurex.chemistry import Chemistry
from taurex.planet import Planet
from taurex.pressure import PressureProfile
from taurex.stellar import Star
from taurex.temperature import TemperatureProfile
from taurex.types import get_float_dtype
from .simplemodel import OneDForwardModel
if t.TYPE_CHECKING:
from taurex.contributions import Contribution
else:
Contribution = object
[docs]
class TransmissionModel(OneDForwardModel):
"""A forward model for transits."""
def __init__(
self,
planet: t.Optional[Planet] = None,
star: t.Optional[Star] = None,
pressure_profile: t.Optional[PressureProfile] = None,
temperature_profile: t.Optional[TemperatureProfile] = None,
chemistry: t.Optional[Chemistry] = None,
nlayers: t.Optional[int] = 100,
atm_min_pressure: t.Optional[float] = 1e-4,
atm_max_pressure: t.Optional[float] = 1e6,
contributions: t.Optional[t.List[Contribution]] = None,
new_path_method: t.Optional[bool] = False,
) -> None:
"""Initialize transit forward model.
Parameters
----------
name: str
Name to use in logging
planet:
Planet model, default planet is Jupiter
star:
Star model, default star is Sun-like
pressure_profile:
Pressure model, alternative is to set ``nlayers``, ``atm_min_pressure``
and ``atm_max_pressure``
temperature_profile:
Temperature model, default is an
:class:`~taurex.data.profiles.temperature.isothermal.Isothermal`
profile at 1500 K
chemistry:
Chemistry model, default is
:class:`~taurex.data.profiles.chemistry.taurexchemistry.TaurexChemistry`
with ``H2O`` and ``CH4``
nlayers: int, optional
Number of layers. Used if ``pressure_profile`` is not defined.
atm_min_pressure: float, optional
Pressure at TOA. Used if ``pressure_profile`` is not defined.
atm_max_pressure: float, optional
Pressure at BOA. Used if ``pressure_profile`` is not defined.
contributions: list, optional
List of contributions to include
new_path_method: bool, optional
Use new path length computation method
"""
super().__init__(
self.__class__.__name__,
planet,
star,
pressure_profile,
temperature_profile,
chemistry,
nlayers,
atm_min_pressure,
atm_max_pressure,
contributions,
)
self.new_method = new_path_method
[docs]
def compute_path_length_old(
self, dz: npt.NDArray[np.float64]
) -> t.List[npt.NDArray[np.float64]]:
"""Compute path length for each layer."""
dl = []
planet_radius = self._planet.fullRadius
total_layers = self.nLayers
z = self.altitudeProfile
self.debug("Computing path_length: \n z=%s \n dz=%s", z, dz)
# Pre-allocate max-size k array, reuse with views to avoid
# repeated allocation in the loop.
_k_buf = np.empty(total_layers, dtype=get_float_dtype())
for layer in range(0, total_layers):
p = (planet_radius + dz[0] / 2 + z[layer]) ** 2
n_remaining = total_layers - layer
k = _k_buf[:n_remaining]
k[0] = np.sqrt(
(planet_radius + dz[0] / 2.0 + z[layer] + dz[layer] / 2.0) ** 2 - p
)
k[1:] = np.sqrt(
(planet_radius + dz[0] / 2 + z[layer + 1 :] + dz[layer + 1 :] / 2) ** 2
- p
)
k[1:] -= np.sqrt(
(
planet_radius
+ dz[0] / 2
+ z[layer : self.nLayers - 1]
+ dz[layer : self.nLayers - 1] / 2
)
** 2
- p
)
dl.append(k * 2.0)
return dl
[docs]
def compute_path_length(self) -> t.List[npt.NDArray[np.float64]]:
"""Compute path length for each layer, new method."""
from taurex.util.geometry import parallel_vector
altitude_boundaries = self.altitude_boundaries
radius = self.planet.fullRadius
# Generate our line of sight paths
viewer, tangent = parallel_vector(
radius, self.altitude_profile + self.deltaz / 2, altitude_boundaries.max()
)
path_lengths = self.planet.compute_path_length(
altitude_boundaries, viewer, tangent
)
return [l for _, l in path_lengths]
[docs]
def path_integral(
self, wngrid: npt.NDArray[np.float64], return_contrib: t.Optional[bool] = False
) -> t.Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
"""Compute path integral.
Calculates the absorption and optical depth for each layer assuming
hemispherical geometry.
"""
dz = self.deltaz
total_layers = self.nLayers
wngrid_size = wngrid.shape[0]
density_profile = self.densityProfile
if self.new_method:
path_length = self.compute_path_length()
else:
path_length = self.compute_path_length_old(dz)
self.path_length = path_length
tau_dtype = get_float_dtype()
tau = np.zeros(shape=(total_layers, wngrid_size), dtype=tau_dtype)
# One scratch buffer for the cross-sections, reused by every
# contribution in turn. Each contribution would otherwise allocate
# its own (nlayers, ngrid) array on every evaluation, which costs a
# full pass of allocation and zeroing per contribution and leaves
# the freed pages behind in the allocator.
sigma_scratch = None
# Memory-efficient: prepare each contribution just before use,
# then clean up its sigma_xsec immediately after all layers.
for contrib in self.contribution_list:
if contrib.sigma_xsec is None:
if sigma_scratch is None:
sigma_scratch = np.zeros(
shape=(total_layers, wngrid_size), dtype=tau_dtype
)
else:
sigma_scratch.fill(0.0)
contrib.prepare(self, wngrid, _out=sigma_scratch)
for layer in range(total_layers):
self.debug("Computing layer %s", layer)
dl = path_length[layer]
end_k = total_layers - layer
if tau[layer].min() > 10:
continue
self.debug("Adding contribution from %s", contrib.name)
contrib.contribute(
self, 0, end_k, layer, layer, density_profile, tau, path_length=dl
)
# Free the large sigma_xsec array for this contribution
del contrib.sigma_xsec
contrib.sigma_xsec = None
self.debug("tau %s %s", tau, tau.shape)
absorption, tau = self.compute_absorption(tau, dz)
return absorption, tau
[docs]
def compute_absorption(
self, tau: npt.NDArray[np.float64], dz: npt.NDArray[np.float64]
) -> t.Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
"""Compute final absorption and optical depth."""
# In-place exp to avoid temporary array allocation
np.exp(-tau, out=tau)
pradius = self._planet.fullRadius
sradius = self._star.radius
# Per-layer integration weights. Keeping these 1D lets the integral
# below avoid the full (nlayers, ngrid) temporaries that broadcasting
# ``(pradius + ap) * (1 - tau) * _dz`` would allocate.
weights = (pradius + self.altitudeProfile) * dz * 2.0
# sum_i w_i * (1 - tau_ij) == sum_i w_i - sum_i w_i * tau_ij
integral = weights.sum() - np.einsum("i,ij->j", weights, tau)
return ((pradius**2.0) + integral) / (sradius**2), tau