"""
Utility functions for atmospheric retrieval with ``petitRADTRANS``.
This module was put together many contributions by Paul Mollière
(MPIA).
"""
import copy
import warnings
from numbers import Real
import matplotlib.pyplot as plt
import numpy as np
from beartype import beartype
from beartype.typing import Dict, List, Optional, Tuple, Union
from scipy.interpolate import interp1d, PchipInterpolator
from scipy.ndimage import gaussian_filter
from species.core import constants
[docs]
@beartype
def get_line_species() -> List[str]:
"""
Function to get the list of the molecular and atomic line species.
This function is not used anywhere so could be removed.
Returns
-------
list(str)
List with the line species.
"""
return [
"CH4",
"CO",
"CO_all_iso",
"CO_all_iso_HITEMP",
"CO_all_iso_Chubb",
"CO2",
"H2O",
"H2O_HITEMP",
"H2S",
"HCN",
"K",
"K_lor_cut",
"K_allard",
"K_burrows",
"NH3",
"Na",
"Na_lor_cut",
"Na_allard",
"Na_burrows",
"OH",
"PH3",
"TiO",
"TiO_all_Exomol",
"TiO_all_Plez",
"VO",
"VO_Plez",
"FeH",
"H2O_main_iso",
"CH4_main_iso",
]
[docs]
@beartype
def pt_ret_model(
temp_3: Optional[np.ndarray],
delta: Real,
alpha: Real,
tint: Real,
press: np.ndarray,
metallicity: Real,
c_o_ratio: Real,
conv: bool = True,
eq_chem: Optional = None,
) -> Tuple[Optional[np.ndarray], Optional[Real], Optional[Real]]:
"""
Pressure-temperature profile for a self-luminous atmosphere (see
Mollière et al. 2020).
Parameters
----------
temp_3 : np.ndarray, None
Array with three temperature points that are added on top of
the radiative Eddington structure (i.e. above tau = 0.1). The
temperature nodes are connected with a spline interpolation
and a prior is used such that t1 < t2 < t3 < t_connect. The
three temperature points are not used if set to ``None``.
delta : float
Proportionality factor in tau = delta * press_cgs**alpha.
alpha : float
Power law index in
:math:`\\tau = \\delta * P_\\mathrm{cgs}**\\alpha`.
For the tau model: use the proximity to the
:math:`\\kappa_\\mathrm{rosseland}` photosphere as prior.
tint : float
Internal temperature for the Eddington model.
press : np.ndarray
Pressure profile (bar).
metallicity : float
Metallicity [Fe/H]. Required for the ``nabla_ad``
interpolation.
c_o_ratio : float
Carbon-to-oxygen ratio. Required for the ``nabla_ad``
interpolation.
conv : bool
Enforce a convective adiabat.
eq_chem : PreCalculatedEquilibriumChemistryTable, None
Instance of the equilibrium chemistry table from
``petitRADTRANS``. The instance is created if the
arguments is set to ``None``.
Returns
-------
np.ndarray
Temperature profile (K) for ``press``.
float
Pressure (bar) where the optical depth is 1.
float, None
Pressure (bar) at the radiative-convective boundary.
"""
# Convert pressures from bar to cgs units
press_cgs = press * 1e6
# Calculate the optical depth
tau = delta * press_cgs**alpha
# Calculate the Eddington temperature
tedd = (3.0 / 4.0 * tint**4.0 * (2.0 / 3.0 + tau)) ** 0.25
# Interpolate the abundances, following chemical equilibrium
if eq_chem is None:
from petitRADTRANS.chemistry.pre_calculated_chemistry import (
PreCalculatedEquilibriumChemistryTable,
)
eq_chem = PreCalculatedEquilibriumChemistryTable()
_, _, nabla_ad = eq_chem.interpolate_mass_fractions(
np.full(tedd.size, c_o_ratio),
np.full(tedd.size, metallicity),
tedd,
press,
carbon_pressure_quench=None,
full=True,
)
# Enforce convective adiabat
if conv:
# Calculate the current, radiative temperature gradient
nab_rad = np.diff(np.log(tedd)) / np.diff(np.log(press_cgs))
# Extend to array of same length as pressure structure
nabla_rad = np.ones_like(tedd)
nabla_rad[0] = nab_rad[0]
nabla_rad[-1] = nab_rad[-1]
nabla_rad[1:-1] = (nab_rad[1:] + nab_rad[:-1]) / 2.0
# Where is the atmosphere convectively unstable?
conv_index = nabla_rad > nabla_ad
if np.argwhere(conv_index).size == 0:
conv_press = None
else:
conv_bound = np.amin(np.argwhere(conv_index))
conv_press = press[conv_bound]
tfinal = None
for i in range(10):
if i == 0:
t_take = copy.copy(tedd)
else:
t_take = copy.copy(tfinal)
_, _, nabla_ad = eq_chem.interpolate_mass_fractions(
np.full(t_take.size, c_o_ratio),
np.full(t_take.size, metallicity),
t_take,
press,
carbon_pressure_quench=None,
full=True,
)
# Calculate the average nabla_ad between the layers
nabla_ad_mean = nabla_ad
nabla_ad_mean[1:] = (nabla_ad[1:] + nabla_ad[:-1]) / 2.0
# What are the increments in temperature due to convection
tnew = nabla_ad_mean[conv_index] * np.mean(np.diff(np.log(press_cgs)))
# What is the last radiative temperature?
tstart = np.log(t_take[~conv_index][-1])
# Integrate and translate to temperature
# from log(temperature)
tnew = np.exp(np.cumsum(tnew) + tstart)
# Add upper radiative and lower covective
# part into one single array
tfinal = copy.copy(t_take)
tfinal[conv_index] = tnew
if np.max(np.abs(t_take - tfinal) / t_take) < 0.01:
# print('n_ad', 1./(1.-nabla_ad[conv_index]))
break
else:
tfinal = tedd
conv_press = None
# Add the three temperature-point P-T description above tau = 0.1
@beartype
def press_tau(tau: Real) -> Real:
"""
Function to return the pressure in cgs units at a given
optical depth.
Parameters
----------
tau : float
Optical depth.
Returns
-------
float
Pressure (cgs) at optical depth ``tau``.
"""
return (tau / delta) ** (1.0 / alpha)
# Where is the uppermost pressure of the
# Eddington radiative structure?
p_bot_spline = press_tau(0.1)
if temp_3 is None:
tret = tfinal
else:
for i_intp in range(2):
if i_intp == 0:
# Create the pressure coordinates for the spline
# support nodes at low pressure
support_points_low = np.logspace(
np.log10(press_cgs[0]), np.log10(p_bot_spline), 4
)
# Create the pressure coordinates for the spline
# support nodes at high pressure, the corresponding
# temperatures for these nodes will be taken from
# the radiative-convective solution
support_points_high = 10.0 ** np.arange(
np.log10(p_bot_spline),
np.log10(press_cgs[-1]),
np.diff(np.log10(support_points_low))[0],
)
# Combine into one support node array, don't add
# the p_bot_spline point twice.
support_points = np.zeros(
len(support_points_low) + len(support_points_high) - 1
)
support_points[:4] = support_points_low
support_points[4:] = support_points_high[1:]
else:
# Create the pressure coordinates for the spline
# support nodes at low pressure
support_points_low = np.logspace(
np.log10(press_cgs[0]), np.log10(p_bot_spline), 7
)
# Create the pressure coordinates for the spline
# support nodes at high pressure, the corresponding
# temperatures for these nodes will be taken from
# the radiative-convective solution
support_points_high = np.logspace(
np.log10(p_bot_spline), np.log10(press_cgs[-1]), 7
)
# Combine into one support node array, don't add
# the p_bot_spline point twice.
support_points = np.zeros(
len(support_points_low) + len(support_points_high) - 1
)
support_points[:7] = support_points_low
support_points[7:] = support_points_high[1:]
# Define the temperature values at the node points
t_support = np.zeros_like(support_points)
if i_intp == 0:
tfintp = interp1d(press_cgs, tfinal)
# The temperature at p_bot_spline (from the
# radiative-convective solution)
t_support[len(support_points_low) - 1] = tfintp(p_bot_spline)
# if temp_3 is not None:
# The temperature at pressures below
# p_bot_spline (free parameters)
t_support[: len(support_points_low) - 1] = temp_3
# else:
# t_support[:3] = tfintp(support_points_low[:3])
# The temperature at pressures above p_bot_spline
# (from the radiative-convective solution)
t_support[len(support_points_low) :] = tfintp(
support_points[len(support_points_low) :]
)
else:
tfintp1 = interp1d(press_cgs, tret)
t_support[: len(support_points_low) - 1] = tfintp1(
support_points[: len(support_points_low) - 1]
)
tfintp = interp1d(press_cgs, tfinal)
# The temperature at p_bot_spline (from
# the radiative-convective solution)
t_support[len(support_points_low) - 1] = tfintp(p_bot_spline)
try:
t_support[len(support_points_low) :] = tfintp(
support_points[len(support_points_low) :]
)
except ValueError:
return None, None, None
# Make the temperature spline interpolation to be returned
# to the user tret = spline(np.log10(support_points),
# t_support, np.log10(press_cgs), order = 3)
cs = PchipInterpolator(np.log10(support_points), t_support)
tret = cs(np.log10(press_cgs))
# Return the temperature, the pressure at tau = 1
# The temperature at the connection point: tfintp(p_bot_spline)
# The last two are needed for the priors on the P-T profile.
return tret, press_tau(1.0) / 1e6, conv_press
[docs]
@beartype
def pt_spline_interp(
knot_press: np.ndarray,
knot_temp: np.ndarray,
pressure: np.ndarray,
pt_smooth: Optional[Real] = 0.3,
) -> np.ndarray:
"""
Function for interpolating the P-T nodes with a PCHIP 1-D monotonic
cubic interpolation. The interpolated temperature is smoothed with
a Gaussian kernel of width 0.3 dex in pressure (see Piette &
Madhusudhan 2020).
Parameters
----------
knot_press : np.ndarray
Pressure knots (bar).
knot_temp : np.ndarray
Temperature knots (K).
pressure : np.ndarray
Pressure points (bar) at which the temperatures is
interpolated.
pt_smooth : float, dict, None
Standard deviation of the Gaussian kernel that is used for
smoothing the P-T profile, after the temperature nodes
have been interpolated to a higher pressure resolution.
The argument should be given as
:math:`\\log10{P/\\mathrm{bar}}`, with the default value
set to 0.3 dex. No smoothing is applied if the argument
is set to ``None``.
Returns
-------
np.ndarray
Interpolated, smoothed temperature points (K).
"""
pt_interp = PchipInterpolator(np.log10(knot_press), knot_temp)
temp_interp = pt_interp(np.log10(pressure))
log_press = np.log10(pressure)
log_diff = np.mean(np.diff(log_press))
if np.std(np.diff(log_press)) / log_diff > 1e-6:
raise ValueError("Expecting equally spaced pressures in log space.")
if pt_smooth is not None:
temp_interp = gaussian_filter(
temp_interp, sigma=pt_smooth / log_diff, mode="nearest"
)
return temp_interp
[docs]
@beartype
def create_pt_profile(
cube,
cube_index: Dict[str, Real],
pt_profile: str,
pressure: np.ndarray,
knot_press: Optional[np.ndarray],
metallicity: Real,
c_o_ratio: Real,
pt_smooth: Optional[Union[Real, Dict[str, Real]]] = 0.3,
eq_chem: Optional = None,
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[Real], Optional[Real]]:
"""
Function for creating a pressure-temperature profile.
Parameters
----------
cube : LP_c_double
Unit cube.
cube_index : dict
Dictionary with the index of each parameter in the ``cube``.
pt_profile : str
The parametrization for the pressure-temperature profile
('molliere', 'free', 'monotonic', 'eddington', 'gradient').
pressure : np.ndarray
Pressure points (bar) at which the temperatures is
interpolated.
knot_press : np.ndarray, None
Pressure knots (bar), which are required when the argument of
``pt_profile`` is either 'free' or 'monotonic'.
metallicity : float
Metallicity [Fe/H].
c_o_ratio : float, None
Carbon-to-oxgen ratio.
pt_smooth : float, dict, None
Standard deviation of the Gaussian kernel that is used for
smoothing the P-T profile, after the temperature nodes
have been interpolated to a higher pressure resolution.
The argument should be given as
:math:`\\log10{P/\\mathrm{bar}}`, with the default value
set to 0.3 dex. No smoothing is applied if the argument
is set to ``None``.
eq_chem : PreCalculatedEquilibriumChemistryTable, None
Instance of the equilibrium chemistry table from
``petitRADTRANS``. Only required when fitting abundances
following equilibrium chemistry and when ``pt_profile``
is either 'molliere' or 'mod-molliere'.
Returns
-------
np.ndarray
Temperatures (K).
np.ndarray, None
Temperature at the knots (K). A ``None`` is returned if
``pt_profile`` is set to 'molliere' or 'eddington'.
float, None
Pressure (bar) where the optical depth is 1.
float, None
Pressure (bar) at the radiative-convective boundary.
"""
knot_temp = None
if pt_profile == "molliere":
temp, phot_press, conv_press = pt_ret_model(
np.array(
[cube[cube_index["t1"]], cube[cube_index["t2"]], cube[cube_index["t3"]]]
),
10.0 ** cube[cube_index["log_delta"]],
cube[cube_index["alpha"]],
cube[cube_index["tint"]],
pressure,
metallicity,
c_o_ratio,
eq_chem,
)
elif pt_profile == "mod-molliere":
temp, phot_press, conv_press = pt_ret_model(
None,
10.0 ** cube[cube_index["log_delta"]],
cube[cube_index["alpha"]],
cube[cube_index["tint"]],
pressure,
metallicity,
c_o_ratio,
eq_chem,
)
elif pt_profile == "gradient":
num_layer = 6 # could make a variable in the future
layer_pt_slopes = np.ones(num_layer) * np.nan
for index in range(num_layer):
layer_pt_slopes[index] = cube[cube_index[f"PTslope_{num_layer - index}"]]
try:
from petitRADTRANS.physics import dTdP_temperature_profile
except ImportError as import_error:
raise ImportError(
"Can't import the dTdP profile function from "
"petitRADTRANS, check that the version of pRT "
"includes this function in petitRADTRANS.physics",
import_error,
)
temp = dTdP_temperature_profile(
pressure,
num_layer, # could change in the future
layer_pt_slopes,
cube[cube_index["T_bottom"]],
)
phot_press = None
conv_press = None
elif pt_profile in ["free", "monotonic"]:
knot_temp = []
for i in range(knot_press.size):
knot_temp.append(cube[cube_index[f"t{i}"]])
knot_temp = np.asarray(knot_temp)
nan_count = np.sum(np.isnan(knot_temp))
if nan_count > 0:
# TODO Not clear why this can happen
warnings.warn(f"Found {nan_count} NaN values in sampled temperature nodes.")
return None, None, None, None
temp = pt_spline_interp(knot_press, knot_temp, pressure, pt_smooth)
phot_press = None
conv_press = None
elif pt_profile == "eddington":
# Eddington approximation
# delta = kappa_ir/gravity
tau = pressure * 1e6 * 10.0 ** cube[cube_index["log_delta"]]
temp = (0.75 * cube[cube_index["tint"]] ** 4.0 * (2.0 / 3.0 + tau)) ** 0.25
phot_press = None
conv_press = None
return temp, knot_temp, phot_press, conv_press
[docs]
@beartype
def make_half_pressure_better(
p_base: Dict[str, Real], pressure: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for reducing the number of pressure layers from 1440 to
~100 (depending on the number of cloud species) with a refinement
around the cloud decks.
Parameters
----------
p_base : dict
Dictionary with the base of the cloud deck for all cloud
species. The keys in the dictionary are included for example
as MgSiO3(c).
pressure : np.ndarray
Pressures (bar) at high resolution (1440 points).
Returns
-------
np.ndarray
Pressures (bar) at lower resolution (60 points) but with a
refinement around the position of the cloud decks.
np.ndarray, None
The indices of the pressures that have been selected from
the input array ``pressure``.
"""
press_plus_index = np.zeros(len(pressure) * 2).reshape(len(pressure), 2)
press_plus_index[:, 0] = pressure
press_plus_index[:, 1] = range(len(pressure))
press_small = press_plus_index[::24, :]
press_plus_index = press_plus_index[::2, :]
indexes_small = press_small[:, 0] > 0.0
indexes = press_plus_index[:, 0] > 0.0
for P_cloud in p_base.values():
indexes_small = indexes_small & (
(np.log10(press_small[:, 0] / P_cloud) > 0.05)
| (np.log10(press_small[:, 0] / P_cloud) < -0.3)
)
indexes = indexes & (
(np.log10(press_plus_index[:, 0] / P_cloud) > 0.05)
| (np.log10(press_plus_index[:, 0] / P_cloud) < -0.3)
)
press_cut = press_plus_index[~indexes, :]
press_small_cut = press_small[indexes_small, :]
press_out = np.zeros((len(press_cut) + len(press_small_cut)) * 2).reshape(
(len(press_cut) + len(press_small_cut)), 2
)
press_out[: len(press_small_cut), :] = press_small_cut
press_out[len(press_small_cut) :, :] = press_cut
press_out = np.sort(press_out, axis=0)
return press_out[:, 0], press_out[:, 1].astype("int")
[docs]
@beartype
def create_abund_dict(
abund_in: Dict[str, np.ndarray],
line_species: List[str],
cloud_species: List[str],
pressure_grid: str = "smaller",
indices: Optional[np.ndarray] = None,
eq_chem: Optional = None,
) -> Dict[str, np.ndarray]:
"""
Function to update the names in the abundance dictionary.
Parameters
----------
abund_in : dict
Dictionary with the mass fractions.
line_species : list
List with the line species.
cloud_species : list
List with the cloud species.
pressure_grid : str
The type of pressure grid that is used for the radiative
transfer. Either 'standard', to use 180 layers both for the
atmospheric structure (e.g. when interpolating the abundances)
and 180 layers with the radiative transfer, or 'smaller' to
use 60 (instead of 180) with the radiative transfer, or 'clouds'
to start with 1440 layers but resample to ~100 layers (depending
on the number of cloud species) with a refinement around the
cloud decks. For cloudless atmospheres it is recommended to use
'smaller', which runs faster than 'standard' and provides
sufficient accuracy. For cloudy atmosphere, one can test with
'smaller' but it is recommended to use 'clouds' for improved
accuracy fluxes.
indices : np.ndarray, None
Pressure indices from the adaptive refinement in a cloudy
atmosphere. Only required with ``pressure_grid='clouds'``.
Otherwise, the argument can be set to ``None``.
eq_chem : PreCalculatedEquilibriumChemistryTable, None
Instance of the equilibrium chemistry table from
``petitRADTRANS``. This is simply to check if the
abundances are from the chemical equilibrium table.
Returns
-------
dict
Dictionary with the updated names of the abundances.
"""
from petitRADTRANS.chemistry.utils import simplify_species_list
line_simple = simplify_species_list(line_species)
# Create a dictionary with the mass fractions
abund_out = {}
if indices is not None:
# Line species
for i, item in enumerate(line_species):
if eq_chem is not None:
abund_out[item] = abund_in[line_simple[i]][indices]
else:
abund_out[item] = abund_in[item][indices]
# Cloud species
for i, item in enumerate(cloud_species):
abund_out[item] = abund_in[item][indices]
abund_out["H2"] = abund_in["H2"][indices]
abund_out["He"] = abund_in["He"][indices]
elif pressure_grid == "smaller":
# Line species
for i, item in enumerate(line_species):
if eq_chem is not None:
abund_out[item] = abund_in[line_simple[i]][::3]
else:
abund_out[item] = abund_in[item][::3]
# Cloud species
for i, item in enumerate(cloud_species):
abund_out[item] = abund_in[item][::3]
abund_out["H2"] = abund_in["H2"][::3]
abund_out["He"] = abund_in["He"][::3]
else:
# Line species
for i, item in enumerate(line_species):
if eq_chem is not None:
abund_out[item] = abund_in[line_simple[i]][:]
else:
abund_out[item] = abund_in[item][:]
# Cloud species
for i, item in enumerate(cloud_species):
abund_out[item] = abund_in[item][:]
abund_out["H2"] = abund_in["H2"][:]
abund_out["He"] = abund_in["He"][:]
# Correction for the nuclear spin degeneracy that was not included
# in the partition function. See Charnay et al. (2018)
if "FeH" in abund_out:
abund_out["FeH"] = abund_out["FeH"] / 2.0
return abund_out
[docs]
@beartype
def calc_spectrum_clear(
rt_object,
pressure: np.ndarray,
temperature: np.ndarray,
log_g: Real,
c_o_ratio: Optional[Real],
metallicity: Optional[Real],
p_quench: Optional[Real],
log_x_abund: Optional[dict],
knot_press_abund: Optional[np.ndarray],
abund_smooth: Optional[Real],
pressure_grid: str = "smaller",
contribution: bool = False,
return_opacities: bool = False,
eq_chem: Optional = None,
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[Dict]]:
"""
Function to simulate an emission spectrum of a clear atmosphere.
rt_object : petitRADTRANS.radtrans.Radtrans
Instance of ``Radtrans``.
pressure : np.ndarray
Array with the pressure points (bar).
temperature : np.ndarray
Array with the temperature points (K) corresponding to
``pressure``.
log_g : float
Log10 of the surface gravity (cm s-2).
c_o_ratio : float, None
Carbon-to-oxygen ratio.
metallicity : float, None
Metallicity.
p_quench : float, None
Quenching pressure (bar).
log_x_abund : dict, None
Dictionary with the log10 of the abundances. Only required when
``eq_chem`` is not ``None``, so when fitting free abundances.
knot_press_abund : np.ndarray, None
Pressure nodes at which the abundances are sampled. Only
required when ``eq_chem`` is not ``None``, so when fitting
free abundances.
abund_smooth : float, None
Standard deviation of the Gaussian kernel that is used for
smoothing the abundance profiles, after the abundance nodes
have been interpolated to a higher pressure resolution.
Only required with ``eq_chem`` and ``knot_press_abund``
are both not ``None``. The argument should be given as
:math:`\\log10{P/\\mathrm{bar}}`. No smoothing is applied
if the argument if set to 0 or ``None``.
pressure_grid : str
The type of pressure grid that is used for the radiative
transfer. Either 'standard', to use 180 layers both for the
atmospheric structure (e.g. when interpolating the abundances)
and 180 layers with the radiative transfer, or 'smaller' to use
60 (instead of 180) with the radiative transfer, or 'clouds' to
start with 1440 layers but resample to ~100 layers (depending
on the number of cloud species) with a refinement around the
cloud decks. For cloudless atmospheres it is recommended to use
'smaller', which runs faster than 'standard' and provides
sufficient accuracy. For cloudy atmosphere, one can test with
'smaller' but it is recommended to use 'clouds' for improved
accuracy fluxes.
contribution : bool
Calculate the emission contribution.
return_opacities : bool
Return opacities and optical depth.
eq_chem : PreCalculatedEquilibriumChemistryTable, None
Instance of the equilibrium chemistry table from
``petitRADTRANS``. Only required when fitting abundances
following equilibrium chemistry. With free abundances,
the argument should be set to ``None``.
Returns
-------
np.ndarray
Wavelength (um).
np.ndarray
Flux (W m-2 um-1).
np.ndarray, None
Emission contribution.
dict, None
Dictionary with extra output.
"""
if knot_press_abund is None:
abund_nodes = None
else:
abund_nodes = knot_press_abund.size
if eq_chem is not None:
# Equilibrium chemistry
abund_in, mmw, _ = eq_chem.interpolate_mass_fractions(
np.full(pressure.shape, c_o_ratio),
np.full(pressure.shape, metallicity),
temperature,
pressure,
carbon_pressure_quench=p_quench,
full=True,
)
else:
# Free abundances
# Create a dictionary with all mass fractions
abund_in = mass_frac_dict(log_x_abund, rt_object.line_species, abund_nodes)
# Create list of all species
abund_species = rt_object.line_species.copy()
abund_species.append("H2")
abund_species.append("He")
if abund_nodes is None:
# Mean molecular weight
mmw_float = mean_molecular_weight(abund_in)
mmw = np.full(pressure.size, mmw_float)
# Create arrays of constant abundances with pressure
for abund_item in abund_species:
abund_in[abund_item] = np.full(pressure.size, abund_in[abund_item])
else:
# Create arrays of pressure-dependent abundances
for abund_item in abund_species:
knot_abund = np.zeros(knot_press_abund.shape)
for node_idx in range(abund_nodes):
knot_abund[node_idx] = abund_in[f"{abund_item}_{node_idx}"]
del abund_in[f"{abund_item}_{node_idx}"]
nan_count = np.sum(np.isnan(knot_abund))
if nan_count > 0:
warnings.warn(
f"Found {nan_count} NaN values in sampled abundance nodes."
)
return None, None, None
abund_in[abund_item] = pt_spline_interp(
knot_press_abund, knot_abund, pressure, pt_smooth=abund_smooth
)
# Mean molecular weight
mmw = np.zeros(pressure.size)
for press_idx in range(pressure.size):
abund_dict = {}
for abund_item, abund_value in abund_in.items():
abund_dict[abund_item] = abund_value[press_idx]
mmw[press_idx] = mean_molecular_weight(abund_dict)
# Extract every three levels when pressure_grid is set to 'smaller'
if pressure_grid == "smaller":
temperature = temperature[::3]
pressure = pressure[::3]
mmw = mmw[::3]
mass_fractions = create_abund_dict(
abund_in,
rt_object.line_species,
rt_object.cloud_species,
pressure_grid=pressure_grid,
indices=None,
eq_chem=eq_chem,
)
# Calculate the emission spectrum
# wavelength in cm and flux in erg cm-2 s-1 cm-1
rad_wavel, rad_flux, extra_out = rt_object.calculate_flux(
temperatures=temperature,
mass_fractions=mass_fractions,
mean_molar_masses=mmw,
reference_gravity=10.0**log_g,
return_contribution=contribution,
return_opacities=return_opacities,
)
# Return wavelength (um), flux (W m-2 um-1), and extra output
return 1e4 * rad_wavel, 1e-7 * rad_flux, extra_out
[docs]
@beartype
def calc_spectrum_clouds(
rt_object,
pressure: np.ndarray,
temperature: np.ndarray,
c_o_ratio: Real,
metallicity: Real,
p_quench: Optional[Real],
log_x_abund: Optional[dict],
log_x_base: Optional[dict],
cloud_dict: Dict[str, Optional[Real]],
log_g: Real,
knot_press_abund: Optional[np.ndarray],
abund_smooth: Optional[Real],
pressure_grid: str = "smaller",
plotting: bool = False,
contribution: bool = False,
tau_cloud: Optional[Real] = None,
cloud_wavel: Optional[Tuple[Real, Real]] = None,
return_opacities: bool = False,
eq_chem: Optional = None,
) -> Tuple[
Optional[np.ndarray], Optional[np.ndarray], Optional[Dict], Optional[np.ndarray]
]:
"""
Function to simulate an emission spectrum of a cloudy atmosphere.
Parameters
----------
rt_object : petitRADTRANS.radtrans.Radtrans
Instance of ``Radtrans``.
pressure : np.ndarray
Array with the pressure points (bar).
temperature : np.ndarray
Array with the temperature points (K) corresponding to
``pressure``.
c_o_ratio : float
Carbon-to-oxygen ratio.
metallicity : float
Metallicity.
p_quench : float, None
Quenching pressure (bar).
log_x_abund : dict, None
Dictionary with the log10 of the abundances. Only required when
``eq_chem`` is not ``None``, so when fitting free abundances.
log_x_base : dict, None
Dictionary with the log10 of the mass fractions at the cloud
base. Only required when the ``cloud_dict`` contains ``fsed``,
``log_kzz``, and ``sigma_lnorm``.
cloud_dict : dict
Dictionary with the cloud parameters.
log_g : float
Log10 of the surface gravity (cm s-2).
knot_press_abund : np.ndarray, None
Pressure nodes at which the abundances are sampled. Only
required when ``eq_chem`` is not ``None``, so when fitting
free abundances.
abund_smooth : float, None
Standard deviation of the Gaussian kernel that is used for
smoothing the abundance profiles, after the abundance nodes
have been interpolated to a higher pressure resolution.
Only required with ``eq_chem`` and ``knot_press_abund``
are both not ``None``. The argument should be given as
:math:`\\log10{P/\\mathrm{bar}}`. No smoothing is applied
if the argument if set to 0 or ``None``.
pressure_grid : str
The type of pressure grid that is used for the radiative
transfer. Either 'standard', to use 180 layers both for the
atmospheric structure (e.g. when interpolating the abundances)
and 180 layers with the radiative transfer, or 'smaller' to
use 60 (instead of 180) with the radiative transfer, or
'clouds' to start with 1440 layers but resample to ~100 layers
(depending on the number of cloud species) with a refinement
around the cloud decks. For cloudless atmospheres it is
recommended to use 'smaller', which runs faster than 'standard'
and provides sufficient accuracy. For cloudy atmosphere, one
can test with 'smaller' but it is recommended to use 'clouds'
for improved accuracy fluxes.
plotting : bool
Create plots.
contribution : bool
Calculate the emission contribution.
tau_cloud : float, None
Total cloud optical that will be used for scaling the cloud
mass fractions. The mass fractions will not be scaled if the
parameter is set to ``None``.
cloud_wavel : tuple(float, float), None
Tuple with the wavelength range (um) that is used for
calculating the median optical depth of the clouds at the
gas-only photosphere and then scaling the cloud optical
depth to the value of ``log_tau_cloud``. The range of
``cloud_wavel`` should be encompassed by the range of
``wavel_range``. The full wavelength range (i.e.
``wavel_range``) is used if the argument is set to ``None``.
return_opacities : bool
Return opacities and optical depth.
eq_chem : PreCalculatedEquilibriumChemistryTable, None
Instance of the equilibrium chemistry table from
``petitRADTRANS``. Only required when fitting abundances
following equilibrium chemistry. With free abundances,
the argument should be set to ``None``.
Returns
-------
np.ndarray, None
Wavelength (um).
np.ndarray, None
Flux (W m-2 um-1).
dict, None
Dictionary with extra output.
np.ndarray, None
Array with mean molecular weight.
"""
if knot_press_abund is None:
abund_nodes = None
else:
abund_nodes = knot_press_abund.size
if eq_chem is not None:
# Equilibrium chemistry
abund_in, mmw, _ = eq_chem.interpolate_mass_fractions(
np.full(pressure.shape, c_o_ratio),
np.full(pressure.shape, metallicity),
temperature,
pressure,
carbon_pressure_quench=p_quench,
full=True,
)
else:
# Free abundances
# Create a dictionary with all mass fractions
abund_in = mass_frac_dict(log_x_abund, rt_object.line_species, abund_nodes)
# Create list of all species
abund_species = rt_object.line_species.copy()
abund_species.append("H2")
abund_species.append("He")
if abund_nodes is None:
# Mean molecular weight
mmw_float = mean_molecular_weight(abund_in)
mmw = np.full(pressure.size, mmw_float)
# Create arrays of constant abundances with pressure
for abund_item in abund_species:
abund_in[abund_item] = np.full(pressure.size, abund_in[abund_item])
else:
# Create arrays of pressure-dependent abundances
for abund_item in abund_species:
knot_abund = np.zeros(knot_press_abund.shape)
for node_idx in range(abund_nodes):
knot_abund[node_idx] = abund_in[f"{abund_item}_{node_idx}"]
del abund_in[f"{abund_item}_{node_idx}"]
nan_count = np.sum(np.isnan(knot_abund))
if nan_count > 0:
warnings.warn(
f"Found {nan_count} NaN values in sampled abundance nodes."
)
return None, None, None, None
abund_in[abund_item] = pt_spline_interp(
knot_press_abund, knot_abund, pressure, pt_smooth=abund_smooth
)
# Mean molecular weight
mmw = np.zeros(pressure.size)
for press_idx in range(pressure.size):
abund_dict = {}
for abund_item, abund_value in abund_in.items():
abund_dict[abund_item] = abund_value[press_idx]
mmw[press_idx] = mean_molecular_weight(abund_dict)
if log_x_base is not None:
p_base = {}
for cloud_item in log_x_base:
if f"log_p_base_{cloud_item}" in cloud_dict:
p_base_item = 10.0 ** cloud_dict[f"log_p_base_{cloud_item}"]
else:
p_base_item = find_cloud_deck(
cloud_item,
pressure,
temperature,
metallicity,
c_o_ratio,
mmw=np.mean(mmw),
plotting=plotting,
)
p_base[f"{cloud_item}"] = p_base_item
abund_in[f"{cloud_item}"] = np.zeros_like(temperature)
if "fsed" in cloud_dict:
f_sed = cloud_dict["fsed"]
else:
f_sed = cloud_dict[f"fsed_{cloud_item}"]
abund_in[f"{cloud_item}"][pressure < p_base_item] = (
10.0 ** log_x_base[cloud_item]
* (pressure[pressure <= p_base_item] / p_base_item) ** f_sed
)
# Adaptive pressure refinement around the cloud base
if pressure_grid == "clouds":
_, indices = make_half_pressure_better(p_base, pressure)
else:
indices = None
# Create dictionary with mass fractions
mass_fractions = create_abund_dict(
abund_in,
rt_object.line_species,
rt_object.cloud_species,
pressure_grid=pressure_grid,
indices=indices,
eq_chem=eq_chem,
)
# Create dictionary with sedimentation parameters
# Use the same value for all cloud species
fseds = {}
for cloud_item in rt_object.cloud_species:
# The cloud_item has the form of e.g. MgSiO3(c). For
# parametrized cloud opacities, then the number of
# cloud_species is zero so the fseds dictionary is empty
if "fsed" in cloud_dict:
fseds[cloud_item] = cloud_dict["fsed"]
else:
fseds[cloud_item] = cloud_dict[f"fsed_{cloud_item}"]
# Create an array with a constant eddy diffusion coefficient (cm2 s-1)
if "log_kzz" in cloud_dict:
Kzz_use = np.full(pressure.shape, 10.0 ** cloud_dict["log_kzz"])
else:
Kzz_use = None
# Adjust number of atmospheric levels
if pressure_grid == "smaller":
temperature = temperature[::3]
pressure = pressure[::3]
mmw = mmw[::3]
if "log_kzz" in cloud_dict:
Kzz_use = Kzz_use[::3]
elif pressure_grid == "clouds":
temperature = temperature[indices]
pressure = pressure[indices]
mmw = mmw[indices]
if "log_kzz" in cloud_dict:
Kzz_use = Kzz_use[indices]
# Optionally plot the cloud properties
if plotting and Kzz_use is not None:
for key, value in mass_fractions.items():
plt.plot(value, pressure, label=key)
plt.xlim(1e-10, 1.0)
plt.ylim(pressure[-1], pressure[0])
plt.yscale("log")
plt.xscale("log")
plt.xlabel("Mass fraction")
plt.ylabel("Pressure (bar)")
if p_quench is not None:
plt.axhline(p_quench, ls="--", color="black")
plt.legend(loc="best", fontsize=8.0)
plt.savefig("abundances.png", bbox_inches="tight")
plt.clf()
plt.plot(temperature, pressure, "o", ls="none", ms=2.0)
for item in log_x_base:
plt.axhline(
p_base[f"{item}"], label=f"Cloud deck {item}", ls="--", color="black"
)
plt.yscale("log")
plt.ylim(1e3, 1e-6)
plt.xlim(0.0, 4000.0)
plt.savefig("pt_cloud_deck.png", bbox_inches="tight")
plt.clf()
for item in log_x_base:
plt.plot(mass_fractions[f"{item}"], pressure)
plt.axhline(p_base[f"{item}"])
plt.yscale("log")
if np.count_nonzero(mass_fractions[f"{item}"]) > 0:
plt.xscale("log")
plt.ylim(1e3, 1e-6)
plt.xlim(1e-10, 1.0)
log_x_base_item = log_x_base[item]
if "fsed" in cloud_dict:
fsed = cloud_dict["fsed"]
else:
fsed = cloud_dict[f"fsed_{item}"]
log_kzz = cloud_dict["log_kzz"]
plt.title(
f"fsed = {fsed:.2f}, log(Kzz) = {log_kzz:.2f}, "
+ f"X_b = {log_x_base_item:.2f}"
)
plt.savefig(f"{item.lower()}_clouds.png", bbox_inches="tight")
plt.clf()
# Reinitiate the pressure layers after make_half_pressure_better
if pressure_grid == "clouds":
rt_object.setup_opa_structure(pressure)
# Width of cloud particle distribution
if "sigma_lnorm" in cloud_dict:
sigma_lnorm = cloud_dict["sigma_lnorm"]
else:
sigma_lnorm = None
if "log_kappa_0" in cloud_dict:
# Cloud model 2
@beartype
def kappa_abs(wavel_micron: np.ndarray, press_bar: np.ndarray) -> np.ndarray:
p_base = 10.0 ** cloud_dict["log_p_base"] # (bar)
kappa_0 = 10.0 ** cloud_dict["log_kappa_0"] # (cm2 g-1)
# Opacity at 1 um (cm2 g-1) as function of pressure (bar)
# See Eq. 5 in Mollière et al. 2020
kappa_p = kappa_0 * (press_bar / p_base) ** cloud_dict["fsed"]
# Opacity (cm2 g-1) as function of wavelength (um)
# See Eq. 4 in Mollière et al. 2020
kappa_grid, wavel_grid = np.meshgrid(kappa_p, wavel_micron, sparse=True)
kappa_tot = kappa_grid * wavel_grid ** cloud_dict["opa_index"]
kappa_tot[:, press_bar > p_base] = 0.0
# if (
# cloud_dict["opa_knee"] > wavel_micron[0]
# and cloud_dict["opa_knee"] < wavel_micron[-1]
# ):
# indices = np.where(wavel_micron > cloud_dict["opa_knee"])[0]
# for i in range(press_bar.size):
# kappa_tot[indices, i] = (
# kappa_tot[indices[0], i]
# * (wavel_micron[indices] / wavel_micron[indices[0]]) ** -4.0
# )
return (1.0 - cloud_dict["albedo"]) * kappa_tot
@beartype
def kappa_scat(wavel_micron: np.ndarray, press_bar: np.ndarray):
p_base = 10.0 ** cloud_dict["log_p_base"] # (bar)
kappa_0 = 10.0 ** cloud_dict["log_kappa_0"] # (cm2 g-1)
# Opacity at 1 um (cm2 g-1) as function of pressure (bar)
# See Eq. 5 in Mollière et al. 2020
kappa_p = kappa_0 * (press_bar / p_base) ** cloud_dict["fsed"]
# Opacity (cm2 g-1) as function of wavelength (um)
# See Eq. 4 in Mollière et al. 2020
kappa_grid, wavel_grid = np.meshgrid(kappa_p, wavel_micron, sparse=True)
kappa_tot = kappa_grid * wavel_grid ** cloud_dict["opa_index"]
kappa_tot[:, press_bar > p_base] = 0.0
# if (
# cloud_dict["opa_knee"] > wavel_micron[0]
# and cloud_dict["opa_knee"] < wavel_micron[-1]
# ):
# indices = np.where(wavel_micron > cloud_dict["opa_knee"])[0]
# for i in range(press_bar.size):
# kappa_tot[indices, i] = (
# kappa_tot[indices[0], i]
# * (wavel_micron[indices] / wavel_micron[indices[0]]) ** -4.0
# )
return cloud_dict["albedo"] * kappa_tot
elif "log_kappa_abs" in cloud_dict:
# Powerlaw absorption and scattering opacities
@beartype
def kappa_abs(wavel_micron: np.ndarray, press_bar: np.ndarray) -> np.ndarray:
p_base = 10.0 ** cloud_dict["log_p_base"] # (bar)
kappa_0 = 10.0 ** cloud_dict["log_kappa_abs"] # (cm2 g-1)
# Opacity at 1 um (cm2 g-1) as function of pressure (bar)
kappa_p = kappa_0 * (press_bar / p_base) ** cloud_dict["fsed"]
# Opacity (cm2 g-1) as function of wavelength (um)
kappa_grid, wavel_grid = np.meshgrid(kappa_p, wavel_micron, sparse=True)
kappa_abs = kappa_grid * wavel_grid ** cloud_dict["opa_abs_index"]
kappa_abs[:, press_bar > p_base] = 0.0
return kappa_abs
if "log_kappa_sca" in cloud_dict:
@beartype
def kappa_scat(wavel_micron: np.ndarray, press_bar: np.ndarray):
p_base = 10.0 ** cloud_dict["log_p_base"] # (bar)
kappa_0 = 10.0 ** cloud_dict["log_kappa_sca"] # (cm2 g-1)
# Opacity at 1 um (cm2 g-1) as function of pressure (bar)
kappa_p = kappa_0 * (press_bar / p_base) ** cloud_dict["fsed"]
# Opacity (cm2 g-1) as function of wavelength (um)
kappa_grid, wavel_grid = np.meshgrid(kappa_p, wavel_micron, sparse=True)
kappa_sca = kappa_grid * wavel_grid ** cloud_dict["opa_sca_index"]
kappa_sca[:, press_bar > p_base] = 0.0
if (
cloud_dict["lambda_ray"] > wavel_micron[0]
and cloud_dict["lambda_ray"] < wavel_micron[-1]
):
indices = np.where(wavel_micron > cloud_dict["lambda_ray"])[0]
for i in range(press_bar.size):
kappa_sca[indices, i] = (
kappa_sca[indices[0], i]
* (wavel_micron[indices] / wavel_micron[indices[0]]) ** -4.0
)
return kappa_sca
else:
kappa_scat = None
elif "log_kappa_gray" in cloud_dict:
# Gray clouds with cloud top
@beartype
def kappa_abs(wavel_micron: np.ndarray, press_bar: np.ndarray) -> np.ndarray:
p_top = 10.0 ** cloud_dict["log_cloud_top"] # (bar)
kappa_gray = 10.0 ** cloud_dict["log_kappa_gray"] # (cm2 g-1)
opa_abs = np.full((wavel_micron.size, press_bar.size), kappa_gray)
opa_abs[:, press_bar < p_top] = 0.0
return opa_abs
# Add optional scattering opacity
if "albedo" in cloud_dict:
@beartype
def kappa_scat(
wavel_micron: np.ndarray, press_bar: np.ndarray
) -> np.ndarray:
# Absorption opacity (cm2 g-1)
opa_abs = kappa_abs(wavel_micron, press_bar)
# Scattering opacity (cm2 g-1)
opa_scat = cloud_dict["albedo"] * opa_abs / (1.0 - cloud_dict["albedo"])
return opa_scat
else:
kappa_scat = None
else:
kappa_abs = None
kappa_scat = None
# Calculate the emission spectrum
# TODO make cloud_wlen work again
rad_wavel, rad_flux, extra_out = rt_object.calculate_flux(
temperatures=temperature,
mass_fractions=mass_fractions,
mean_molar_masses=mmw,
reference_gravity=10.0**log_g,
cloud_particle_radius_distribution_std=sigma_lnorm,
eddy_diffusion_coefficients=Kzz_use,
cloud_f_sed=fseds,
cloud_photosphere_median_optical_depth=tau_cloud,
additional_absorption_opacities_function=kappa_abs,
additional_scattering_opacities_function=kappa_scat,
# cloud_wlen=cloud_wavel,
return_contribution=contribution,
return_opacities=return_opacities,
)
# if rt_object.flux is None:
# rad_wavel = None
# rad_flux = None
# extra_out = None
# if (
# hasattr(rt_object, "scaling_physicality")
# and rt_object.scaling_physicality > 1.0
# ):
# # cloud_scaling_factor > 2 * (fsed + 1)
# # Set to None such that -inf will be returned as ln_like
# wavel = None
# f_lambda = None
# contr_em = None
#
# if plotting and Kzz_use is None and hasattr(rt_object, "continuum_opa"):
# plt.plot(wavel, rt_object.continuum_opa[:, 0], label="Total continuum opacity")
# plt.plot(
# wavel,
# rt_object.continuum_opa[:, 0] - rt_object.continuum_opa_scat[:, 0],
# label="Absorption continuum opacity",
# )
# plt.plot(
# wavel,
# rt_object.continuum_opa_scat[:, 0],
# label="Scattering continuum opacity",
# )
# plt.xlabel(r"Wavelength ($\mu$m)")
# plt.ylabel("Opacity at smallest pressure")
# plt.yscale("log")
# plt.legend(loc="best")
# plt.savefig("continuum_opacity.png", bbox_inches="tight")
# plt.clf()
# Return wavelength (um), flux (W m-2 um-1), and extra output
return 1e4 * rad_wavel, 1e-7 * rad_flux, extra_out, mmw
[docs]
@beartype
def mass_frac_dict(
log_x_abund: Dict[str, Real],
line_species: List[str],
abund_nodes: Optional[int] = None,
) -> Dict[str, Real]:
"""
Function to return a dictionary with the mass fractions of
all species, including molecular hydrogen and helium.
Parameters
----------
log_x_abund : dict
Dictionary with the log10 of the mass fractions of metals.
line_species : list(str)
List with the line species.
abund_nodes : int, None
Number of abundance nodes. The argument can be set to
``None`` if abundances are constant with pressure.
Returns
-------
dict
Dictionary with the mass fractions of all species.
"""
if abund_nodes is None:
# Initiate abundance dictionary
abund = {}
# Initiate the total mass fraction of the metals
metal_sum = 0.0
for line_item in line_species:
if line_item in log_x_abund:
# Add the mass fraction to the dictionary
abund[line_item] = 10.0 ** log_x_abund[line_item]
# Update the total mass fraction of the metals
metal_sum += abund[line_item]
# Mass fraction of H2 and He
ab_h2_he = 1.0 - metal_sum
# Add H2 and He mass fraction to the dictionary
abund["H2"] = ab_h2_he * 0.75
abund["He"] = ab_h2_he * 0.25
else:
# Initiate abundance dictionary
abund = {}
for node_idx in range(abund_nodes):
# Initiate the total mass fraction of the metals
metal_sum = 0.0
for line_item in line_species:
if f"{line_item}_{node_idx}" in log_x_abund:
# Add the mass fraction to the dictionary
abund[f"{line_item}_{node_idx}"] = (
10.0 ** log_x_abund[f"{line_item}_{node_idx}"]
)
# Update the total mass fraction of the metals
metal_sum += abund[f"{line_item}_{node_idx}"]
# Mass fraction of H2 and He
ab_h2_he = 1.0 - metal_sum
# Add H2 and He mass fraction to the dictionary
abund[f"H2_{node_idx}"] = ab_h2_he * 0.75
abund[f"He_{node_idx}"] = ab_h2_he * 0.25
return abund
[docs]
@beartype
def mean_molecular_weight(mass_frac: Dict[str, Real]) -> Real:
"""
Function to calculate the mean molecular weight from a
dictionary with mass fractions.
Parameters
----------
mass_frac : dict
Dictionary with the mass fraction of each species.
Returns
-------
float
Mean molecular weight in atomic mass units.
"""
from petitRADTRANS.chemistry.utils import simplify_species_list
masses = atomic_masses()
mmw_sum = 0.0
line_orig = list(mass_frac.keys())
line_simple = simplify_species_list(line_orig)
for abund_idx, abund_item in enumerate(line_simple):
if "_" in abund_item:
mmw_sum += mass_frac[abund_item] / masses[abund_item.split("_")[0]]
else:
mmw_sum += mass_frac[line_orig[abund_idx]] / masses[abund_item]
return 1.0 / mmw_sum
[docs]
@beartype
def potassium_abundance(
log_x_abund: Dict[str, Real],
line_species: List[str],
abund_nodes: Optional[int] = None,
) -> Union[Real, List[Real]]:
"""
Function to calculate the mass fraction of potassium at a solar
ratio of the sodium and potassium abundances.
Parameters
----------
log_x_abund : dict
Dictionary with the log10 of the mass fractions.
line_species : list(str)
List with the line species.
abund_nodes : int, None
Number of abundance nodes. The argument can be set to
``None`` if abundances are constant with pressure.
Returns
-------
float, list(float)
Log10 of the mass fraction of potassium or a list with the
log10 mass fraction of potassium in case ``abund_nodes`` is
not ``None``. The length of the list is equal to the argument
of ``abund_nodes`` in that case.
"""
# Solar volume mixing ratios of Na and K (Asplund et al. 2009)
n_na_solar = 1.60008694353205e-06
n_k_solar = 9.86605611925677e-08
# Get the atomic masses
masses = atomic_masses()
# Create a dictionary with all mass fractions
x_abund = mass_frac_dict(log_x_abund, line_species, abund_nodes)
if abund_nodes is None:
# Calculate mean molecular weight from the mass fractions
mmw = mean_molecular_weight(x_abund)
# Volume mixing ratio of sodium
if "Na" in log_x_abund:
n_na_abund = x_abund["Na"] * mmw / masses["Na"]
elif "Na_lor_cut" in log_x_abund:
n_na_abund = x_abund["Na_lor_cut"] * mmw / masses["Na"]
elif "Na_allard" in log_x_abund:
n_na_abund = x_abund["Na_allard"] * mmw / masses["Na"]
elif "Na_burrows" in log_x_abund:
n_na_abund = x_abund["Na_burrows"] * mmw / masses["Na"]
# Volume mixing ratio of potassium
n_k_abund = n_na_abund * n_k_solar / n_na_solar
# Mass fraction of potassium
log_x_k = np.log10(n_k_abund * masses["K"] / mmw)
else:
log_x_k = []
for node_idx in range(abund_nodes):
# Calculate mean molecular weight from the mass fractions
abund_dict = {}
for abund_item in line_species:
abund_dict[abund_item] = x_abund[f"{abund_item}_{node_idx}"]
mmw = mean_molecular_weight(abund_dict)
# Volume mixing ratio of sodium
if f"Na_{node_idx}" in log_x_abund:
n_na_abund = x_abund[f"Na_{node_idx}"] * mmw / masses["Na"]
elif f"Na_lor_cut_{node_idx}" in log_x_abund:
n_na_abund = x_abund[f"Na_lor_cut_{node_idx}"] * mmw / masses["Na"]
elif f"Na_allard_{node_idx}" in log_x_abund:
n_na_abund = x_abund[f"Na_allard_{node_idx}"] * mmw / masses["Na"]
elif f"Na_burrows_{node_idx}" in log_x_abund:
n_na_abund = x_abund[f"Na_burrows_{node_idx}"] * mmw / masses["Na"]
# Volume mixing ratio of potassium
n_k_abund = n_na_abund * n_k_solar / n_na_solar
# Mass fraction of potassium
log_x_k.append(np.log10(n_k_abund * masses["K"] / mmw))
return log_x_k
[docs]
@beartype
def log_x_cloud_base(
c_o_ratio: Real, metallicity: Real, cloud_fractions: Dict[str, Real]
) -> Dict[str, Real]:
"""
Function for returning a dictionary with the log10 mass fractions
at the cloud base.
Parameters
----------
c_o_ratio : float
C/O ratio.
metallicity : float
Metallicity, [Fe/H].
cloud_fractions : dict
Dictionary with the log10 mass fractions at the cloud base,
relative to the maximum values allowed from elemental
abundances. The dictionary keys are the cloud species without
the structure and shape index (e.g. Na2S(c) instead of
Na2S(c)_cd).
Returns
-------
dict
Dictionary with the log10 mass fractions at the cloud base.
Compared to the keys of ``cloud_fractions``, the keys in the
returned dictionary are provided without ``(c)`` (e.g. Na2S
instead of Na2S(c)).
"""
from petitRADTRANS.chemistry.utils import simplify_species_list
log_x_base = {}
# cloud_simple = simplify_species_list(cloud_fractions.keys())
for key, value in cloud_fractions.items():
# Simplify cloud species name
cloud_simple = simplify_species_list([key])
# Mass fraction
x_cloud = cloud_mass_fraction(cloud_simple[0], metallicity, c_o_ratio)
# Log10 of the mass fraction at the cloud base
log_x_base[f"{key}"] = np.log10(10.0**value * x_cloud)
return log_x_base
[docs]
@beartype
def solar_mixing_ratios() -> dict:
"""
Function which returns the volume mixing ratios for solar elemental
abundances (i.e. [Fe/H] = 0); adopted from Asplund et al. (2009).
Returns
-------
dict
Dictionary with the solar number fractions (i.e. volume
mixing ratios).
"""
n_fracs = {}
n_fracs["H"] = 0.9207539305
n_fracs["He"] = 0.0783688694
n_fracs["C"] = 0.0002478241
n_fracs["N"] = 6.22506056949881e-05
n_fracs["O"] = 0.0004509658
n_fracs["Na"] = 1.60008694353205e-06
n_fracs["Mg"] = 3.66558742055362e-05
n_fracs["Al"] = 2.595e-06
n_fracs["Si"] = 2.9795e-05
n_fracs["P"] = 2.36670201997668e-07
n_fracs["S"] = 1.2137900734604e-05
n_fracs["Cl"] = 2.91167958499589e-07
n_fracs["K"] = 9.86605611925677e-08
n_fracs["Ca"] = 2.01439011429255e-06
n_fracs["Ti"] = 8.20622804366359e-08
n_fracs["V"] = 7.83688694089992e-09
n_fracs["Fe"] = 2.91167958499589e-05
n_fracs["Ni"] = 1.52807116806281e-06
return n_fracs
[docs]
@beartype
def atomic_masses() -> dict:
"""
Function which returns the atomic and molecular masses.
Returns
-------
dict
Dictionary with the atomic and molecular masses.
"""
masses = {}
# Atoms
masses["H"] = 1.0
masses["He"] = 4.0
masses["C"] = 12.0
masses["N"] = 14.0
masses["O"] = 16.0
masses["Na"] = 23.0
masses["Na_lor_cur"] = 23.0
masses["Na_allard"] = 23.0
masses["Na_burrows"] = 23.0
masses["Mg"] = 24.3
masses["Al"] = 27.0
masses["Si"] = 28.0
masses["P"] = 31.0
masses["S"] = 32.0
masses["Cl"] = 35.45
masses["K"] = 39.1
masses["K_lor_cut"] = 39.1
masses["K_allard"] = 39.1
masses["K_burrows"] = 39.1
masses["Ca"] = 40.0
masses["Ti"] = 47.9
masses["V"] = 51.0
masses["Fe"] = 55.8
masses["Ni"] = 58.7
# Molecules
masses["H2"] = 2.0
masses["H2O"] = 18.0
masses["H2O_HITEMP"] = 18.0
masses["H2O_main_iso"] = 18.0
masses["CH4"] = 16.0
masses["CH4_main_iso"] = 16.0
masses["CO2"] = 44.0
masses["CO2_main_iso"] = 44.0
masses["CO"] = 28.0
masses["CO_all_iso"] = 28.0
masses["CO_all_iso_Chubb"] = 28.0
masses["CO_all_iso_HITEMP"] = 28.0
masses["NH3"] = 17.0
masses["NH3_main_iso"] = 17.0
masses["HCN"] = 27.0
masses["C2H2,acetylene"] = 26.0
masses["PH3"] = 34.0
masses["PH3_main_iso"] = 34.0
masses["H2S"] = 34.0
masses["H2S_main_iso"] = 34.0
masses["VO"] = 67.0
masses["VO_Plez"] = 67.0
masses["TiO"] = 64.0
masses["TiO_all_Exomol"] = 64.0
masses["TiO_all_Plez"] = 64.0
masses["FeH"] = 57.0
masses["FeH_main_iso"] = 57.0
masses["OH"] = 17.0
return masses
[docs]
@beartype
def cloud_mass_fraction(composition: str, metallicity: Real, c_o_ratio: Real) -> Real:
"""
Function to calculate the mass fraction for a cloud species.
Parameters
----------
composition : str
Cloud composition ('Fe', 'MgSiO3', 'Mg2SiO4', 'Al2O3',
'Na2S', or 'KCL').
metallicity : float
Metallicity [Fe/H].
c_o_ratio : float
Carbon-to-oxygen ratio.
Returns
-------
float
Mass fraction.
"""
# Solar fractional number densities (i.e. volume mixing ratios; VMR)
nfracs = solar_mixing_ratios()
# Atomic masses
masses = atomic_masses()
# Make a copy of the dictionary with the solar number densities
nfracs_use = copy.copy(nfracs)
# Scale the solar number densities by the [Fe/H], except H and He
for item in nfracs:
if item not in ["H", "He"]:
nfracs_use[item] = nfracs[item] * 10.0**metallicity
# Adjust the VMR of O with the C/O ratio
nfracs_use["O"] = nfracs_use["C"] / c_o_ratio
if composition == "Fe":
nfrac_cloud = nfracs_use["Fe"]
mass_cloud = masses["Fe"]
elif composition == "MgSiO3":
nfrac_cloud = np.min(
[nfracs_use["Mg"], nfracs_use["Si"], nfracs_use["O"] / 3.0]
)
mass_cloud = masses["Mg"] + masses["Si"] + 3.0 * masses["O"]
elif composition == "Mg2SiO4":
nfrac_cloud = np.min(
[nfracs_use["Mg"] / 2.0, nfracs_use["Si"], nfracs_use["O"] / 4.0]
)
mass_cloud = 2.0 * masses["Mg"] + masses["Si"] + 4.0 * masses["O"]
elif composition == "Al2O3":
nfrac_cloud = np.min([nfracs_use["Al"] / 2.0, nfracs_use["O"] / 3.0])
mass_cloud = 2.0 * masses["Al"] + 3.0 * masses["O"]
elif composition == "Na2S":
nfrac_cloud = np.min([nfracs_use["Na"] / 2.0, nfracs_use["S"]])
mass_cloud = 2.0 * masses["Na"] + masses["S"]
elif composition == "KCL":
nfrac_cloud = np.min([nfracs_use["K"], nfracs_use["Cl"]])
mass_cloud = masses["K"] + masses["Cl"]
# Cloud mass fraction
x_cloud = mass_cloud * nfrac_cloud
mass_norm = 0.0
for item in nfracs_use:
# Sum up the mass fractions of all species
mass_norm += masses[item] * nfracs_use[item]
# Normalize the cloud mass fraction by the total mass fraction
return x_cloud / mass_norm
[docs]
@beartype
def get_condensation_curve(
composition: str,
press: np.ndarray,
metallicity: Real,
c_o_ratio: Real,
mmw: Real = 2.33,
) -> np.ndarray:
"""
Function to find the base of the cloud deck by intersecting the
P-T profile with the saturation vapor pressure.
Parameters
----------
composition : str
Cloud composition ('Fe', 'MgSiO3', 'Mg2SiO4', 'Al2O3',
'Na2S', or 'KCL').
press : np.ndarray
Pressures (bar).
metallicity : float
Metallicity [Fe/H].
c_o_ratio : float
Carbon-to-oxygen ratio.
mmw : float
Mean molecular weight.
Returns
-------
np.array
Condensation temperatures (K) for the provided input pressures.
"""
if composition == "Fe":
Pc, Tc = return_T_cond_Fe_comb(metallicity, c_o_ratio, mmw)
elif composition == "MgSiO3":
Pc, Tc = return_T_cond_MgSiO3(metallicity, c_o_ratio, mmw)
elif composition == "Mg2SiO4":
Pc, Tc = return_T_cond_Mg2SiO4(metallicity)
elif composition == "Al2O3":
Pc, Tc = return_T_cond_Al2O3(metallicity)
elif composition == "Na2S":
Pc, Tc = return_T_cond_Na2S(metallicity, c_o_ratio, mmw)
elif composition == "KCL":
Pc, Tc = return_T_cond_KCl(metallicity, c_o_ratio, mmw)
else:
raise ValueError(
f"The '{composition}' composition is not "
"supported by get_condensation_curve."
)
index = (Pc > 1e-8) & (Pc < 1e5)
Pc, Tc = Pc[index], Tc[index]
tcond_p = interp1d(Pc, Tc)
return tcond_p(press)
[docs]
@beartype
def find_cloud_deck(
composition: str,
press: np.ndarray,
temp: np.ndarray,
metallicity: Real,
c_o_ratio: Real,
mmw: Real = 2.33,
plotting: bool = False,
) -> Real:
"""
Function to find the base of the cloud deck by intersecting the
P-T profile with the saturation vapor pressure.
Parameters
----------
composition : str
Cloud composition ('Fe', 'MgSiO3', 'Mg2SiO4', 'Al2O3',
'Na2S', or 'KCL').
press : np.ndarray
Pressures (bar).
temp : np.ndarray
Temperatures (K).
metallicity : float
Metallicity [Fe/H].
c_o_ratio : float
Carbon-to-oxygen ratio.
mmw : float
Mean molecular weight.
plotting : bool
Create a plot.
Returns
-------
float
Pressure (bar) at the base of the cloud deck.
"""
from petitRADTRANS.chemistry.utils import simplify_species_list
cloud_simple = simplify_species_list([composition])
Tcond_on_input_grid = get_condensation_curve(
composition=cloud_simple[0],
press=press,
metallicity=metallicity,
c_o_ratio=c_o_ratio,
mmw=mmw,
)
Tdiff = Tcond_on_input_grid - temp
diff_vec = Tdiff[1:] * Tdiff[:-1]
ind_cdf = diff_vec < 0.0
if len(diff_vec[ind_cdf]) > 0:
P_clouds = (press[1:] + press[:-1])[ind_cdf] / 2.0
P_cloud = float(P_clouds[-1])
else:
P_cloud = 1e-8
if plotting:
plt.plot(temp, press)
plt.plot(Tcond_on_input_grid, press)
plt.axhline(P_cloud, color="red", linestyle="--")
plt.yscale("log")
plt.xlim(0.0, 3000.0)
plt.ylim(1e2, 1e-6)
plt.savefig(f"{composition.lower()}_clouds_cdf.png", bbox_inches="tight")
plt.clf()
return P_cloud
[docs]
@beartype
def scale_cloud_abund(
params: Dict[str, Real],
rt_object,
pressure: np.ndarray,
temperature: np.ndarray,
mmw: np.ndarray,
abund_in: Dict[str, np.ndarray],
composition: str,
tau_cloud: Real,
pressure_grid: str,
) -> Real:
"""
Function to scale the mass fraction of a cloud species to the
requested optical depth.
Parameters
----------
params : dict
Dictionary with the model parameters.
rt_object : petitRADTRANS.radtrans.Radtrans
Instance of ``Radtrans``.
pressure : np.ndarray
Array with the pressure points (bar).
temperature : np.ndarray
Array with the temperature points (K) corresponding
to ``pressure``.
mmw : np.ndarray
Array with the mean molecular weights corresponding
to ``pressure``.
abund_in : dict
Dictionary with arrays that contain the pressure-dependent,
equilibrium mass fractions of the line species.
composition : sr
Cloud composition ('Fe(c)', 'MgSiO3(c)', 'Mg2SiO4(c)',
'Al2O3(c)', 'Na2S(c)', 'KCl(c)').
tau_cloud : float
Optical depth of the clouds. The returned mass fraction is
scaled such that the optical depth at the shortest wavelength
is equal to ``tau_cloud``.
pressure_grid : str
The type of pressure grid that is used for the radiative
transfer. Either 'standard', to use 180 layers both for the
atmospheric structure (e.g. when interpolating the abundances)
and 180 layers with the radiative transfer, or 'smaller' to
use 60 (instead of 180) with the radiative transfer, or
'clouds' to start with 1440 layers but resample to ~100 layers
(depending on the number of cloud species) with a refinement
around the cloud decks. For cloudless atmospheres it is
recommended to use 'smaller', which runs faster than 'standard'
and provides sufficient accuracy. For cloudy atmosphere, one
can test with 'smaller' but it is recommended to use 'clouds' for
improved accuracy fluxes.
Returns
-------
float
Mass fraction relative to the maximum value allowed from
elemental abundances. The value has been scaled to the
requested optical depth ``tau_cloud`` (at the shortest
wavelength).
"""
# Dictionary with the requested cloud composition and setting the
# log10 of the mass fraction (relative to the maximum value
# allowed from elemental abundances) equal to zero
cloud_fractions = {composition: 0.0}
# Create a dictionary with the log10 of
# the mass fraction at the cloud base
log_x_base = log_x_cloud_base(
params["c_o_ratio"], params["metallicity"], cloud_fractions
)
# Get the pressure (bar) of the cloud base
p_base = find_cloud_deck(
composition,
pressure,
temperature,
params["metallicity"],
params["c_o_ratio"],
mmw=np.mean(mmw),
plotting=False,
)
# Initialize the cloud abundance in
# the dictionary with mass fractions
abund_in[composition] = np.zeros_like(temperature)
# Set the cloud abundances by scaling
# from the base with the f_sed parameter
abund_in[composition][pressure < p_base] = (
10.0 ** log_x_base[composition]
* (pressure[pressure <= p_base] / p_base) ** params["fsed"]
)
# Adaptive pressure refinement around the cloud base
if pressure_grid == "clouds":
_, indices = make_half_pressure_better({composition: p_base}, pressure)
else:
indices = None
# Create the dictionary with mass fractions
mass_fractions = create_abund_dict(
abund_in,
rt_object.line_species,
rt_object.cloud_species,
pressure_grid=pressure_grid,
indices=indices,
eq_chem=True,
)
# Interpolate the line opacities to the temperature structure
if pressure_grid == "standard":
rt_object.interpolate_species_opa(temperature)
mmw_select = mmw.copy()
if "log_kzz" in params:
kzz_select = np.full(pressure.size, 10.0 ** params["log_kzz"])
else:
# Backward compatibility
kzz_select = np.full(pressure.size, 10.0 ** params["kzz"])
elif pressure_grid == "smaller":
rt_object.interpolate_species_opa(temperature[::3])
mmw_select = mmw[::3]
if "log_kzz" in params:
kzz_select = np.full(pressure[::3].size, 10.0 ** params["log_kzz"])
else:
# Backward compatibility
kzz_select = np.full(pressure[::3].size, 10.0 ** params["kzz"])
elif pressure_grid == "clouds":
# Reinitiate the pressure structure
# after make_half_pressure_better
rt_object.setup_opa_structure(pressure[indices])
rt_object.interpolate_species_opa(temperature[indices])
mmw_select = mmw[indices]
if "log_kzz" in params:
kzz_select = np.full(pressure[indices].size, 10.0 ** params["log_kzz"])
else:
# Backward compatibility
kzz_select = np.full(pressure[indices].size, 10.0 ** params["kzz"])
# Set the continuum opacities to zero because
# calc_cloud_opacity adds to existing opacities
rt_object.continuum_opa = np.zeros_like(rt_object.continuum_opa)
rt_object.continuum_opa_scat = np.zeros_like(rt_object.continuum_opa_scat)
rt_object.continuum_opa_scat_emis = np.zeros_like(rt_object.continuum_opa_scat_emis)
# Calculate the cloud opacities for
# the defined atmospheric structure
rt_object.calc_cloud_opacity(
mass_fractions,
mmw_select,
10.0 ** params["logg"],
params["sigma_lnorm"],
fsed=params["fsed"],
Kzz=kzz_select,
radius=None,
add_cloud_scat_as_abs=False,
)
# Calculate the cloud optical depth and set the tau_cloud attribute
rt_object.calc_tau_cloud(10.0 ** params["logg"])
# Extract the wavelength-averaged optical
# depth at the largest pressure
tau_current = np.mean(rt_object.tau_cloud[0, :, 0, -1])
# Set the continuum opacities again to zero
rt_object.continuum_opa = np.zeros_like(rt_object.continuum_opa)
rt_object.continuum_opa_scat = np.zeros_like(rt_object.continuum_opa_scat)
rt_object.continuum_opa_scat_emis = np.zeros_like(rt_object.continuum_opa_scat_emis)
if tau_current > 0.0:
# Scale the mass fraction
log_x_scaled = np.log10(tau_cloud / tau_current)
else:
log_x_scaled = 100.0
return log_x_scaled
[docs]
@beartype
def cube_to_dict(cube, cube_index: Dict[str, Real]) -> Dict[str, Real]:
"""
Function to convert the parameter cube into a dictionary.
Parameters
----------
cube : LP_c_double
Cube with the parameters.
cube_index : dict
Dictionary with the index of each parameter in the ``cube``.
Returns
-------
dict
Dictionary with the parameters.
"""
params = {}
for key, value in cube_index.items():
params[key] = cube[value]
return params
[docs]
@beartype
def list_to_dict(param_list: List[str], sample_val: np.ndarray) -> Dict[str, Real]:
"""
Function to convert the parameter cube into a dictionary.
Parameters
----------
param_list : list(str)
List with the parameter labels.
sample_val : np.ndarray
Array with the parameter values, in the same order as
``param_list``.
Returns
-------
dict
Dictionary with the parameters.
"""
sample_dict = {}
for item in param_list:
sample_dict[item] = sample_val[param_list.index(item)]
return sample_dict
[docs]
@beartype
def return_T_cond_Fe(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for solid Fe.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
masses = atomic_masses()
T = np.linspace(100.0, 10000.0, 1000)
# Taken from Ackerman & Marley (2001)
# including erratum (P_vap is in bar, not cgs!)
P_vap = lambda x: np.exp(15.71 - 47664.0 / x)
x_cloud = cloud_mass_fraction("Fe", FeH, CO)
return P_vap(T) / (x_cloud * MMW / masses["Fe"]), T
[docs]
@beartype
def return_T_cond_Fe_l(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for liquid Fe.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
masses = atomic_masses()
T = np.linspace(100.0, 10000.0, 1000)
# Taken from Ackerman & Marley (2001)
# including erratum (P_vap is in bar, not cgs!)
P_vap = lambda x: np.exp(9.86 - 37120.0 / x)
x_cloud = cloud_mass_fraction("Fe", FeH, CO)
return P_vap(T) / (x_cloud * MMW / masses["Fe"]), T
[docs]
@beartype
def return_T_cond_Fe_comb(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for Fe.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
P1, T1 = return_T_cond_Fe(FeH, CO, MMW)
P2, T2 = return_T_cond_Fe_l(FeH, CO, MMW)
retP = np.zeros_like(P1)
index = P1 < P2
retP[index] = P1[index]
retP[~index] = P2[~index]
return retP, T2
[docs]
@beartype
def return_T_cond_MgSiO3(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for MgSiO3.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
masses = atomic_masses()
T = np.linspace(100.0, 10000.0, 1000)
# Taken from Ackerman & Marley (2001)
# including erratum (P_vap is in bar, not cgs!)
P_vap = lambda x: np.exp(25.37 - 58663.0 / x)
x_cloud = cloud_mass_fraction("MgSiO3", FeH, CO)
m_mgsio3 = masses["Mg"] + masses["Si"] + 3.0 * masses["O"]
return P_vap(T) / (x_cloud * MMW / m_mgsio3), T
[docs]
@beartype
def return_T_cond_Mg2SiO4(FeH: Real) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for Mg2SiO4.
Parameters
----------
FeH : float
Metallicity.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
# Condensation temperature of Mg2SiO4
# See Eq. 18 in Visscher et al. 2010
temp = np.linspace(100.0, 10000.0, 1000)
P_vap = lambda x: 10.0 ** ((5.89 - 1e4 / temp - 0.73 * FeH) / 0.37)
return P_vap(temp), temp
[docs]
@beartype
def return_T_cond_Al2O3(FeH: Real) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the condensation temperature for Al2O3.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
# Return dictionary with atomic masses
# masses = atomic_masses()
# Create pressures (bar)
pressure = np.logspace(-6, 3, 1000)
# Equilibrium mass fraction of Al2O3
# x_cloud = cloud_mass_fraction('Al2O3', FeH, CO)
# Molecular mass of Al2O3
# m_al2o3 = 3. * masses['Al'] + 2. * masses['O']
# Partial pressure of Al2O3
# part_press = pressure/(x_cloud*MMW/m_al2o3)
# Condensation temperature of Al2O3
# (see Eq. 4 in Wakeford et al. 2017)
t_cond = 1e4 / (
5.014
- 0.2179 * np.log10(pressure)
+ 2.264e-3 * np.log10(pressure) ** 2
- 0.580 * FeH
)
return pressure, t_cond
[docs]
@beartype
def return_T_cond_Na2S(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for Na2S.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
masses = atomic_masses()
# Taken from Charnay+2018
T = np.linspace(100.0, 10000.0, 1000)
# This is the partial pressure of Na, so divide by factor 2 to get
# the partial pressure of the hypothetical Na2S gas particles, this
# is OK: there are more S than Na atoms at solar abundance ratios.
P_vap = lambda x: 1e1 ** (8.55 - 13889.0 / x - 0.5 * FeH) / 2.0
x_cloud = cloud_mass_fraction("Na2S", FeH, CO)
m_na2s = 2.0 * masses["Na"] + masses["S"]
return P_vap(T) / (x_cloud * MMW / m_na2s), T
[docs]
@beartype
def return_T_cond_KCl(
FeH: Real, CO: Real, MMW: Real = 2.33
) -> Tuple[np.ndarray, np.ndarray]:
"""
Function for calculating the saturation pressure for KCl.
Parameters
----------
FeH : float
Metallicity.
CO : float
Carbon-to-oxygen ratio.
MMW : float
Mean molecular weight.
Returns
-------
np.ndarray
Saturation pressure (bar).
np.ndarray
Temperature (K).
"""
masses = atomic_masses()
T = np.linspace(100.0, 10000.0, 1000)
# Taken from Charnay+2018
P_vap = lambda x: 1e1 ** (7.611 - 11382.0 / T)
x_cloud = cloud_mass_fraction("KCL", FeH, CO)
m_kcl = masses["K"] + masses["Cl"]
return P_vap(T) / (x_cloud * MMW / m_kcl), T
[docs]
@beartype
def convolve_spectrum(
input_wavel: np.ndarray, input_flux: np.ndarray, spec_res: Real
) -> np.ndarray:
"""
Function to convolve a spectrum with a Gaussian filter.
Parameters
----------
input_wavel : np.ndarray
Input wavelengths.
input_flux : np.ndarray
Input flux
spec_res : float
Spectral resolution of the Gaussian filter.
Returns
-------
np.ndarray
Convolved spectrum.
"""
# delta_lambda of resolution element is
# FWHM of the LSF's standard deviation
sigma_lsf = 1.0 / spec_res / (2.0 * np.sqrt(2.0 * np.log(2.0)))
# The input spacing of petitRADTRANS is 1e3, but just compute
# it to be sure, or more versatile in the future.
# Also, we have a log-spaced grid, so the spacing is constant
# as a function of wavelength
spacing = np.mean(2.0 * np.diff(input_wavel) / (input_wavel[1:] + input_wavel[:-1]))
# Calculate the sigma to be used with the Gaussian filter
# in units of the input wavelength bins
sigma_filter = sigma_lsf / spacing
return gaussian_filter(input_flux, sigma=sigma_filter, mode="nearest")
[docs]
@beartype
def quench_pressure(
pressure: np.ndarray,
temperature: np.ndarray,
metallicity: Real,
c_o_ratio: Real,
log_g: Real,
log_kzz: Real,
eq_chem,
) -> Optional[Real]:
"""
Function to determine the CO/CH$_4$ quenching pressure by
intersecting the pressure-dependent timescales of the
vertical mixing and the CO/CH$_4$ reaction rates.
Parameters
----------
pressure : np.ndarray
Array with the pressures (bar).
temperature : np.ndarray
Array with the temperatures (K) corresponding to ``pressure``.
metallicity : float
Metallicity [Fe/H].
c_o_ratio : float
Carbon-to-oxygen ratio.
log_g : float
Log10 of the surface gravity (cm s-2).
log_kzz : float
Log10 of the eddy diffusion coefficient (cm2 s-1).
eq_chem : PreCalculatedEquilibriumChemistryTable
Instance of the equilibrium chemistry table from
``petitRADTRANS``.
Returns
-------
float, None
Quenching pressure (bar).
"""
# Interpolate the equilibbrium abundances
co_array = np.full(pressure.size, c_o_ratio)
feh_array = np.full(pressure.size, metallicity)
_, mmw, _ = eq_chem.interpolate_mass_fractions(
co_array,
feh_array,
temperature,
pressure,
carbon_pressure_quench=None,
full=True,
)
# Surface gravity (m s-2)
gravity = 1e-2 * 10.0**log_g
# Mean molecular weight (kg)
mmw *= constants.ATOMIC_MASS
# Pressure scale height (m)
h_scale = constants.BOLTZMANN * temperature / (mmw * gravity)
# Diffusion coefficient (m2 s-1)
kzz = 1e-4 * 10.0**log_kzz
# Mixing timescale (s)
t_mix = h_scale**2 / kzz
# Chemical timescale (see Eq. 12 from Zahnle & Marley 2014)
metal = 10.0**metallicity
t_chem = 1.5e-6 * pressure**-1.0 * metal**-0.7 * np.exp(42000.0 / temperature)
# Determine pressure at which t_mix = t_chem
t_diff = t_mix - t_chem
diff_product = t_diff[1:] * t_diff[:-1]
# If t_mix and t_chem intersect then there
# is 1 negative value in diff_product
indices = diff_product < 0.0
if np.sum(indices) == 1:
p_quench = (pressure[1:] + pressure[:-1])[indices] / 2.0
p_quench = p_quench[0]
elif np.sum(indices) == 0:
p_quench = None
else:
raise ValueError(
f"Encountered unexpected number of indices "
f"({np.sum(indices)}) when determining the "
f"intersection of t_mix and t_chem."
)
return p_quench
[docs]
def convective_flux(
press: np.ndarray,
temp: np.ndarray,
mmw: np.ndarray,
nabla_ad: np.ndarray,
kappa_r: np.ndarray,
density: np.ndarray,
c_p: np.ndarray,
gravity: Real,
f_bol: Real,
mix_length: Real = 1.0,
) -> np.ndarray:
"""
Function for calculating the convective flux with mixing-length
theory. This function has been adopted from petitCODE (Paul
Mollière, MPIA) and was converted from Fortran to Python.
Parameters
----------
press : np.ndarray
Array with the pressures (Pa).
temp : np.ndarray
Array with the temperatures (K) at ``pressure``.
mmw : np.ndarray
Array with the mean molecular weights at ``pressure``.
nabla_ad : np.ndarray
Array with the adiabatic temperature gradient at ``pressure``.
kappa_r : np.ndarray
Array with the Rosseland mean opacity (m2 kg-1) at
``pressure``.
density : np.ndarray
Array with the density (kg m-3) at ``pressure``.
c_p : np.ndarray
Array with the specific heat capacity (J kg-1 K-1) at
constant pressure, ``pressure``.
gravity : float
Surface gravity (m s-2).
f_bol : float
Bolometric flux (W m-2) at the top of the atmosphere,
calculated from the low-resolution spectrum.
mix_length : float
Mixing length for the convection in units of the pressure
scale height (default: 1.0).
Returns
-------
np.ndarray
Convective flux (W m-2) at each pressure.
"""
t_transp = (f_bol / constants.SIGMA_SB) ** 0.25 # (K)
nabla_rad = (
3.0 * kappa_r * press * t_transp**4.0 / 16.0 / gravity / temp**4.0
) # (dimensionless)
h_press = (
constants.BOLTZMANN * temp / (mmw * constants.ATOMIC_MASS * gravity)
) # (m)
l_mix = mix_length * h_press # (m)
U = (
(12.0 * constants.SIGMA_SB * temp**3.0)
/ (c_p * density**2.0 * kappa_r * l_mix**2.0)
* np.sqrt(8.0 * h_press / gravity)
)
W = nabla_rad - nabla_ad
# FIXME thesis: 2336U^4W
A = (
1168.0 * U**3.0
+ 2187 * U * W
+ 27.0
* np.sqrt(
3.0 * (2048.0 * U**6.0 + 2236.0 * U**4.0 * W + 2187.0 * U**2.0 * W**2.0)
)
) ** (1.0 / 3.0)
xi = (
19.0 / 27.0 * U
- 184.0 / 27.0 * 2.0 ** (1.0 / 3.0) * U**2.0 / A
+ 2.0 ** (2.0 / 3.0) / 27.0 * A
)
nabla = xi**2.0 + nabla_ad - U**2.0
nabla_e = nabla_ad + 2.0 * U * xi - 2.0 * U**2.0
f_conv = (
density
* c_p
* temp
* np.sqrt(gravity)
* (mix_length * h_press) ** 2.0
/ (4.0 * np.sqrt(2.0))
* h_press**-1.5
* (nabla - nabla_e) ** 1.5
)
f_conv[np.isnan(f_conv)] = 0.0
return f_conv # (W m-2)
