"""
NotRetrieval class module for handling atmospheric retrieval calculations.
This module contains the NotRetrieval class which processes atmospheric data
including volume mixing ratios (VMR), spectra, and opacity contributions for
both high resolution (HR) and low resolution (LR) data.
"""
import numpy as np
from numpy import mean
from petitRADTRANS.spectral_model import SpectralModel
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
get_line_lists,
get_keys_by_value,
get_condensed_line_list
)
[docs]
class NotRetrieval:
"""
A class for handling atmospheric retrieval calculations and data processing.
This class manages atmospheric data including volume mixing ratios (VMR),
spectral data at different resolutions, and opacity contributions. It processes
molecular species data and creates dictionaries for easy access to VMR values.
"""
[docs]
def __init__(
self,
mass_fraction_names_hr,
mass_fraction_names_lr,
MMW,
vmr,
mean_VMR_and_MF_string,
mean_VMR_and_MF_dict,
wl_full_resolution_HR,
depth_full_resolution_HR,
wlen_mu_contribution_HR,
contribution_HR,
wl_full_resolution_LR,
depth_full_resolution_LR,
wl_binned_LR,
spectrum_binned_LR,
offsetLRarr,
wlen_mu_contribution_LR,
contribution_LR,
chemistry,
params_only_molec,
species,
opacity_contribution_HR,
opacity_contribution_LR,
instruments_LR,
rp,
rad_mode,
stellar_radius,
stellar_spectrum,
lh_HR,
lh_LR,
HR_res_present,
LR_res_present,
use_hr_linelists_for_lr
):
"""
Initialize the NotRetrieval class with atmospheric and spectral data.
Args:
mass_fraction_names_hr: High resolution mass fraction names
mass_fraction_names_lr: Low resolution mass fraction names
MMW: Mean molecular weight data
vmr: Volume mixing ratio data
mean_VMR_and_MF_string: Mean VMR and Molecular Fraction string
mean_VMR_and_MF_dict: Mean VMR and Molecular Fraction dictionary
wl_full_resolution_HR: Full resolution wavelength data
depth_full_resolution_HR: Full resolution spectrum data
wlen_mu_contribution_HR: High resolution wavelength contribution
contribution_HR: High resolution contribution data
wl_full_resolution_LR: Low resolution full wavelength data
depth_full_resolution_LR: Low resolution full spectrum data
wl_binned_LR: Low resolution binned wavelength data
spectrum_binned_LR: Low resolution binned spectrum data
offsetLRarr: Low resolution offset array
wlen_mu_contribution_LR: Low resolution wavelength contribution
contribution_LR: Low resolution contribution data
chemistry: Chemistry table selection
params_only_molec: Molecular parameters
species: List of chemical species
opacity_contribution_HR: High resolution opacity contribution
opacity_contribution_LR: Low resolution opacity contribution
instruments_LR: List of instruments for low resolution data
rp: Planet radius parameter
rad_mode: Radiative transfer mode selection
stellar_radius: Stellar radius parameter
stellar_spectrum: Stellar spectrum for resolution data
lh_HR: Likelihood dictionary for high resolution data
lh_LR: Likelihood dictionary for low resolution data
HR_res_present: Boolean indicating if high resolution data is present
LR_res_present: Boolean indicating if low resolution data is present
use_hr_linelists_for_lr: Use high resolution linelists for low resolution data
"""
# Store all input parameters as instance attributes
self.mass_fraction_names_hr = mass_fraction_names_hr
self.mass_fraction_names_lr = mass_fraction_names_lr
self.MMW = MMW
self.vmr = vmr
self.mean_VMR_and_MF_string = mean_VMR_and_MF_string
self.mean_VMR_and_MF_dict = mean_VMR_and_MF_dict
self.wl_full_resolution_HR = wl_full_resolution_HR
self.depth_full_resolution_HR = depth_full_resolution_HR
self.wlen_mu_contribution_HR = wlen_mu_contribution_HR
self.contribution_HR = contribution_HR
self.wlen_mu_contribution_LR = wlen_mu_contribution_LR
self.contribution_LR = contribution_LR
self.wl_full_resolution_LR = wl_full_resolution_LR
if use_hr_linelists_for_lr:
self.depth_full_resolution_LR = SpectralModel.convolve(
wl_full_resolution_LR, depth_full_resolution_LR, 5000,
)
else:
self.depth_full_resolution_LR = depth_full_resolution_LR
self.wl_binned_LR = wl_binned_LR
self.spectrum_binned_LR = spectrum_binned_LR
self.offsetLRarr = offsetLRarr
self.instruments_LR = instruments_LR
self.rp = rp
self.rad_mode = rad_mode
self.stellar_radius = stellar_radius
self.lh_HR = lh_HR
self.lh_LR = lh_LR
# Initialize VMR-related attributes
self.vmr_dict = dict()
self.vmr_mean = []
self.opacity_contribution_HR = None
self.opacity_contribution_LR = None
# Process molecular species and create VMR dictionary
self._process_molecular_species(chemistry, params_only_molec, species)
if self.rad_mode == "Emission":
# TODO opacities contribution for Emission"
pass
else:
if HR_res_present and opacity_contribution_HR is not None:
self.opacity_contribution_HR = self._process_opacities(opacity_contribution_HR, stellar_spectrum.get("model_PRT_HR", None))
if LR_res_present and opacity_contribution_LR is not None:
self.opacity_contribution_LR = self._process_opacities(opacity_contribution_LR, stellar_spectrum.get("model_PRT_LR", None))
def _process_opacities(self, opacities_contributions, stellar_spectrum):
"""
Process opacity contributions by converting units and applying stellar correction.
Args:
opacities_contributions: Dictionary of opacity contribution data, where values
can be tuples or nested dicts of species-level contributions
stellar_spectrum: Stellar spectrum array used for emission mode correction
Returns:
The processed opacity contributions dictionary with converted units,
or None if input is None.
"""
if opacities_contributions is None:
return None
for key, value in opacities_contributions.copy().items():
if value is None:
continue
if isinstance(value, dict):
specie_dict = {}
for species_name, species_data in value.copy().items():
specie_dict[species_name] = self.adapt_opacities(species_data, stellar_spectrum)
opacities_contributions[key] = specie_dict
else:
opacities_contributions[key] = self.adapt_opacities(value, stellar_spectrum)
return opacities_contributions
[docs]
def adapt_opacities(self, value, stellar_spectrum=None):
"""
Convert a single opacity contribution tuple to display-ready units.
Wavelength is converted from cm to nm. In transmission mode, the contribution
is converted from cm to Jupiter radii.
Args:
value: Tuple of (wavelength, contribution, extra_data) arrays in cgs units
stellar_spectrum: Stellar spectrum array (reserved for emission mode)
Returns:
Tuple of (wavelength_nm, contribution_converted, extra_data) arrays.
"""
temp_value1 = value[0].copy() * 1e7 # wl in nanometers
temp_value2 = value[1].copy() # contribution
temp_value3 = value[2].copy()
if self.rad_mode == "Transmission":
temp_value2 = temp_value2 / ConstantVariables.R_JUP_MEAN # from cm to Rjup
else:
# TODO opacities contribution for Emission
pass
# temp_value2 = (
# (temp_value2 / (np.pi * stellar_spectrum)) *
# (self.rp * ConstantVariables.R_JUP_MEAN / (self.stellar_radius * ConstantVariables.R_SUN)) ** 2
# ) + 1
return temp_value1, temp_value2, temp_value3
def _process_molecular_species(self, chemistry, params_only_molec, species):
"""
Process molecular species data and populate VMR dictionary.
This method handles the creation of VMR dictionaries based on the chemistry
type and molecular parameters. It processes species names and creates
mean VMR values for display.
Args:
chemistry: Chemistry table selection
params_only_molec: Molecular parameters
species: List of chemical species
"""
# Get molecular line lists and condensed data
molecs, isotopes, opacity_line_lists, _, isotope_masses = get_line_lists()
(
condensed_molecs_list, condensed_names_list,
condensed_masses, condensed_visualization_name
) = get_condensed_line_list()
# Process species based on chemistry type
if chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[1]:
self._process_chemistry_table_species(
species, opacity_line_lists, isotopes,
condensed_molecs_list, condensed_names_list
)
else:
self._process_molecular_parameters(
params_only_molec, species, opacity_line_lists, isotopes
)
# Convert vmr mean list to formatted string
self.vmr_mean = '\n'.join(self.vmr_mean)
def _process_chemistry_table_species(self, species, opacity_line_lists, isotopes,
condensed_molecs_list, condensed_names_list):
"""
Process species from chemistry table.
Args:
species: List of chemical species
opacity_line_lists: opacity_line_list mapping
isotopes: Isotope mapping
condensed_molecs_list: List of condensed molecules
condensed_names_list: List of condensed names
"""
# List of species to exclude from processing
excluded_species = ["H2", "He", "H", "H-", "e-"]
for index, key in enumerate(species):
if key not in excluded_species:
# Try to find isotope mapping
isotop = get_keys_by_value(opacity_line_lists, key)
if isotop is not None:
key = get_keys_by_value(isotopes, isotop)
else:
# Use condensed molecule list if isotope not found
key = condensed_molecs_list[np.where(condensed_names_list == key)][0]
# Store VMR data and create mean VMR string
self.vmr_dict[key] = self.vmr[:, index]
self.vmr_mean.append(f"VMR({key}) = {mean(self.vmr[:, index]):.7e}")
def _process_molecular_parameters(self, params_only_molec, species,
opacity_line_lists, isotopes):
"""
Process molecular parameters for VMR calculation.
Args:
params_only_molec: Molecular parameters
species: List of chemical species
opacity_line_lists: opacity_line_list mapping
isotopes: Isotope mapping
"""
# List of species to exclude from processing
excluded_species = ["H2", "He", "H", "H-", "e-"]
for mol in params_only_molec:
if mol is not None:
name = mol.name_for_list_molec
if name not in excluded_species:
# Find isotope mapping for the molecule
isotop = get_keys_by_value(opacity_line_lists, name)
name = get_keys_by_value(isotopes, isotop)
# Find species index and store VMR data
i = species.index(mol.molec_formula)
self.vmr_dict[name] = self.vmr[:, i]
self.vmr_mean.append(f"VMR({name}) = {mean(self.vmr[:, i]):.7e}")