"""
High-resolution likelihood calculations for atmospheric retrieval.
Manages mass fraction computation, prior determinant calculation, and
the full high-resolution spectral likelihood evaluation. Extracted from
ModelData to isolate the HR likelihood logic.
"""
import os
import pickle
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from petitRADTRANS import physical_constants as phys_const
import pyratbay.atmosphere as pa
from scipy.integrate import trapezoid
from scipy.interpolate import splrep, splev
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.Retrieval.ModelCalculation.ParameterHandler import get_param_array_initial_value
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
convolve_resolution, trpca, compute_adjusted_abundance,
convolve_solid_body_rotation, kernel_solid_body_rotation,
rv_planet_and_star, estimate_continuum,
)
from GUIBRUSHR.core.types import slice_section
def _ctx_suffix(chain, step):
"""Return ``" [chain X, step Y]"`` when both are set, else ``""``.
Used to annotate per-step diagnostic prints with their MCMC origin so
log triage can point to the offending chain. Returns the empty string
when called outside an MCMC context (e.g., MAP optimizer, standalone).
"""
if chain is None or step is None:
return ""
return f" [chain {chain}, step {step}]"
[docs]
class LikelihoodHR:
"""
High-resolution likelihood evaluator for atmospheric retrieval.
Handles mass fraction computation, Bayesian prior determinant
calculation, and the full high-resolution spectral likelihood
(PCA telluric removal, chi-square, diagnostic plots).
Attributes:
atmosphere: Atmosphere object with species, pressure, resolution data.
retrieval_data: Retrieval configuration (paths, methods, masks).
param_handler: ParameterHandler for index/value lookups.
bestpars_data: Bestpars object with list_bestpars and nfit.
beta_list: Array of beta values per night (or 1).
model_type: Model type string ("Retrieval", "Model", etc.).
clight: Speed of light constant.
start_molec: Index where molecules start in params_list.
start_elements: Index where elements start in params_list.
start_condensed: Index where condensed molecules start in params_list.
"""
[docs]
def __init__(
self,
atmosphere,
retrieval_data,
param_handler,
bestpars_data,
# jitter_list,
beta_list,
model_type,
clight,
start_molec,
start_elements,
start_condensed,
):
self.atmosphere = atmosphere
self.retrieval_data = retrieval_data
self.param_handler = param_handler
self.bestpars_data = bestpars_data
# self.jitter_list = jitter_list
self.beta_list = beta_list
self.model_type = model_type
self.clight = clight
self.start_molec = start_molec
self.start_elements = start_elements
self.start_condensed = start_condensed
# Reason of the most recent high_resolution_lhood failure (None on
# success). Read by ModelData.lh_function_gib to report a distinct
# reject reason: "hr_model_nan", "hr_negative_spectrum", "hr_trpca_failed".
self.failure_reason = None
# ------------------------------------------------------------------
# Convenience accessors delegating to param_handler
# ------------------------------------------------------------------
@property
def initial_param_array(self):
return self.param_handler.initial_param_array
[docs]
def get_value(self, params, name, single_value=True):
return self.param_handler.get_value(params, name, single_value)
# ------------------------------------------------------------------
# Mass fraction calculation
# ------------------------------------------------------------------
# noinspection PyUnresolvedReferences
[docs]
def calculate_mass_fraction(
self,
temperature,
metallicity,
c_o_ratio,
si_o_ratio_final,
vmr_peak_arr,
pressure_peak_arr,
width_peak_arr,
chain=None,
step=None,
):
"""
Compute mass mixing ratios for atmospheric species.
This method calculates mass mixing ratios based on the selected chemistry
model (equilibrium, free chemistry, or hybrid). It handles volume mixing
ratio calculations, mean molecular weight computation, and species abundance
adjustments for dissociation effects.
Args:
temperature: 1D temperature profile array
metallicity: log(metallicity) enhancement factor for equilibrium chemistry
c_o_ratio: Carbon-to-oxygen ratio for equilibrium chemistry
si_o_ratio_final: Si-to-oxygen ratio for equilibrium chemistry
vmr_peak_arr: VMR peak values for variable abundance profiles
pressure_peak_arr: Pressure peak positions for variable profiles
width_peak_arr: Width parameters for variable abundance profiles
chain: MCMC chain index, used only to annotate warning prints
with their origin. None outside an MCMC context.
step: MCMC outer step index, used together with ``chain`` to
annotate warning prints. None outside an MCMC context.
Returns:
Tuple containing:
- mass_fraction: Dictionary of mass mixing ratio profiles
- MMW: 1D mean molecular weight profile array
- vmr: 2D volume mixing ratio array of shape [nlayers, nmol]
- mean_VMR_and_MF_string: String containing mean VMR and mass fractions per species
- mean_VMR_and_MF_dict: Dictionary mapping species to their mean VMR, MF, and linelist name
- err: Boolean, True if negative mass fractions were found
"""
# Get parameters and molecular masses
params = self.initial_param_array
masses = self.atmosphere.species_obj.mass_vector
# Calculate VMR based on chemistry model. Slice by NAME anchor so
# the call works for both the legacy positional list and a
# ParamRegistry container, and so any future reorder of
# ConstantVariables.params_list cannot silently shift the section.
params_list = ConstantVariables.params_list
array_molecules_full = np.array(slice_section(
params, params_list,
ConstantVariables.FIRST_MOLEC_NAME,
ConstantVariables.FIRST_ELEMENT_NAME,
))
array_condensed_molecules = np.array(slice_section(
params, params_list,
ConstantVariables.FIRST_CONDENSED_NAME,
))
# H2/He are at the start of the molecules section; skip them by
# taking a Python slice on the already-named molecule array (avoids
# carrying integer offsets across the boundary).
skip = self.atmosphere.species_obj.skip_molec_param_array
array_molecules_excluded = array_molecules_full[skip:]
e_ratio = {}
if c_o_ratio is not None:
e_ratio["C_O"] = c_o_ratio
if si_o_ratio_final is not None:
e_ratio["Si_O"] = si_o_ratio_final
# EQUILIBRIUM CHEMISTRY: Use thermochemical equilibrium
if self.atmosphere.chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[0]: # Equilibrium
species = list(self.atmosphere.species_compatible_with_prt)
vmr = self.atmosphere.chemcat_obj.thermochemical_equilibrium(
temperature=temperature,
metallicity=metallicity,
e_ratio=e_ratio,
)
elif self.atmosphere.chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[1]:
pressure = self.atmosphere.pressure_data.pressures
nlayers = len(pressure)
species = list(self.atmosphere.species_obj.composition)
nmol = len(species)
vmr = np.zeros((nlayers, nmol))
# To correct indices of vmr_peak, pressure_peak, width_peak, a4
i_peak = 0
index_metals = []
# He and H2 are NOT included (+2 reason)
for mol in array_molecules_excluded:
if mol is None:
continue
imol = species.index(mol.name)
if mol.molec_formula not in ConstantVariables.NOT_VMR_METALS:
index_metals.append(imol)
vmr_value = 10 ** self.get_value(params, mol.molec_formula)
if mol.is_VMR_costant:
vmr[:, imol] = vmr_value
else:
vmr_peak_par = vmr_peak_arr[i_peak]
pressure_peak_par = pressure_peak_arr[i_peak]
width_peak_par = width_peak_arr[i_peak]
vmr_pf_log = np.log10(vmr_value) - vmr_peak_par * np.exp(
-np.power(
(np.log10(pressure) - pressure_peak_par)
/ width_peak_par,
2,
)
)
vmr_pf_log[np.where(vmr_pf_log > -1)] = -1
vmr_pf = 10 ** vmr_pf_log
i_peak += 1
vmr[:, imol] = vmr_pf
# condensed
for condensed_mol in array_condensed_molecules:
if condensed_mol is None:
continue
imol = species.index(condensed_mol.name)
index_metals.append(imol)
vmr_value = 10 ** self.get_value(params, condensed_mol.name)
vmr[:, imol] = vmr_value
vmr_metals = np.sum(vmr[:, index_metals], axis=1)
if self.atmosphere.species_obj.include_h_m:
# Compute the H2/He ratio from the provided parameters
H2_He_ratio = self.get_value(params, "H2") / self.get_value(params, "He")
# Get indices for the relevant species in the vmr array
i_H2 = species.index('H2')
i_He = species.index('He')
i_H = species.index('H')
i_em = species.index('e-')
i_hm = species.index('H-')
vmr[:, i_He] = (1.0 - vmr_metals) / (1.0 + H2_He_ratio)
# --- H2 Abundance ---
coeff_h2 = [1, 2.41e04, 6.5]
base_vmr_H2 = H2_He_ratio * vmr[-1, i_He]
A_H2 = compute_adjusted_abundance(pressure, temperature, coeff_h2, base_vmr_H2)
vmr[:, i_H2] = A_H2
vmr[:, i_H] = A_H2[::-1]
# --- Free Electrons (e-) ---
coeff_em = [-0.4, 0, 0]
base_vmr_em = vmr[-1, i_em]
A_em = compute_adjusted_abundance(pressure, temperature, coeff_em, base_vmr_em)
vmr[:, i_em] = A_em
# --- Negative Hydrogen Ions (H-) ---
coeff_hm = [0.6, -0.14e4, 7.7]
base_vmr_hm = vmr[-1, i_hm]
A_hm = compute_adjusted_abundance(pressure, temperature, coeff_hm, base_vmr_hm)
vmr[:, i_hm] = A_hm
# --- Water Vapor (H2O) ---
if 'H2O' in species:
i_h2o = species.index('H2O')
coeff_h2o = [2, 4.83e4, 15.9]
base_vmr_h2o = vmr[-1, i_h2o]
A_h2o = compute_adjusted_abundance(pressure, temperature, coeff_h2o, base_vmr_h2o)
vmr[:, i_h2o] = A_h2o
else:
H2_He_ratio = self.get_value(params, "H2") / self.get_value(params, "He")
i_H2 = species.index('H2')
i_He = species.index('He')
vmr[:, i_He] = (1.0 - vmr_metals) / (1.0 + H2_He_ratio)
vmr[:, i_H2] = H2_He_ratio * vmr[:, i_He]
else:
species = list(self.atmosphere.species_compatible_with_prt)
e_abundances = dict()
element_slots = slice_section(
params, ConstantVariables.params_list,
ConstantVariables.FIRST_ELEMENT_NAME,
)
for element in element_slots:
if element is None:
continue
# Extract bare element symbol (e.g. "C/H" -> "C", "Fe/H" -> "Fe")
elem_symbol = element.name.split("/")[0]
e_abundances[elem_symbol] = self.get_value(params, element.name)
# In hybrid chemistry, e_abundances drive C and O directly;
# C/O ratio is never passed — drop it, keep other ratios (e.g. Si_O).
hybrid_e_ratio = dict(e_ratio)
hybrid_e_ratio.pop("C_O", None)
vmr = self.atmosphere.chemcat_obj.thermochemical_equilibrium(
temperature=temperature,
metallicity=metallicity,
e_abundances=e_abundances,
# e_ratio=hybrid_e_ratio,
)
# Mean molecular weight
MMW = pa.mean_weight(vmr, species, mass=masses)
# Mass-mixing ratio abundances:
mass_fractions_PTR = {}
mean_VMR_and_MF_dict = {}
mean_VMR_and_MF_string = ""
err = False
for mol in array_molecules_full:
if mol is not None:
name = mol.name_for_list_molec
i = species.index(mol.molec_formula)
mass_fractions_PTR[name] = vmr[:, i] * masses[i] / MMW
mean_VMR_and_MF_dict[mol.molec_formula] = {
"VMR": np.mean(vmr[:, i]),
"MF": np.mean(mass_fractions_PTR[name]),
"LineList": name
}
mean_VMR_and_MF_string += f"{mol.molec_formula}: Mean VMR = {np.mean(vmr[:, i])}:.7e; Mean MF {np.mean(mass_fractions_PTR[name]):.7e}\n"
if np.any(mass_fractions_PTR[name] < 0):
print(f"Warning{_ctx_suffix(chain, step)}: negative mass fraction values found for {name}")
err = True
for condensed_mol in array_condensed_molecules:
if condensed_mol is not None:
name = condensed_mol.name_for_list_molec
i = species.index(name)
mass_fractions_PTR[name] = vmr[:, i] * masses[i] / MMW
mean_VMR_and_MF_dict[condensed_mol.molec_formula] = {
"VMR": np.mean(vmr[:, i]),
"MF": np.mean(mass_fractions_PTR[name]),
"LineList": name
}
mean_VMR_and_MF_string += f"{condensed_mol.molec_formula}: Mean VMR = {np.mean(vmr[:, i])}:.7e; Mean MF {np.mean(mass_fractions_PTR[name]):.7e}\n"
if np.any(mass_fractions_PTR[name] < 0):
print(f"Warning{_ctx_suffix(chain, step)}: negative mass fraction values found for {name}")
err = True
return mass_fractions_PTR, MMW, vmr, mean_VMR_and_MF_string, mean_VMR_and_MF_dict, err
# ------------------------------------------------------------------
# Prior determinant
# ------------------------------------------------------------------
[docs]
def calculate_log_prior(self, params):
"""
Calculate the log of the Gaussian prior product for Bayesian parameter estimation.
IDL called this function/value 'determinant' (keyword: determinant=det), but it
computes π(θ) = Π_i exp[-0.5*((θ_i-μ_i)/σ_i)²] — the product of Gaussian priors.
This is not a matrix determinant; the name is a historical misnomer from IDL exofast.
IDL computed the product in linear space::
det = 1.0
det *= exp(-((value - initial_value)**2) / (2*sigma**2))
which underflows to 0.0 for N ≳ 150 parameters with deviations of a few sigma,
silently blocking all MCMC steps (acceptance ratio C = 0/0 → nan → never accepted).
This implementation computes the equivalent in log-space (sum instead of product),
which is numerically exact for any number of parameters::
log_prior = Σ -0.5 * ((θ_i - μ_i) / σ_i)²
Args:
params: Array of parameter objects
Returns:
log_prior: Log of the Gaussian prior product (float, always finite)
"""
# IDL: det = 1.0 → log-space equivalent: log_prior = 0.0 (log(1) = 0)
log_prior = 0.0
for iel, param in enumerate(params):
# Skip parameters without priors
if param is None or param.sigma_prior is None:
continue
# Determine if parameter has single or multiple values
single_value = param.name not in self.list_multiple_param
value = self.get_value(params, param.name, single_value)
# Get initial parameter value for comparison
initial_value = get_param_array_initial_value(
self.initial_param_array, iel, single_value
)
sigma = params[iel].get_sigma_prior()
# IDL: det *= exp(-((value - initial_value)**2) / (2*sigma**2))
# Log-space equivalent — avoids underflow, mathematically identical:
if single_value:
log_prior += -((value - initial_value) ** 2) / (2 * sigma ** 2)
else:
# Handle array parameters (e.g., jitter, f_rot, beta)
for index_arr in range(len(value)):
log_prior += -(
(value[index_arr] - initial_value[index_arr]) ** 2
/ (2 * sigma[index_arr] ** 2)
)
return log_prior
# ------------------------------------------------------------------
# TRPCA error debug info
# ------------------------------------------------------------------
def _save_trpca_error_debug_info(
self,
lcrm_mask,
lcrm_mask_nomask,
good_pixel,
night,
h,
dict_calc_model,
temperature,
mass_fraction,
vmr,
MMW,
wl_full_resolution_HR,
model_full_resolution_HR,
depth_full_resolution_HR,
):
"""
Save detailed debug information when TRPCA fails.
Args:
lcrm_mask: Masked telluric-removed model spectrum
lcrm_mask_nomask: Unmasked telluric-removed model spectrum
good_pixel: Boolean array of valid pixels
night: Night object containing observation data
h: Current spectral order index
dict_calc_model: Dictionary with calculated model parameters
temperature: Temperature profile array
mass_fraction: Mass fraction dictionary
vmr: Volume mixing ratio array
MMW: Mean molecular weight array
wl_full_resolution_HR: High-resolution wavelength array
model_full_resolution_HR: High-resolution model spectrum
depth_full_resolution_HR: High-resolution transit depth spectrum
"""
error_trpca_dictionary = {
"lcrm_mask": lcrm_mask,
"lcrm_mask_nomask": lcrm_mask_nomask,
"good_pixel": good_pixel,
"nfc": night.nfc[h],
"tell_rm_method": self.retrieval_data.tell_rm_method,
"smooth_on": night.smooth_on,
"smooth_size": night.smooth_size,
"model_reprocessing": self.retrieval_data.model_reprocessing,
"subtraction_of_the_average_spectrum": night.subtraction_of_the_average_spectrum,
"dict_calc_model": dict_calc_model,
"temperature": temperature,
"mass_fraction": mass_fraction,
"vmr": vmr,
"MMW": MMW,
"LIST_PT_PROFILE_TABLE": ConstantVariables.LIST_PT_PROFILE_TABLE,
"instruments_HR": self.atmosphere.resolution_obj.instruments_HR,
"retrieval_data": self.retrieval_data,
"wl_full_resolution_HR": wl_full_resolution_HR,
"model_full_resolution_HR": model_full_resolution_HR,
"depth_full_resolution_HR": depth_full_resolution_HR,
"stellar_spectrum": self.atmosphere.stellar_spectrum,
"stellar_spline_model_HR": self.atmosphere.stellar_spline_model_hr,
"atmospherehr": {
"pressures": self.atmosphere.pressure_data.pressures,
"line_species_complete_name_hr": self.atmosphere.species_obj.line_species_complete_name_hr,
"rayleigh_species": self.atmosphere.species_obj.rayleigh_species,
"gas_continuum_contributors": self.atmosphere.species_obj.continum_opacity,
"wavelength_boundaries": np.array([
self.atmosphere.wavelength.min_wl_hr,
self.atmosphere.wavelength.max_wl_hr
]),
"line_opacity_mode": "lbl",
"line_by_line_opacity_sampling": self.atmosphere.resolution_obj.lbl_high_res
},
"SOLAR_TO_JUPITER_MASSES": ConstantVariables.SOLAR_TO_JUPITER_MASSES,
"clight": ConstantVariables.CLIGHT,
"stellar_spectrum_type_hr": self.atmosphere.stellar_spectrum_type_hr,
"beta_list": self.beta_list
# "jitter_list": self.jitter_list
}
error_file_path = f"{self.retrieval_data.path_results}/error_trpca.pkl"
with open(error_file_path, "ab") as fo:
pickle.dump(error_trpca_dictionary, fo)
print(
f"TRPCA FAILED, check file: {self.retrieval_data.path_results}"
"/error_trpca.pkl\n"
)
# ------------------------------------------------------------------
# High-resolution likelihood
# ------------------------------------------------------------------
# fmt: off
[docs]
def high_resolution_lhood(
self,
temperature,
mass_fraction,
vmr,
MMW,
dict_calc_model,
chain=None,
step=None,
):
"""
Compute high-resolution likelihood with advanced spectral processing.
This method performs the most complex spectral processing in the atmospheric
retrieval system. It handles:
1. Full-resolution atmospheric model calculation
2. Instrumental resolution convolution
3. Stellar rotation broadening (solid body rotation)
4. Orbital dynamics and Doppler shifts
5. Telluric removal via Principal Component Analysis (PCA)
6. Chi-square likelihood evaluation against observations
The method preserves exact mathematical operations critical for
high-precision atmospheric spectroscopy and exoplanet detection.
Args:
temperature: Atmospheric temperature profile
mass_fraction: Mass fraction profiles for all species
vmr: Volume mixing ratios for all species
MMW: Mean molecular weight profile
dict_calc_model: Dictionary containing all model parameters
chain: MCMC chain index, used only to annotate diagnostic
prints (negative spectrum, TRPCA error). None outside an
MCMC context.
step: MCMC outer step index, used together with ``chain`` to
annotate diagnostic prints. None outside an MCMC context.
Returns:
Tuple containing:
- status: True if successful, False if errors occurred
- lhood: Log-likelihood value
- wl_full_resolution_HR: Full resolution wavelength array
- depth_full_resolution_HR: Full resolution spectrum
- opacity_contribution_HR: Opacity contributions (for plotting)
- lh_high_resolution: Dictionary of per-instrument per-night likelihood details
"""
# EXTRACT PARAMETERS FROM MODEL CALCULATION DICTIONARY
rp = dict_calc_model["rp"] # Planet radius
gravity = dict_calc_model["gravity"] # Surface gravity
omegad = dict_calc_model["omegad"] # Rotation velocity (log scale)
rv = dict_calc_model["rv"] # Systemic radial velocity
sf = dict_calc_model["sf"] # Single scaling factor
sf_arr = dict_calc_model["sf_arr"] # Multiple scaling factors
f_rot_arr = dict_calc_model["f_rot_arr"] # Rotation factors per night
T0 = dict_calc_model["T0"] # Reference temperature
T_low = dict_calc_model["T_low"] # Low pressure temperature
T3_node = dict_calc_model["T3_node"] # Temperature node 3
ecc = dict_calc_model["ecc"] # Orbital eccentricity
opi = dict_calc_model["opi"] # Argument of periastron
kp = dict_calc_model["kp"] # Planet velocity semi-amplitude
# jitter_arr = dict_calc_model["jitter_arr"] # Noise jitter per night
beta_arr = dict_calc_model["beta_arr"] # Noise beta per night
dVsys_arr = dict_calc_model["dVsys_arr"] # Systemic velocity variations
# INITIALIZE VARIABLES
lhood = 0 # Total log-likelihood
csscaled = None # Spline for convolved spectrum
opacity_contribution_HR = None
lh_high_resolution = {}
self.failure_reason = None # cleared on every evaluation
# CALCULATE FULL-RESOLUTION ATMOSPHERIC MODEL
# This is the core atmospheric model calculation producing the theoretical spectrum
wl_full_resolution_HR, model_full_resolution_HR, depth_full_resolution_HR = self.atmosphere.calc_model(
temperature, mass_fraction, vmr, MMW, dict_calc_model, True,
)
# MODEL FAILURE GUARD: stop here if the HR model spectrum is NaN/Inf.
# This is a model-evaluation failure (e.g. petitRADTRANS returning NaN for
# an out-of-range temperature profile), upstream of the PCA/trpca step.
# Returning now makes the cause explicit instead of surfacing later as the
# misleading "Input X contains NaN" PCA error.
if (not np.all(np.isfinite(model_full_resolution_HR))
or not np.all(np.isfinite(depth_full_resolution_HR))):
self.failure_reason = "hr_model_nan"
print(f"High-resolution model failed: NaN/Inf in the model spectrum "
f"(model evaluation, not PCA){_ctx_suffix(chain, step)}.")
return False, -np.inf, wl_full_resolution_HR, depth_full_resolution_HR, opacity_contribution_HR, lh_high_resolution
# SETUP PLOTTING FOR MODEL ANALYSIS (if needed)
pdf = None
if "Model" in self.model_type:
os.system(f"mkdir -p {self.retrieval_data.path_results}/model_high_low_res/")
pdf = PdfPages(
f"{self.retrieval_data.path_results}/model_high_low_res/model_reprocessing.pdf"
)
if "Opacity" in self.model_type:
# Calculate opacity contributions for analysis
opacity_contribution_HR = self.atmosphere.opacity_contribution(
temperature, mass_fraction, MMW, dict_calc_model, True
)
# PROCESS EACH HIGH-RESOLUTION INSTRUMENT
counter_nights_total = -1 # Global night counter across all instruments
counter_derivative_params = -1
for index_instrument, instrument in enumerate(self.atmosphere.resolution_obj.instruments_HR):
lh_night = {}
# INSTRUMENTAL RESOLUTION CONVOLUTION (if no stellar rotation)
if omegad is None:
# Convolve theoretical spectrum with instrument resolution function.
# The kernel is NOT shifted by rv: the radial velocity (including the
# retrieval rv) is applied as a single term on the data wavelength
# grid below, via rv_dynamic in rv_planet_and_star.
wl_convolved_resolution, model_convolved_resolution = convolve_resolution(
wl_full_resolution_HR, model_full_resolution_HR, instrument.hwhm_km_s,
self.atmosphere.rv_sampling
)
# Create spline interpolation for fast evaluation at arbitrary wavelengths
csscaled = splrep(wl_convolved_resolution, model_convolved_resolution)
# PROCESS EACH OBSERVATION NIGHT
for index_night_in_current_instrument, night in enumerate(instrument.night_arr):
counter_nights_total += 1
if index_night_in_current_instrument != 0:
counter_derivative_params += 1
# DETERMINE SCALING FACTOR FOR THIS NIGHT
# sf and sf_multi arrive already in linear units (Log10ToLinear applied
# by to_linear at extraction time); the neutral default is 1 (= 10**0).
scale_HR = 1
if sf is not None:
scale_HR = sf
elif sf_arr is not None:
scale_HR = sf_arr[counter_nights_total]
# STELLAR ROTATION BROADENING (if included)
# omega arrives already in linear units (Log10ToLinear applied by to_linear).
if omegad is not None:
# Determine reference temperature for rotation kernel calculation
# Different temperature profile formats use different reference temperatures
if self.retrieval_data.format_temperature == ConstantVariables.LIST_PT_PROFILE_TABLE[3]:
T0 = T_low
elif self.retrieval_data.format_temperature == ConstantVariables.LIST_PT_PROFILE_TABLE[4]:
T0 = T3_node
# Calculate solid body rotation kernel
ker_rot = kernel_solid_body_rotation(
omegad, f_rot_arr, T0, np.mean(MMW), gravity, rp,
self.atmosphere.rad_mode, counter_nights_total, self.atmosphere.rv_sampling
)
# Apply rotation broadening to the spectrum
wl_convolved_rot, model_convolved_rot = convolve_solid_body_rotation(
wl_full_resolution_HR, model_full_resolution_HR, ker_rot, self.atmosphere.rv_sampling
)
# Apply instrumental resolution convolution after rotation.
# As above, the kernel is not shifted by rv: the radial velocity
# is applied through vtot (rv_dynamic) on the data grid below.
wl_convolved_resolution_after_rot, model_convolved_resolution_after_rot = convolve_resolution(
wl_convolved_rot, model_convolved_rot, instrument.hwhm_km_s, self.atmosphere.rv_sampling
)
# Create spline for the rotation+resolution convolved spectrum
csscaled = splrep(wl_convolved_resolution_after_rot, model_convolved_resolution_after_rot)
# ORBITAL DYNAMICS AND RADIAL VELOCITY CALCULATION
# Single radial-velocity computation: the retrieval rv is passed in
# as rv_dynamic and folded into vtot here (it is no longer applied as
# a convolution-kernel shift). With coeff_vtot=1 (Retrieval) rv_dynamic
# enters subtracted, so vtot = vtot_without_rv - rv, which reproduces
# the previous "+rv" kernel shift exactly.
vtot, vtot_star, _ = rv_planet_and_star(
self.retrieval_data.eccentricity, self.atmosphere.target, ecc, opi, kp, night, "Retrieval",
index_night_in_current_instrument, dVsys_arr, counter_derivative_params,
rv_dynamic=rv
)
lh_current_night = 0
# PROCESS EACH SPECTRAL ORDER
for h in range(night.n_good_orders):
# Extract observational data for this order
wavelengths_night = night.lambdas[:, h, :] # Wavelength array
spectrum_only_planet_night = night.spectra[:, h, :] # Observed spectra
error_spectrum_only_planet_night = night.sigma_spectra_lin[:, h, :] # Linear-space error bars
good_pixels = night.maskinvm[:, h, :] == 0 # Valid pixel mask
# APPLY DOPPLER SHIFTS TO WAVELENGTH GRID
# Reshape wavelength array to match good pixels structure
wl_data_masked = wavelengths_night[np.where(good_pixels)].reshape(
-1, np.shape(good_pixels)[1]
)
# Calculate Doppler shifts for planet and star
dl_planet = wl_data_masked * vtot / self.clight # Planet velocity shift
wlen_shifted_planet = dl_planet + wl_data_masked # Planet rest frame wavelengths
dl_star = wl_data_masked * vtot_star / self.clight # Stellar velocity shift
wlen_shifted_star = dl_star + wl_data_masked # Stellar rest frame wavelengths
Fscaled = None
mol_modf = None
stellar_spectrum_HR = None
# CALCULATE OBSERVED SPECTRUM BASED ON OBSERVATION MODE
if self.atmosphere.rad_mode == 'Transmission':
# TRANSMISSION SPECTROSCOPY
# Evaluate atmospheric model at Doppler-shifted wavelengths
model_eval = splev(wlen_shifted_planet, csscaled, der=0)
# Convert from altitude to transit depth
mol_mod = model_eval / phys_const.r_jup_mean
mol_modf = (mol_mod / (
self.atmosphere.target.stellar_radius * ConstantVariables.RATIO_RSUN_RJUP_MEAN)) ** 2
spectrum = 1 - mol_modf # Transit depth spectrum
else:
# EMISSION SPECTROSCOPY
# Calculate stellar spectrum contribution
stellar_spectrum_HR = splev(wlen_shifted_star * 1e-7, self.atmosphere.stellar_spline_model_hr)
# Evaluate planetary emission model
model_eval = splev(wlen_shifted_planet, csscaled, der=0)
# Calculate planet-to-star flux ratio
# Stellar spc is already multiplied by np.pi to go from erg/cm^2/cm/s/sr to erg/cm^2/cm/s
Fscaled = (model_eval / stellar_spectrum_HR) * (
rp / (self.atmosphere.target.stellar_radius * phys_const.r_sun)) ** 2
spectrum = 1 + Fscaled # Combined stellar + planetary flux
if np.sum(spectrum < 0) > 0:
if "Model" in self.model_type:
os.system(f"mkdir -p {self.retrieval_data.path_results}/error_spectrum/")
with open(f"{self.retrieval_data.path_results}/error_spectrum/spectrum.pkl",
"wb") as pkl_spectrum:
error_spc_pkl = {
"spectrum": spectrum,
"Fscaled": Fscaled,
"mol_modf": mol_modf,
"stellar_spectrum_HR": stellar_spectrum_HR
}
pickle.dump(error_spc_pkl, pkl_spectrum)
print(
f"Negative values in spectrum for order {h} of night "
f"{index_night_in_current_instrument}{_ctx_suffix(chain, step)}.")
self.failure_reason = "hr_negative_spectrum"
return False, -np.inf, wl_full_resolution_HR, depth_full_resolution_HR, opacity_contribution_HR, lh_high_resolution
# PREPARE DATA FOR TELLURIC REMOVAL VIA PCA
model_in_log = np.log10(spectrum) # Convert to log scale for PCA processing
good_pixel = np.where(good_pixels[:, -1])
good_pixel = good_pixel[0]
# SETUP PCA PROCESSING BASED ON METHOD # TODO can be moved in night_extraction
lcrm_mask_nomask = np.zeros_like(night.array_reconstructed_from_PCA_components[:, h, :])
if self.retrieval_data.model_reprocessing == "hard":
# Hard PCA: Use existing telluric removal data
lcrm_mask = night.array_reconstructed_from_PCA_components[good_pixel, h, :]
if night.subtraction_of_the_average_spectrum:
lcrm_mask = lcrm_mask + night.average_spectrum_to_be_added[h]
else: # "soft" method
# Soft PCA: Simplified approach with minimal telluric removal
lcrm_mask = np.zeros_like(night.array_reconstructed_from_PCA_components[good_pixel, h, :])
# INJECT ATMOSPHERIC MODEL INTO IN-TRANSIT DATA
# Add theoretical atmospheric signal to in-transit observations
lcrm_mask[:, night.intransit] = (
model_in_log + (lcrm_mask[:, night.intransit])
)
# EXECUTE TELLURIC REMOVAL PCA
# Apply Time-Resolved Principal Component Analysis (TRPCA)
# This removes telluric contamination while preserving the atmospheric signal
lcrm_pca, error_trpca = trpca(
lcrm_mask,
lcrm_mask_nomask,
good_pixel,
night.nfc[h], # Number of components
night.smooth_on, # Smoothing flag
night.smooth_size, # Smoothing size
night.apply_standardize,
self.retrieval_data.model_reprocessing, # Processing method
night.subtraction_of_the_average_spectrum,
night.pca_mode # PCA mode: spatial or temporal
)
if error_trpca:
if "Model" in self.model_type:
print(f"Error in TRPCA processing{_ctx_suffix(chain, step)}.")
# Handle PCA failure with detailed error logging
# self._save_trpca_error_debug_info(
# lcrm_mask, lcrm_mask_nomask, good_pixel, night, h,
# dict_calc_model, temperature, mass_fraction, vmr, MMW,
# wl_full_resolution_HR, model_full_resolution_HR,
# depth_full_resolution_HR,
# )
self.failure_reason = "hr_trpca_failed"
return False, -np.inf, wl_full_resolution_HR, depth_full_resolution_HR, opacity_contribution_HR, lh_high_resolution
# CALCULATE CHI-SQUARE LIKELIHOOD
# Extract only in-transit data for likelihood calculation
lcrm_pca = lcrm_pca[:, night.intransit]
# Beta noise scale factor for this night, if the beta
# parameter is configured (beta_arr is None otherwise).
# beta_arr originates from beta_HR (legacy bare "beta" rows are
# remapped to beta_HR at read time); see ModelData.lh_function_gib.
if beta_arr is not None and np.size(beta_arr) > 0:
beta_factor = beta_arr[counter_nights_total]
else:
beta_factor = 1
# # Add jitter noise to error bars if specified
# if self.jitter_list is not None and self.jitter_list.size > 0:
# errplusjit = np.sqrt(
# error_spectrum_only_planet_night[good_pixel, :] ** 2 + jitter_arr[counter_nights_total] ** 2
# )
# else:
# errplusjit = error_spectrum_only_planet_night[good_pixel, :]
final_error = error_spectrum_only_planet_night[good_pixel, :]
n_masked_points = 0
if (
self.retrieval_data.mask_phase_enabled and
self.retrieval_data.mask_night_index[counter_nights_total]
):
if self.retrieval_data.mask_inside_range:
mask_phases_indices_current_night = np.where(np.logical_and(
night.phase_new_range_considered >= self.retrieval_data.min_phase,
night.phase_new_range_considered <= self.retrieval_data.max_phase
))[0]
else:
mask_phases_indices_current_night = np.where(np.logical_or(
night.phase_new_range_considered <= self.retrieval_data.min_phase,
night.phase_new_range_considered >= self.retrieval_data.max_phase
))[0]
spectrum_only_planet_night[:, mask_phases_indices_current_night] = 0
lcrm_pca[:, mask_phases_indices_current_night] = 0
final_error[:, mask_phases_indices_current_night] = 1
# errplusjit[:, mask_phases_indices_current_night] = 1
# Masked points contribute 0 to chi2 and to sum(ln sigma),
# so they must not be counted in N either
n_masked_points = len(good_pixel) * len(mask_phases_indices_current_night)
# Calculate chi-square for this spectral order
chiquadro_HR = np.sum(
(
(
# Linear-space chi2: data and model exponentiated back from log.
# scale_HR multiplies the line CONTRAST (10**lcrm_pca - 1). A
# negative model for large scale_HR is benign (no log of model).
10 ** spectrum_only_planet_night[good_pixel, :] # observed data -> linear
- (1 + (10 ** lcrm_pca - 1) * scale_HR) # linear model: 1 + contrast*scale
)
/ (beta_factor * final_error) # linear-space sigma
) ** 2
)
n_good_pixels = len(good_pixel)
n_spectra = spectrum_only_planet_night.shape[1] # or the temporal dimension
# Total number of data points
dof = n_good_pixels * n_spectra
chi2_reduced = chiquadro_HR / dof
# Add to total log-likelihood, Gibson et al. 2021:
# ln L = -(N/2) ln 2pi - N ln beta - sum(ln sigma_i) - chi2 / 2
# chi2 above is already computed with beta-scaled errors;
# masked phase points are excluded from N
n_data = dof - n_masked_points
lh_current_order = (
-0.5 * n_data * np.log(2.0 * np.pi)
- n_data * np.log(beta_factor)
- np.sum(np.log(final_error))
- 0.5 * chiquadro_HR
)
lh_current_night += lh_current_order
lh_night[
night.date] = f"{chiquadro_HR};{lh_current_night};{chi2_reduced};{dof};{self.bestpars_data.nfit}"
lhood += lh_current_order
# GENERATE DIAGNOSTIC PLOTS (for model analysis mode)
if "Model" in self.model_type:
# Create two-panel plot showing model injection and PCA result
fig, axs = plt.subplots(2, 1, figsize=(12, 5))
axs[0].set_title(
"Night: "
+ instrument.nights[index_night_in_current_instrument]
+ " Order: "
+ str(h)
)
# Prepare phase array for plotting (center around transit)
phforplot = night.phases_considered
phforplot[phforplot > 0.5] -= 1
# TOP PANEL: Show injected atmospheric model
model_in_logforplot = np.zeros(
np.shape(night.array_reconstructed_from_PCA_components[:, h, night.intransit])
)
# Fill good pixels with model signal, bad pixels with mean
model_in_logforplot[np.where(good_pixels[:, 0] == 1)] = model_in_log
model_in_logforplot[np.where(good_pixels[:, 0] == 0)] = np.mean(model_in_log)
# Plot model as 2D image (wavelength vs. orbital phase)
result_im = axs[0].imshow(
model_in_logforplot.T[:, 500:701],
aspect="auto",
origin="lower",
interpolation="None",
extent=[
wavelengths_night[500, 0],
wavelengths_night[701, 0],
phforplot[0],
phforplot[-1],
],
)
_ = fig.colorbar(result_im, ax=[axs[0]])
# BOTTOM PANEL: Show PCA-processed result
lcrm_pcaforplot = np.zeros(
np.shape(night.array_reconstructed_from_PCA_components[:, h, night.intransit])
)
# Fill good pixels with PCA result, bad pixels with mean
lcrm_pcaforplot[np.where(good_pixels[:, 0] == 1)] = lcrm_pca
lcrm_pcaforplot[np.where(good_pixels[:, 0] == 0)] = np.mean(lcrm_pca)
# Plot PCA result as 2D image
result_im2 = axs[1].imshow(
lcrm_pcaforplot.T[:, 500:701],
aspect="auto",
origin="lower",
interpolation="None",
extent=[
wavelengths_night[500, 0],
wavelengths_night[701, 0],
phforplot[0],
phforplot[-1],
],
)
_ = fig.colorbar(result_im2, ax=[axs[1]])
# Save plot and clean up
pdf.savefig(fig)
fig.clear()
plt.close(fig)
lh_high_resolution[instrument.name] = lh_night
# RETURN RESULTS
return True, lhood, wl_full_resolution_HR, depth_full_resolution_HR, opacity_contribution_HR, lh_high_resolution