"""
Spectroscopy utility functions for GUIBRUSHR atmospheric retrieval.
This module contains functions for spectroscopic data processing extracted from
general_functions.py:
- kernel_solid_body_rotation: Compute a rotational broadening kernel.
- convolve_solid_body_rotation: Convolve a spectrum with the rotational kernel.
- convolve_resolution: Degrade a spectrum to instrumental resolution with Doppler shift.
- calculate_mass_molecule: Calculate molecular mass from elemental composition.
- generate_chemcat: Generate a chemical catalog for atmospheric modeling.
- skewed_gaussian: Evaluate a skewed Gaussian (helper for estimate_continuum).
- estimate_continuum: Estimate and remove the continuum from a spectrum.
- compute_adjusted_abundance: Compute abundance accounting for dissociation effects.
- spectral_convolution_ptr: Convolve a spectrum using petitRADTRANS.
"""
from pathlib import Path
from re import findall, sub
from typing import List, Tuple, Union
import chemcat
import numpy as np
import pyratbay.io as io
from petitRADTRANS import physical_constants as pc # pc = petitRADTRANS physical_constants
from petitRADTRANS.spectral_model import SpectralModel
from scipy.optimize import curve_fit
from scipy.interpolate import splrep, splev
from scipy.signal import convolve
from scipy.stats import skewnorm
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
[docs]
def kernel_solid_body_rotation(
omegad: float,
f_rot_arr: np.ndarray,
Teq: float,
mu: float,
gs: float,
Rpj: float,
rad_mode: str,
index_nights_total: int,
rv_sampling: float = 0.1
) -> np.ndarray:
"""
Compute a rotational broadening kernel for spectroscopic observations.
A planet rotating as a rigid body with angular frequency omega broadens
spectral lines because different surface elements of the visible disk move
at different line-of-sight velocities: an element at projected distance x
from the rotation axis contributes flux at a radial velocity offset
v = omega * x. The rotational kernel K(v) is the disk-integrated weight
of all surface elements emitting at velocity v, and convolving a model
spectrum with K(v) reproduces the line broadening induced by planetary
rotation. The overall width of the kernel is set by the equatorial
rotational velocity v_eq = omega * Rp.
The rotation rate omega is a free parameter of the retrieval and is not
constrained to any particular value. For a tidally locked planet it equals
2*pi / P_orb, but the kernel is valid for any rotation rate.
The kernel geometry differs between emission and transmission spectroscopy:
- In emission spectroscopy, only the planetary disk itself contributes.
Each surface strip at projected distance x has a chord length
2*sqrt(Rp^2 - x^2), so K(v) is proportional to sqrt(1 - (v/v_eq)^2)
for \|v\| < v_eq, and zero outside. This is the classical solid-body
rotational kernel for a uniformly emitting disk: it peaks at v=0
(sub-observer point, longest chord) and drops to zero at v=±v_eq
(planetary limb, chord=0).
- In transmission spectroscopy, the relevant region is the atmospheric
annulus between Rp and Rp + z, where z = 5*H is the effective
atmospheric thickness and H = k_B * T_eq / (mu * m_H * g) is the
scale height. Defining v_eq_atm = omega * (Rp + z), two velocity
regions contribute differently:
- v_eq <= \|v\| < v_eq_atm (annulus only): the line of sight passes
only through the annulus, with weight
sqrt((Rp+z)^2 - (v/omega)^2) / z.
- \|v\| < v_eq (annulus above opaque disk): the line of sight passes
through the full annulus minus the opaque planetary disk, with
weight (sqrt((Rp+z)^2 - (v/omega)^2) - sqrt(Rp^2 - (v/omega)^2)) / z.
An asymmetry parameter f_rot (stored and passed in linear units in f_rot_arr,
converted from log10 by to_linear at extraction time) is applied to the
receding hemisphere (v > 0) to allow for day/night asymmetries in the
transmission signal. It is a free parameter of the retrieval.
In both modes the kernel is normalized to unit sum before being returned,
so that convolution preserves the total flux of the spectrum.
Parameters
----------
omegad : float
Rotational angular frequency of the planet, in inverse days (d^-1).
The equatorial rotational velocity is v_eq = omega * Rp, where omega
is omegad converted to s^-1. This is a free parameter of the retrieval
and is not assumed to equal any particular value.
f_rot_arr : np.ndarray
Array of linear rotational asymmetry scaling factors (converted from
log10 by to_linear at extraction time), one per observing night. Used
only in transmission mode: f_rot_arr[index_nights_total] is multiplied
onto the kernel for v > 0 to model day/night hemisphere asymmetries.
It is a free parameter of the retrieval.
Teq : float
Equilibrium temperature of the planet, in Kelvin. Used to compute
the atmospheric scale height H = k_B * T_eq / (mu * m_H * g),
which defines the vertical extent z = 5*H of the transmission
annulus. Only used in transmission mode.
mu : float
Mean molecular weight of the atmosphere, in units of the hydrogen
mass m_H. Used to compute the scale height. Only used in transmission
mode.
gs : float
Surface gravity of the planet, in cgs units (cm s^-2). Used to
compute the scale height. Only used in transmission mode.
Rpj : float
Planetary radius, in centimeters. Defines v_eq = omega * Rp (emission)
and the inner boundary of the transmission annulus (transmission).
rad_mode : str
Observing mode. Must be one of ConstantVariables.LIST_RAD_MODE:
- LIST_RAD_MODE[0]: transmission spectroscopy (annulus kernel).
- any other value: emission spectroscopy (solid disk kernel).
index_nights_total : int
Index into f_rot_arr selecting the asymmetry factor for the current
observing night. Only used in transmission mode.
rv_sampling : float, optional
Velocity step of the radial velocity grid on which the kernel is
computed, in km/s. Default is 0.1 km/s, fine enough to well-sample
the kernel profile for any realistic v_eq.
Returns
-------
np.ndarray
Normalized rotational broadening kernel, sampled on a velocity grid
spanning -30 to +30 km/s with step rv_sampling. The kernel sums to
unity by construction.
"""
# Physical constants
kb = pc.kB # Boltzmann constant (erg/K)
mH = 1.6738e-24 # Hydrogen mass (g)
# --- Atmospheric scale height ---
# H = k_B * T_eq / (mu * m_H * g), in cm.
# Used only in transmission mode to define the vertical extent of the
# absorbing annulus: z = 5*H is a standard approximation for the
# effective atmospheric thickness probed in transmission.
H = kb * Teq / (mu * mH * gs)
z = 5 * H
# --- Radial velocity grid ---
# The kernel is computed on a grid from -30 to +30 km/s with step
# rv_sampling (default 0.1 km/s), giving 601 points. This range is
# sufficient for any realistic v_eq: for WASP-77Ab (omega=0.665 d^-1,
# Rpj=1.21 Rj), v_eq ~ 0.67 km/s, well within ±30 km/s.
rv_array = np.arange(-30, 30 + rv_sampling, rv_sampling)
rv_array_cm = rv_array * 1e5 # convert km/s to cm/s
# --- Convert omega from d^-1 to s^-1 ---
# 1 d^-1 = 1/86400 s^-1 = 1.1574e-5 s^-1
omega = omegad * 1.157407407407407e-05 # s^-1
# --- Projected distance x on the disk ---
# For a surface element moving at line-of-sight velocity v = omega * x,
# the projected distance from the rotation axis is x = v / omega.
# Here rv_array_cm plays the role of v (in cm/s), so x is in cm.
# x = rv_array_cm / omega
ker_rot = np.zeros(len(rv_array))
# Equatorial rotational velocity: v_eq = omega * Rp (cm/s)
# Defines the velocity at the planetary limb (x = ±Rp).
# In transmission, also define v_eq_atm = omega * (Rp + z),
# the velocity at the outer edge of the atmospheric annulus.
v_eq = omega * Rpj # cm/s
v_eq_atm = omega * (Rpj + z) # cm/s
if rad_mode == ConstantVariables.LIST_RAD_MODE[0]: # Transmission mode
# f_rot arrives already in linear units (Log10ToLinear applied by to_linear)
f_rot = f_rot_arr[index_nights_total]
# Region 1 (annulus only): v_eq <= |v| < v_eq_atm
# The line of sight passes only through the atmospheric annulus.
mask1 = (np.abs(rv_array_cm) >= v_eq) & (np.abs(rv_array_cm) < v_eq_atm)
# Region 2 (annulus above opaque disk): |v| < v_eq
# The line of sight passes through the full annulus minus the opaque disk.
mask2 = np.abs(rv_array_cm) < v_eq
ker_rot[mask1] = np.sqrt((Rpj + z) ** 2 - (rv_array_cm[mask1] / omega) ** 2) / z
ker_rot[mask2] = (
np.sqrt((Rpj + z) ** 2 - (rv_array_cm[mask2] / omega) ** 2) -
np.sqrt(Rpj ** 2 - (rv_array_cm[mask2] / omega) ** 2)
) / z
ker_rot[rv_array_cm > 0] *= f_rot
else: # Emission mode
# K(v) ∝ sqrt(1 - (v/v_eq)^2), zero for |v| > v_eq
mask = np.abs(rv_array_cm) < v_eq
ker_rot[mask] = np.sqrt(1 - (rv_array_cm[mask] / v_eq) ** 2)
# Normalize to unit sum so that convolution preserves total flux.
return ker_rot / np.sum(ker_rot)
[docs]
def convolve_solid_body_rotation(
wl: np.ndarray,
model: np.ndarray,
ker_rot: np.ndarray,
rv_sampling: float = 0.1,
n_el: int = 321
) -> Tuple[np.ndarray, np.ndarray]:
"""
Convolve a model spectrum with a rotational broadening kernel.
This function is the second convolution step in the model preparation
pipeline, applied after convolve_resolution. It takes the rotational
kernel produced by kernel_solid_body_rotation (defined on a fine velocity
grid with step rv_sampling) and convolves it with the model spectrum.
The procedure mirrors that of convolve_resolution:
1. The pixel velocity step of the model wavelength grid is computed as
dv = c * dlambda / lambda, constant by construction for a log-uniform
grid (constant resolving power R_model).
2. The kernel is resampled from its fine velocity grid (step rv_sampling,
default 0.1 km/s) onto a coarser grid whose spacing matches the model
pixel velocity step, using spline interpolation. This step is mandatory
because discrete convolution requires kernel and signal to share the
same grid spacing.
3. The resampled kernel is convolved with the model in 'valid' mode
(scipy.signal.convolve), discarding the n_el//2 edge pixels on each
side where the kernel does not fully overlap the signal, and the result
is normalized by the kernel sum to ensure flux conservation.
Note: unlike convolve_resolution, this function does not apply a radial
velocity shift. The shift is handled entirely by convolve_resolution via
kernel translation. This function is responsible only for the rotational
line broadening.
Parameters
----------
wl : np.ndarray
1D array of wavelength values for the model spectrum. Must be
log-uniform (constant resolving power R_model), so that the pixel
velocity step dv = c / R_model is constant across the grid.
model : np.ndarray
1D array of spectral flux values corresponding to wl. This is the
spectrum already convolved with the instrumental LSF by
convolve_resolution.
ker_rot : np.ndarray
Rotational broadening kernel computed by kernel_solid_body_rotation,
sampled on a velocity grid from -30 to +30 km/s with step rv_sampling.
Must be normalized to unit sum.
rv_sampling : float, optional
Velocity step of the input kernel grid, in km/s. Must match the
rv_sampling used in kernel_solid_body_rotation. Default is 0.1 km/s.
n_el : int, optional
Number of elements in the resampled kernel. Must be odd to ensure
symmetry around zero and avoid introducing a spurious velocity shift.
The half-extent of the resampled grid is (n_el - 1) / 2 pixels.
Default is 321, giving a half-extent of 160 pixels.
Returns
-------
wavelength : np.ndarray
Wavelength array of the convolved spectrum, shorter than the input
by n_el // 2 = 160 points on each side due to 'valid' mode convolution.
rconv : np.ndarray
Rotationally broadened flux array, flux-conserving by construction.
"""
# --- Radial velocity grid for the input kernel ---
# This must match the grid used in kernel_solid_body_rotation.
rv_array = np.arange(-30, 30 + rv_sampling, rv_sampling)
# --- Pixel velocity step of the model wavelength grid ---
# dv = c * dlambda / lambda, constant by construction for a log-uniform
# grid. Computed on each adjacent pixel pair and averaged as a numerical
# verification (see convolve_resolution Step 2 for a detailed discussion).
dv = ConstantVariables.CLIGHT * (wl[1:] - wl[:-1]) / wl[:-1]
rv_pix = np.mean(dv)
# --- Resampled velocity grid matching the model pixel step ---
# A grid of n_el points (odd, default 321) with spacing rv_pix, spanning
# ±(n_el-1)/2 * rv_pix = ±160 * rv_pix km/s. The odd number of points
# ensures symmetry around zero, preventing spurious velocity offsets.
# (See convolve_resolution Step 3 for a detailed discussion.)
max_val = ((n_el - 1) / 2) * rv_pix
rv_array_mod = np.linspace(-max_val, max_val, n_el)
n_rv = len(rv_array_mod)
# --- Resample the kernel onto the model grid via spline interpolation ---
# splrep builds a B-spline representation of ker_rot on rv_array, and
# splev evaluates it on rv_array_mod. After this step the kernel has
# the same velocity spacing as the model, making discrete convolution
# well-defined. (See convolve_resolution Step 4b for a detailed discussion.)
csscaled = splrep(rv_array, ker_rot)
ker_rot_pix = splev(rv_array_mod, csscaled, der=0)
# --- Convolve in 'valid' mode and normalize ---
# scipy.signal.convolve in 'valid' mode returns only the n - K + 1 central
# output points where the kernel fully overlaps the signal, discarding the
# n_rv // 2 = 160 edge pixels on each side. The wavelength array is trimmed
# by the same amount. A final normalization by np.sum(ker_rot_pix) corrects
# for any residual deviation from unity introduced by the spline resampling,
# ensuring strict flux conservation. (See convolve_resolution Step 5 for a
# detailed discussion.)
ini = int(n_rv / 2)
wavelength = wl[ini:-ini]
rconv = convolve(model, ker_rot_pix, mode="valid") / np.sum(ker_rot_pix)
return wavelength, rconv
[docs]
def convolve_resolution(
wavelength_model: np.ndarray,
flux_model: np.ndarray,
lsf_hwhm: float,
lsf_grid_step: float = 0.1,
kernel_size: int = 321
) -> Tuple[np.ndarray, np.ndarray]:
"""
Degrade a high-resolution model spectrum to instrumental resolution.
A real spectrograph cannot reproduce an infinitely narrow spectral line: even
a perfectly monochromatic source produces a broadened profile on the detector,
due to diffraction, optical aberrations, slit width, and pixel sampling. This
broadening is described by the Line Spread Function (LSF). To compare a
synthetic model with observed data, the model must be degraded to the same
instrumental resolution, i.e. convolved with the LSF.
The LSF is approximated as a Gaussian whose width is expressed in velocity
units. Working in velocity space is convenient because the LSF width is
approximately constant in velocity across the full spectral range, whereas
in wavelength units it varies with lambda. The Half Width at Half Maximum
is HWHM = c / (2R), where R is the nominal resolving power of the instrument,
and the corresponding Gaussian standard deviation is sigma = HWHM / sqrt(2*ln(2)).
The factor sqrt(2*ln(2)) ~ 1.177 comes from setting the Gaussian equal to
half its peak value and solving for v: G(v_1/2) = exp(-v_1/2^2 / 2*sigma^2) = 1/2
gives v_1/2 = sigma * sqrt(2*ln(2)), which is by definition the HWHM.
This function performs ONLY the resolution degradation. It does not apply any
radial velocity shift: the kernel is built symmetric around zero velocity, so
convolution preserves the spectral positions. All Doppler shifts (orbital,
systemic, barycentric and the retrieval rv via rv_dynamic) are applied as a
single velocity term on the data wavelength grid by the caller, through
calculate_rv_planet_and_star (see cross_correlation.py). This keeps the radial
velocity computed in one place rather than split between the kernel and vtot.
Parameters
----------
wavelength_model : np.ndarray
1D array of wavelength values for the model spectrum. Must be computed
at constant resolving power R_model, so that the ratio dlambda/lambda
= 1/R_model is constant by construction (log-uniform grid), yielding
a constant velocity step dv = c / R_model per pixel.
flux_model : np.ndarray
1D array of flux values corresponding to wavelength_model.
lsf_hwhm : float
Half Width at Half Maximum of the instrumental LSF, in km/s.
Related to spectral resolution by: HWHM = c / (2R).
Example: IGRINS (R ~ 45000) has HWHM ~ 3.3 km/s.
lsf_grid_step : float, optional
Velocity step for the initial fine LSF kernel grid, in km/s. Must be
small enough to well-sample the Gaussian peak. Default is 0.1 km/s.
kernel_size : int, optional
Number of elements in the resampled convolution kernel. Must be odd
to ensure the grid is symmetric around zero, which avoids introducing
a spurious velocity shift during convolution. The half-extent of the
resampled grid is (kernel_size - 1) / 2 pixels. Default is 321,
giving a half-extent of 160 pixels.
Returns
-------
wavelength_out : np.ndarray
Wavelength array of the convolved spectrum. Shorter than the input by
kernel_size // 2 = 160 points on each side, because scipy.signal.convolve
in 'valid' mode discards the edge pixels where the kernel does not fully
overlap the signal (see Step 5).
flux_convolved : np.ndarray
Convolved flux array, flux-conserving by construction.
"""
# --- Step 1: Build the Gaussian LSF kernel on a fine velocity grid ---
# The grid spans ±30 km/s with a step of lsf_grid_step (default 0.1 km/s),
# yielding 601 points. This range is intentionally much wider than the
# Gaussian itself: for IGRINS (HWHM ~ 3.3 km/s) the kernel drops to
# negligible values already beyond ±10 km/s, so nothing is lost at the edges.
#
# HWHM_TO_SIGMA = sqrt(2 * ln(2)) ~ 1.177 converts HWHM to Gaussian sigma.
# It comes from solving G(v_1/2) = 1/2:
# exp(-v_1/2^2 / 2*sigma^2) = 1/2 => v_1/2 = sigma * sqrt(2*ln(2)) = HWHM
# therefore: sigma = HWHM / sqrt(2*ln(2))
#
# The kernel is normalized to unit sum so that the convolution preserves
# the total flux of the spectrum.
lsf_velocity_grid = np.arange(-30, 30 + lsf_grid_step, lsf_grid_step)
HWHM_TO_SIGMA = 1.177 # = sqrt(2 * ln(2))
lsf_sigma = lsf_hwhm / HWHM_TO_SIGMA
lsf_kernel = np.exp(-0.5 * (lsf_velocity_grid / lsf_sigma) ** 2)
lsf_kernel = lsf_kernel / np.sum(lsf_kernel)
# --- Step 2: Compute the velocity step of the model wavelength grid ---
# The model is computed at constant resolving power R_model, so its
# wavelength grid has dlambda/lambda = 1/R_model = constant by construction
# (log-uniform grid). The corresponding pixel velocity step is therefore
# dv = c * dlambda/lambda = c / R_model, constant across the full grid.
# We compute it explicitly on each pair of adjacent pixels and take the
# mean as a numerical verification. This value dv defines the spacing of
# the resampled kernel in the next step.
# Example: for R_model = 2.5e5 (petitRADTRANS at R=1e6, downsampled by 4),
# dv = 3e5 / 2.5e5 = 1.2 km/s per pixel.
velocity_step_per_pixel = ConstantVariables.CLIGHT * (wavelength_model[1:] - wavelength_model[:-1]) / wavelength_model[:-1]
mean_velocity_step = np.mean(velocity_step_per_pixel)
# --- Step 3: Build the resampled velocity grid matching the model ---
# The discrete convolution requires that the kernel and the model spectrum
# share the same velocity grid spacing. We therefore create a new grid of
# kernel_size points (odd, default 321) with spacing mean_velocity_step,
# so that each kernel element corresponds to exactly one model pixel.
# The half-extent is (kernel_size - 1) / 2 = 160 pixels, corresponding to
# ±160 * 1.2 = ±192 km/s for R_model = 2.5e5. This is sufficient to
# accommodate both the LSF width (a few km/s) and any realistic planetary
# radial velocity shift.
# The odd kernel_size ensures the grid is symmetric around zero, preventing
# any spurious velocity offset from being introduced during convolution.
half_kernel_extent = ((kernel_size - 1) / 2) * mean_velocity_step
resampled_velocity_grid = np.linspace(-half_kernel_extent, half_kernel_extent, kernel_size)
# --- Step 4: Resample the LSF kernel onto the model velocity grid ---
# splrep/splev interpolates the kernel from the fine grid (step = 0.1 km/s,
# 601 points) onto the coarser grid built in Step 3 (step = mean_velocity_step,
# 321 points). After this step, kernel and model share the same grid spacing
# and discrete convolution is well-defined. The kernel stays centered on zero
# velocity: no radial velocity shift is applied here (all Doppler shifts are
# handled on the data wavelength grid via vtot in calculate_rv_planet_and_star).
kernel_spline = splrep(lsf_velocity_grid, lsf_kernel)
resampled_lsf_kernel = splev(resampled_velocity_grid, kernel_spline, der=0)
# --- Step 5: Convolve and return ---
# scipy.signal.convolve is called in 'valid' mode. In a discrete convolution
# between a signal of N points and a kernel of K points, the output has
# N + K - 1 points in 'full' mode. However, the first and last K//2 = 160
# output points correspond to positions where the kernel extends beyond the
# signal boundary, requiring the signal to be padded with zeros — which
# introduces boundary artefacts. 'valid' mode discards these edge points
# and returns only the N - K + 1 central values where the kernel overlaps
# the signal completely. The wavelength array is trimmed by the same amount.
#
# A final normalization by np.sum(resampled_lsf_kernel) corrects for any
# residual deviation from unity introduced by the spline resampling in
# Step 4, ensuring strict flux conservation in the output.
half_kernel_size = kernel_size // 2
wavelength_out = wavelength_model[half_kernel_size:-half_kernel_size]
flux_convolved = convolve(flux_model, resampled_lsf_kernel, mode="valid") / np.sum(resampled_lsf_kernel)
return wavelength_out, flux_convolved
[docs]
def calculate_mass_molecule(elements, periodic_table):
"""
Calculate the total molecular mass of a molecule from its elemental composition.
This function parses element strings that may contain mass numbers, atom counts,
and ionization states, then calculates the total molecular mass using atomic
masses from the periodic table.
Args:
elements: List of element strings with format like 'H2', '16O', 'C1+', etc.
periodic_table: DataFrame containing periodic table data with columns
'Symbol', 'NumberofProtons', 'NumberofNeutrons'
Returns:
float: Total molecular mass in atomic mass units (u)
"""
total_mass = 0
# Define atomic masses (in atomic mass units, u)
proton_mass = 1.007276 # u
neutron_mass = 1.008665 # u
electron_mass = 0.000548579909 # u
#
for element in elements:
# Count the number of removed electrons indicated by '+' symbols
electrons_removed = element.count("+")
# Remove numerical parts and extra characters to obtain the elemental symbol
symbol = sub(r'[0-9]+', '', element).replace("+", "").replace("_", "")
if len(symbol) > 4:
# Skip elements with an unexpectedly long symbol
continue
# Extract all numeric parts from the element string (e.g., mass numbers and counts)
amounts = findall(r'\d+', element)
if len(amounts) > 1:
# If two numbers are present: the first is the mass number, the second is the count
proton_plus_neutron = amounts[0]
atom_in_molecule = amounts[1]
elif len(amounts) == 1:
# If one number is present, determine its role by checking its position relative to the symbol
if symbol in element[:len(symbol)]:
# Use periodic table values if the number is part of standard element notation
table_row = periodic_table[periodic_table["Symbol"] == element[:len(symbol)]]
proton_plus_neutron = int(table_row["NumberofNeutrons"].iloc[0]) + int(table_row["NumberofProtons"].iloc[0])
atom_in_molecule = amounts[0]
else:
proton_plus_neutron = amounts[0]
atom_in_molecule = "1"
else:
# If no numeric value is found, default to using the periodic table (assume count = 1)
table_row = periodic_table[periodic_table["Symbol"] == element[:len(symbol)]]
proton_plus_neutron = int(table_row["NumberofNeutrons"].iloc[0]) + int(table_row["NumberofProtons"].iloc[0])
atom_in_molecule = "1"
# Retrieve the number of protons (and by extension, electrons) for the element
number_protons = number_electrons = int(
periodic_table[periodic_table["Symbol"] == symbol]["NumberofProtons"].iloc[0]
)
# Calculate the number of neutrons from the total mass number minus the number of protons
number_neutrons = int(proton_plus_neutron) - number_protons
# Add the mass contribution of this element to the total mass of the isotope
total_mass += ((number_protons * proton_mass +
number_neutrons * neutron_mass +
(number_electrons - electrons_removed) * electron_mass) *
int(atom_in_molecule))
return total_mass
[docs]
def generate_chemcat(
min_pressure: float,
max_pressure: float,
nlayersf: int,
path_default: Union[str, Path]
) -> Tuple[np.ndarray, chemcat.Network, List[str]]:
"""
Generate a chemical catalog and retrieve molecular masses for atmospheric modeling.
This function creates a chemical network with specified pressure layers and temperature,
and prepares molecular data for atmospheric modeling.
Parameters:
min_pressure (float): Logarithm of minimum pressure (bar)
max_pressure (float): Logarithm of maximum pressure (bar)
nlayersf (int): Number of pressure layers
path_default (Union[str, Path]): Base path to default files directory
Returns:
Tuple containing:
- np.ndarray: Array of molecular masses for selected molecules
- chemcat.Network: Chemical network object
- List[str]: List of species names compatible with petitRADTRANS
"""
# Construct path to molecules data file
molecules_path = Path(path_default) / "GUIBRUSHR" / "Files" / "Molecules" / "molecules.dat"
names, masses_from_file, diam = None, None, None
try:
# Read molecular data
names, masses_from_file, diam = io.read_molecs(str(molecules_path))
except FileNotFoundError:
print(f"Molecules data file not found at {molecules_path}")
except Exception as e:
print(f"Error reading molecular data: {str(e)}")
# Build list of molecules from various chemical groups
molecules = (
ConstantVariables.HCNO_NEUTRALS.split() +
ConstantVariables.IONS.split() +
ConstantVariables.ALKALI.split() +
ConstantVariables.METALS.split() +
ConstantVariables.METAL_OXIDES.split()
)
# Format species names for petitRADTRANS compatibility
species_compatible_with_prt = [
mol if mol in ConstantVariables.IONS else mol.replace("+", "_+")
for mol in molecules
]
# Generate pressure grid and temperature array
pressure_grid = np.logspace(min_pressure, max_pressure, nlayersf)
temp = np.full(nlayersf, 1000.0) # Constant temperature of 1000 K
# Create chemical network
chemcat_network = chemcat.Network(pressure_grid, temp, molecules)
# Extract molecular masses
try:
molecular_masses = np.array([masses_from_file[names == spec][0] for spec in molecules])
except IndexError:
print("One or more molecules not found in molecular data")
return molecular_masses, chemcat_network, species_compatible_with_prt
[docs]
def skewed_gaussian(x, amp, center, width, skew):
"""
Calculate a skewed Gaussian distribution.
Args:
x: Input values
amp: Amplitude of the distribution
center: Center position
width: Width parameter
skew: Skewness parameter
Returns:
Values of the skewed Gaussian distribution
"""
return amp * skewnorm.pdf(x, skew, loc=center, scale=width)
[docs]
def estimate_continuum(spectrum_new, num_windows=40, top_window=10, check_gaussian=False):
"""
Estimate and remove the continuum from a spectrum using windowed median fitting.
Args:
spectrum_new (np.ndarray): Input spectrum array.
num_windows (int): Number of windows to divide the spectrum into.
top_window (int): Number of top values per window used for continuum estimation.
check_gaussian (bool): If True, fit a skewed Gaussian instead of a polynomial.
Returns:
tuple: (continuum, model_after_processing) where continuum is the estimated
continuum and model_after_processing is the spectrum normalized by it.
"""
window_size = len(spectrum_new) // num_windows
medians_continuum = []
wl_continuum = []
# Calculate medians_continuum for continuum fitting
for i in range(num_windows):
xv = (i * window_size + (i + 1) * window_size) / 2
window_element = spectrum_new[i * window_size: (i + 1) * window_size]
top_10 = np.sort(window_element)[-top_window:]
medians_continuum.append(np.median(top_10))
wl_continuum.append(xv)
wl_continuum = np.array(wl_continuum)
medians_continuum = np.array(medians_continuum)
try:
if check_gaussian:
p0 = [1, np.mean(wl_continuum), np.std(wl_continuum) * 2, 0]
params, _ = curve_fit(skewed_gaussian, wl_continuum, medians_continuum, p0=p0)
# skewed_fit = skewed_gaussian(wl_continuum, *params)
continuum = skewed_gaussian(np.arange(len(spectrum_new)), *params)
else:
par = np.polyfit(wl_continuum, medians_continuum, 4)
continuum = np.polyval(par, np.arange(len(spectrum_new)))
except Exception as e:
par = np.polyfit(wl_continuum, medians_continuum, 4)
continuum = np.polyval(par, np.arange(len(spectrum_new)))
model_after_processing = spectrum_new / continuum # we want the continuum at one, not at 0.9 or similar
return continuum, model_after_processing
[docs]
def compute_adjusted_abundance(
pressure: np.ndarray,
temperature: np.ndarray,
coeff,
base_vmr
) -> np.ndarray:
"""
Compute adjusted abundance accounting for dissociation effects.
This function calculates the adjusted abundance of a species by considering
both the base volume mixing ratio and dissociation effects. The dissociation
is modeled using a pressure and temperature dependent function.
Parameters:
pressure (np.ndarray): Pressure values in bar
temperature (np.ndarray): Temperature values in Kelvin
coeff: Coefficients [a, b, c] for dissociation calculation
base_vmr: Base volume mixing ratio from reference layer
Returns:
np.ndarray: Adjusted abundance values
The calculation follows these steps:
1. Compute dissociation factor (Ad) using pressure and temperature
2. Convert to inverse dissociation factor (1/Ad)
3. Combine with inverse base abundance (1/A0) in quadrature
4. Return final adjusted abundance (A)
"""
# Calculate dissociation factor
log10Ad = coeff[0] * np.log10(pressure) + coeff[1] / temperature - coeff[2]
one_over_Ad = 1.0 / (10 ** log10Ad)
# Calculate inverse of base abundance
one_over_A0 = 1.0 / base_vmr
# Combine terms in quadrature
one_over_A = (np.sqrt(one_over_A0) + np.sqrt(one_over_Ad)) ** 2
# Return final adjusted abundance
return 1.0 / one_over_A
[docs]
def spectral_convolution_ptr(wavelength, spectrum, pixel_dv):
"""Convolve a spectrum to instrumental resolution using petitRADTRANS."""
# Calculate resolution from pixel velocity width
resolution = ConstantVariables.CLIGHT / (2 * pixel_dv)
return SpectralModel.convolve(wavelength, spectrum, resolution)