"""
Atmosphere module for atmospheric modeling and radiative transfer calculations.
This module provides the Atmosphere class for handling atmospheric models,
including support for both petitRADTRANS and PyratBay radiative transfer codes.
"""
from pathlib import Path
from traceback import format_exc
import pyratbay.spectrum as ps
import numpy as np
from petitRADTRANS.plotlib import plot_opacity_contributions
from petitRADTRANS.radtrans import Radtrans
from petitRADTRANS import physical_constants as phys_const
from petitRADTRANS import physics as phys
from petitRADTRANS.stellar_spectra.phoenix import PhoenixStarTable
import pyratbay as pb
from scipy.integrate import trapezoid
from scipy.interpolate import splev
from scipy.stats import norm
import matplotlib
matplotlib.rcParams['agg.path.chunksize'] = 10000
import os
if not os.environ.get('SPHINX_BUILD'):
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Pressure import Pressure
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.SpeciesMolec import SpeciesMolec
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Resolution import Resolution
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Wavelength import Wavelength
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
generate_chemcat,
make_pyratbay_config, generate_star_model_hr, get_stellar_model
)
[docs]
def cloud_opas(lambda0_micron, kappa_cond, kappa_opac, omega, xi, temperatures):
"""
Create a cloud opacity function for atmospheric modeling.
This function returns a closure that calculates cloud opacities based on
wavelength-dependent scattering and absorption properties.
Parameters
----------
lambda0_micron : float
Central wavelength in microns
kappa_cond : float
Condensation opacity coefficient
kappa_opac : float or None
Opacity coefficient (if None, calculated from kappa_cond)
omega : float
Width parameter for wavelength dependence
xi : float
Asymmetry parameter
temperatures : array_like
Temperature profile
Returns
-------
function
Opacity function that takes wavelengths and pressures as input
"""
def get_opacity(wavelengths, pressures):
"""
Calculate cloud opacities for given wavelengths and pressures.
Parameters
----------
wavelengths : array_like
Wavelengths in microns
pressures : array_like
Pressures in bar
Returns
-------
array_like
Cloud opacities
"""
# Convert pressure from bar to dyn/cm^2 for CGS units
# The pressure must be in dyne because the kB constant is in erg/K (CGS)
# Conversion: P_dyn = 1e6 * P_bar
if kappa_opac is None:
pressures_dyn_cm_2 = 1e6 * pressures # dyn/cm^2
# Calculate total number density in cm^-3
ntot = pressures_dyn_cm_2 / (phys_const.kB * temperatures)
kappa_opac_f = kappa_cond * ntot
else:
kappa_opac_f = kappa_opac
# Calculate wavelength-dependent terms
arg = (wavelengths - lambda0_micron) / omega # Dimensionless
phi_min = norm.pdf(arg) # Probability density function (dimensionless)
phi_maiusc = norm.cdf(xi * arg) # Cumulative distribution function (dimensionless)
# Calculate final opacities
opacities = 2 * (phi_min * phi_maiusc)[:, np.newaxis] * kappa_opac_f
return opacities
return get_opacity
[docs]
class Atmosphere:
"""
Atmospheric model class for radiative transfer calculations.
This class handles atmospheric modeling using either petitRADTRANS or PyratBay
radiative transfer codes. It manages pressure profiles, chemical compositions,
wavelength grids, and stellar spectra for both high and low resolution calculations.
"""
[docs]
def __init__(self):
"""Initialize the Atmosphere object with default None values."""
# Stellar spectrum configuration
self.temporary_force_new_binning = True
self.hwhm_hr = None
self.hwhm_lr = None
self.stellar_spectrum_type_lr_mode = None
self.use_hr_linelists_for_lr = None
self.do_mean = False
self.stellar_spline_model_hr = None
self.stellar_spectrum_type_lr = None
self.stellar_spectrum_type_hr = None
self.stellar_spectrum = {}
# Model configuration
self.rv_sampling = None
self.retrieval_model = None
self.rad_mode = None
self.additional_opac = False
# Target and instrument configuration
self.target = None
self.resolution_obj = None
# Atmospheric structure
self.pressure_data = None
self.temp_min = None
self.temp_max = None
self.temp_step = None
# Chemical composition
self.chemistry = None
self.chemcat_obj = None
self.species_compatible_with_prt = None
self.species_obj = None
# Wavelength grids
self.wavelength = None
self.wl_hr_nm = None
self.wl_lr_nm = None
# Radiative transfer objects
self.atmosphere_hr = None # High resolution atmosphere object
self.atmosphere_lr = None # Low resolution atmosphere object
[docs]
def update_params(
self, line_species, line_species_isotope,
line_species_complete_name_hr, line_species_complete_name_lr, list_condensed_molecules,
rayleigh_species, range_min_pressures, range_max_pressures,
layers, lbl_high_res, lbl_low_res, resolution=None, instruments_HR=None,
instruments_LR=None, continum_opacity=None, mass_vector=None,
target=None, chemistry=None, manual_model_obj=None,
temp_min=300, temp_max=3000, temp_step=150,
initial_params=None, path_default="", rad_mode="Transmission",
manual_model_composition=None, retrieval_model="petitRADTRANS",
rv_sampling=0.1, additional_opac=False, start_molecs=0
):
"""
Update atmospheric model parameters.
This method configures the atmospheric model with species, pressure grid,
resolution settings, and other parameters needed for radiative transfer
calculations.
Parameters
----------
line_species : list
List of line species for opacity calculations
line_species_isotope : list
List of isotope-specific line species
line_species_complete_name_hr : list
Complete names of line species HR
line_species_complete_name_lr : list
Complete names of line species LR
list_condensed_molecules : list
List of condensed/cloud species
rayleigh_species : list
List of Rayleigh scattering species
range_min_pressures : float
Minimum pressure in the atmospheric grid
range_max_pressures : float
Maximum pressure in the atmospheric grid
layers : int
Number of atmospheric layers
lbl_high_res : float
Line-by-line high resolution sampling
lbl_low_res: float
Line-by-line low resolution sampling
resolution : float, optional
Spectral resolution
instruments_HR : list, optional
High resolution instruments
instruments_LR : list, optional
Low resolution instruments
continum_opacity : list, optional
Continuum opacity contributors
mass_vector : array_like, optional
Mass fractions for chemical species
target : object, optional
Target object containing system parameters
chemistry : str, optional
Chemistry model type
manual_model_obj : object, optional
Manual model configuration object
temp_min : float, default=300
Minimum temperature for calculations
temp_max : float, default=3000
Maximum temperature for calculations
temp_step : float, default=150
Temperature step for calculations
initial_params : dict, optional
Initial parameter values
path_default : str, default=""
Default path for data files
rad_mode : str, default="Transmission"
Radiative transfer mode
manual_model_composition : dict, optional
Manual composition specification
retrieval_model : str, default="petitRADTRANS"
Radiative transfer code to use
rv_sampling : float, default=0.1
Radial velocity sampling
additional_opac : bool, default=False
Whether to include additional opacity sources
start_molecs : int, default=0
Starting molecule index
"""
# Initialize default parameters if not provided
if initial_params is None:
initial_params = {}
# Unpack manual model object if provided
# This allows overriding individual parameters with a complete model object
if manual_model_obj is not None:
resolution = manual_model_obj.resolution
instruments_HR = manual_model_obj.instruments_HR
instruments_LR = manual_model_obj.instruments_LR
continum_opacity = manual_model_obj.continum_opacity
mass_vector = manual_model_obj.mass_vector
target = manual_model_obj.target
temp_min = manual_model_obj.temp_min
temp_max = manual_model_obj.temp_max
temp_step = manual_model_obj.temp_step
rad_mode = manual_model_obj.rad_mode
retrieval_model = manual_model_obj.retrieval_model
manual_model_composition = manual_model_obj.manual_model_composition
rv_sampling = manual_model_obj.rv_sampling
additional_opac = manual_model_obj.additional_opac
# Set core atmospheric model parameters
self.additional_opac = additional_opac
print(f"Additional opac {self.additional_opac}")
self.retrieval_model = retrieval_model
self.target = target
self.rad_mode = rad_mode
# Initialize resolution object for wavelength grid management
self.resolution_obj = Resolution(
resolution, instruments_HR, instruments_LR,
target.systemic_velocity, lbl_high_res, lbl_low_res
)
# Initialize pressure grid
self.pressure_data = Pressure(
range_min_pressures, range_max_pressures, layers
)
# Set temperature range parameters
self.temp_min = temp_min
self.temp_max = temp_max
self.temp_step = temp_step
self.chemistry = chemistry
self.rv_sampling = rv_sampling
# Initialize chemistry catalog object
self.chemcat_obj = None
# Handle different chemistry models
if chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[1]:
# Use provided mass vector for manual chemistry
mass_vector = np.array(mass_vector)
else:
# Generate chemistry catalog for equilibrium chemistry
mass_vector, self.chemcat_obj, self.species_compatible_with_prt = generate_chemcat(
self.pressure_data.range_min_pressures,
self.pressure_data.range_max_pressures,
self.pressure_data.n_layers,
path_default,
)
# Initialize species and molecular composition object
self.species_obj = SpeciesMolec(
continum_opacity,
line_species,
line_species_isotope,
line_species_complete_name_hr,
line_species_complete_name_lr,
list_condensed_molecules,
rayleigh_species,
mass_vector,
start_molecs,
initial_params,
manual_model_composition,
)
[docs]
def set_wl_range(self, minwlen_lr=None, maxwlen_lr=None,
minwlen_hr=None, maxwlen_hr=None):
"""
Set wavelength ranges for high and low resolution calculations.
This method determines the wavelength boundaries for both high and low
resolution calculations based on instrument specifications and data ranges.
If no explicit ranges are provided, they are automatically determined
from the instrument configurations.
Parameters
----------
minwlen_lr : float, optional
Minimum wavelength for low resolution calculations (microns)
maxwlen_lr : float, optional
Maximum wavelength for low resolution calculations (microns)
minwlen_hr : float, optional
Minimum wavelength for high resolution calculations (microns)
maxwlen_hr : float, optional
Maximum wavelength for high resolution calculations (microns)
"""
# Determine LOW RESOLUTION wavelength range
if minwlen_lr is None and self.resolution_obj.low_resolution():
minwlen_lr = []
maxwlen_lr = []
for instrument in self.resolution_obj.instruments_LR:
limit_ptr3 = 0.3 # Minimum wavelength limit in microns
# Use instrumental wavelength limits if specified
if instrument.wl_min is not None:
minwlen_lr.append(instrument.wl_min)
if instrument.wl_max is not None:
maxwlen_lr.append(instrument.wl_max)
# Determine wavelength limits from actual data range
lr_data = instrument.lr_data
# Find minimum wavelength with buffer
position = np.argmin(lr_data.ldata_LR)
min_wl_temp = lr_data.ldata_LR[position]
step_min_wl = lr_data.step_LR[position]
# Apply lower limit or add buffer to data range
min_wl_temp = (
limit_ptr3 if min_wl_temp < limit_ptr3 else ((min_wl_temp - 2 * step_min_wl) / 1000)
)
minwlen_lr.append(min_wl_temp)
# Find maximum wavelength with buffer
position = np.argmax(lr_data.ldata_LR)
max_wl_temp = lr_data.ldata_LR[position]
step_max_wl = lr_data.step_LR[position]
max_wl_temp = ((max_wl_temp + 2 * step_max_wl) / 1000)
maxwlen_lr.append(max_wl_temp)
# Take the most restrictive (minimum) and extensive (maximum) ranges
minwlen_lr = np.min(minwlen_lr)
maxwlen_lr = np.max(maxwlen_lr)
# Determine HIGH RESOLUTION wavelength range
if minwlen_hr is None and self.resolution_obj.high_resolution():
minwlen_hr = []
maxwlen_hr = []
for instrument in self.resolution_obj.instruments_HR:
minwlen_night = []
maxwlen_night = []
# Collect wavelength ranges from all observation nights
for icheck in range(instrument.ndate):
# Convert from nm to microns and add buffer
minwlen_night.append(
np.min(instrument.night_arr[icheck].lambdas) / 1000 - 0.1
)
maxwlen_night.append(
np.max(instrument.night_arr[icheck].lambdas) / 1000 + 0.1
)
# Choose between instrumental limits and night data limits
# This balances completeness with computational efficiency
minwlen_hr.append(np.maximum(instrument.wl_min, np.min(minwlen_night)))
maxwlen_hr.append(np.minimum(instrument.wl_max, np.max(maxwlen_night)))
# Take the most restrictive ranges across all instruments
minwlen_hr = np.min(minwlen_hr)
maxwlen_hr = np.max(maxwlen_hr)
# Initialize wavelength object with determined ranges
self.wavelength = Wavelength(
minwlen_lr, maxwlen_lr, minwlen_hr, maxwlen_hr
)
[docs]
def init_pyratbay(self, output_folder, initial_params):
"""
Initialize PyratBay radiative transfer module.
This method sets up the PyratBay atmospheric model with the specified
configuration and identifies available opacity models for clouds and
Rayleigh scattering.
Parameters
----------
output_folder : str or Path
Directory for PyratBay output files
initial_params : dict or None
Initial parameter values for the model
"""
# Generate PyratBay configuration file
config_file = make_pyratbay_config(
output_folder,
self.target, self.pressure_data, self.species_obj,
self.wavelength, self.resolution_obj, initial_params,
)
# Initialize PyratBay object
self.atmosphere_hr = pyrat = pb.Pyrat(config_file)
# Set logging verbosity to minimal
pyrat.log.verb = 0
# Identify and store indices for specific opacity models
opacity_models = [model.name for model in pyrat.opacity.models]
# Set cloud opacity model index if available
if 'deck' in opacity_models:
pyrat.icloud = opacity_models.index('deck')
# Set Rayleigh scattering model index if available
if 'lecavelier' in opacity_models:
pyrat.irayleigh = opacity_models.index('lecavelier')
def _read_opacities_low_resolution(self, path_target):
"""
Initialize low-resolution opacity tables and stellar spectra.
This method sets up the petitRADTRANS Radtrans object for low-resolution
calculations. It creates the atmospheric model grid and prepares stellar
spectra for emission calculations using one of three stellar spectrum types,
each with different processing modes.
Steps:
1. Define wavelength boundaries from configured LR range
2. Create Radtrans object:
- If use_hr_linelists_for_lr: Use line-by-line (lbl) opacity mode
- Otherwise: Use correlated-k opacity tables
3. Extract wavelength grid from Radtrans and compute bin widths
4. Load or compute stellar spectrum (for emission mode only):
Stellar Spectrum Types:
- Planck: Blackbody function at stellar effective temperature
(Same processing for both use_hr_linelists_for_lr modes)
- Phoenix HR: High-resolution Phoenix stellar atmosphere model
* If use_hr_linelists_for_lr: Bin HR spectrum to LR grid using ps.bin_spectrum
* Otherwise: Convolve HR spectrum with instrumental profile via spline
- Phoenix LR: Low-resolution Phoenix stellar atmosphere model
* If use_hr_linelists_for_lr: Bin Phoenix spectrum to LR grid using ps.bin_spectrum
* Otherwise: Interpolate Phoenix spectrum to LR wavelength grid
5. Store wavelength grid in nanometers
Parameters
----------
path_target : str
Path to target directory containing stellar spectrum files (for Phoenix HR)
Notes
-----
- For emission mode, stellar spectrum is required for planet-to-star flux ratio
- Sets self.atmosphere_lr, self.stellar_spectrum["model_PRT_LR"], and self.wl_lr_nm
- The use_hr_linelists_for_lr flag determines both opacity mode and stellar spectrum processing
"""
# Define wavelength range for low-resolution calculations
wl_min = self.wavelength.min_wl_lr
wl_max = self.wavelength.max_wl_lr
if self.use_hr_linelists_for_lr:
self.atmosphere_lr = Radtrans(
pressures=self.pressure_data.pressures,
line_species=self.species_obj.line_species_complete_name_lr,
rayleigh_species=self.species_obj.rayleigh_species,
gas_continuum_contributors=self.species_obj.continum_opacity,
wavelength_boundaries=np.array([wl_min, wl_max]),
line_opacity_mode="lbl", # Line-by-line mode for high resolution
line_by_line_opacity_sampling=self.resolution_obj.lbl_low_res,
)
else:
# Create low-resolution Radtrans object with correlated-k opacities
self.atmosphere_lr = Radtrans(
pressures=self.pressure_data.pressures,
line_species=self.species_obj.line_species_complete_name_lr,
cloud_species=self.species_obj.list_condensed_molecules,
rayleigh_species=self.species_obj.rayleigh_species,
gas_continuum_contributors=self.species_obj.continum_opacity,
wavelength_boundaries=np.array([wl_min, wl_max]),
)
# Extract wavelength grid from Radtrans (in cm)
wl_lr = self.atmosphere_lr.get_wavelengths()
# Calcola le differenze tra punti consecutivi
diffs = np.diff(wl_lr)
# Bin width per ogni punto
bin_width = np.zeros_like(wl_lr)
bin_width[0] = diffs[0] / 2 # primo: metà della distanza successiva
bin_width[1:-1] = (diffs[:-1] + diffs[1:]) / 2 # centrali: media delle due distanze
bin_width[-1] = diffs[-1] / 2 # ultimo: metà della distanza precedente
# Setup stellar spectrum for emission spectroscopy
# (Not needed for transmission mode)
if self.rad_mode == "Emission":
# Step 1: Load/compute base stellar spectrum based on type
if self.stellar_spectrum_type_lr == "Planck":
# Planck blackbody - same for both HR/LR processing modes
self.stellar_spectrum["model_PRT_LR"] = np.pi * phys.planck_function_cm(
self.target.stellar_effective_temperature, wl_lr
)
self.star_wavelength_cm = self.stellar_spectrum["model_PRT_LR"]
self.stellar_spectrum_LR = wl_lr
print("Used Planck function for stellar spectrum at LR")
elif self.stellar_spectrum_type_lr == "Phoenix HR":
# Load high-resolution Phoenix spectrum
path_stellar_spectrum = str(Path(
path_target, "Star_Spectrum", "star_spectrum_default.pkl"
))
if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
# Bin HR spectrum to LR grid
self.star_wavelength_cm, self.star_spectrum = get_stellar_model(
path_stellar_spectrum, np.min(wl_lr), np.max(wl_lr)
)
self.stellar_spectrum["model_PRT_LR"] = ps.bin_spectrum(
wl_lr * 1e4, self.star_wavelength_cm * 1e4, self.star_spectrum,
half_widths=bin_width * 1e4, gaps=None
)
print("Using high-resolution stellar spectrum for emission mode (binned)")
else:
# Convolve HR spectrum with instrumental profile
wl_mid = (wl_lr[:-1] + wl_lr[1:]) / 2
delta_wl = np.diff(wl_lr)
resolution = int(np.round(np.mean(wl_mid / delta_wl)))
self.hwhm_lr = ConstantVariables.CLIGHT / (2 * resolution)
self.star_wavelength_cm, self.star_spectrum, self.stellar_spline_model_lr = generate_star_model_hr(
path_stellar_spectrum, self.hwhm_lr, np.min(wl_lr), np.max(wl_lr)
)
self.stellar_spectrum["model_PRT_LR"] = splev(wl_lr, self.stellar_spline_model_lr)
print(f"Using high-resolution stellar spectrum for emission mode at resolution = {resolution}")
else: # Phoenix LR
# Load low-resolution Phoenix model
star_table = PhoenixStarTable()
stellar_spectrum, _ = star_table.compute_spectrum(
temperature=self.target.stellar_effective_temperature
)
self.star_wavelength_cm = stellar_spectrum[:, 0]
flux_star_Fnu = stellar_spectrum[:, 1]
# Convert from F_nu to F_lambda (erg/cm^2/s/cm)
self.star_spectrum = flux_star_Fnu * ConstantVariables.CLIGHT * 1e5 / (self.star_wavelength_cm**2)
# Process based on mode
if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
self.stellar_spectrum["model_PRT_LR"] = ps.bin_spectrum(
wl_lr * 1e4, self.star_wavelength_cm * 1e4, self.star_spectrum,
half_widths=bin_width * 1e4, gaps=None
)
print("Used Phoenix LR function for stellar spectrum at LR (binned)")
else:
self.stellar_spectrum["model_PRT_LR"] = np.interp(
wl_lr, self.star_wavelength_cm, self.star_spectrum
)
print("Used Phoenix LR function for stellar spectrum at LR (interpolated)")
# Convert wavelength grid from cm to nm for convenience
self.wl_lr_nm = wl_lr * 1e7
def _read_opacities_high_resolution(self, path_target):
"""
Initialize high-resolution opacity tables and stellar spectra.
This method sets up the petitRADTRANS Radtrans object for high-resolution
calculations using line-by-line (lbl) opacity mode. It creates the atmospheric
model grid at high spectral resolution and prepares stellar spectra for
emission calculations. Handles special case where HR linelists are used
for LR calculations.
Steps:
1. Define wavelength boundaries from configured HR range
2. Create Radtrans object with line-by-line opacity mode
3. Extract high-resolution wavelength grid from Radtrans
4. Load or compute stellar spectrum (for emission mode only):
- Planck: Use blackbody function at stellar temperature
- Phoenix/File: Load and spline-interpolate high-resolution stellar spectrum
5. If HR linelists used for LR: bin HR stellar spectrum to LR instrument resolution
6. Store wavelength grid in nanometers
Parameters
----------
path_target : str
Path to target-specific data directory containing stellar spectra
Notes
-----
- Only executed if high_resolution() is True or use_hr_linelists_for_lr is True
- Line-by-line mode provides much higher spectral resolution than correlated-k
- When use_hr_linelists_for_lr is True, HR opacities are binned to LR resolution
- Spline interpolation used for stellar spectrum to match Radtrans grid
- Sets self.atmosphere_hr, self.stellar_spectrum["model_PRT_HR"], and self.wl_hr_nm
"""
# Define wavelength range for high-resolution calculations
wl_min = self.wavelength.min_wl_hr
wl_max = self.wavelength.max_wl_hr
# Create high-resolution Radtrans object with line-by-line opacity mode
# This provides much finer spectral resolution than low-resolution correlated-k tables
self.atmosphere_hr = Radtrans(
pressures=self.pressure_data.pressures,
line_species=self.species_obj.line_species_complete_name_hr,
rayleigh_species=self.species_obj.rayleigh_species,
gas_continuum_contributors=self.species_obj.continum_opacity,
wavelength_boundaries=np.array([wl_min, wl_max]),
line_opacity_mode="lbl", # Line-by-line mode for high resolution
line_by_line_opacity_sampling=self.resolution_obj.lbl_high_res,
)
# Extract wavelength grid from Radtrans (in cm)
wl_hr = self.atmosphere_hr.get_wavelengths()
# Setup stellar spectrum for emission spectroscopy
# (Not needed for transmission mode)
if self.rad_mode == "Emission":
if self.stellar_spectrum_type_hr == "Planck":
# METHOD 1: Use Planck blackbody function
# Compute flux at each wavelength using stellar effective temperature
self.stellar_spectrum["model_PRT_HR"] = np.pi * phys.planck_function_cm(
self.target.stellar_effective_temperature, wl_hr
)
print("Used Planck function for stellar spectrum at HR")
else: # Phoenix or custom stellar spectrum from file
# METHOD 2: Load and spline-interpolate stellar spectrum
# Load high-resolution stellar spectrum from file
path_stellar_spectrum = str(Path(
path_target, "Star_Spectrum", "star_spectrum_default.pkl"
))
# Generate stellar model with spline interpolation
# Returns wavelength grid, full-resolution spectrum, and spline coefficients
_, _, self.stellar_spline_model_hr = generate_star_model_hr(
path_stellar_spectrum, self.hwhm_hr, np.min(wl_hr), np.max(wl_hr)
)
# Evaluate spline at Radtrans wavelength points
self.stellar_spectrum["model_PRT_HR"] = splev(wl_hr, self.stellar_spline_model_hr)
print(f"Used star spectrum {path_stellar_spectrum} for stellar spectrum at HR")
# Convert wavelength grid from cm to nm for convenience
self.wl_hr_nm = wl_hr * 1e7
def _stellar_spectrum_instrument_LR(self, instrument_LR):
"""
Bin stellar spectrum to low-resolution instrument wavelength bins.
This method takes either high-resolution or low-resolution stellar spectra
and bins them to match each instrument's wavelength resolution. The binning
approach depends on the use_hr_linelists_for_lr flag.
Two processing modes:
1. **use_hr_linelists_for_lr = True**:
- Uses HR stellar spectrum (self.star_spectrum, self.star_wavelength_cm)
- Bins using petitRADTRANS ps.bin_spectrum for all instruments
- Ignores stellar_spectrum_type_lr_mode setting
2. **use_hr_linelists_for_lr = False**:
- Uses LR stellar spectrum (self.stellar_spectrum["model_PRT_LR"])
- Bins using method specified by stellar_spectrum_type_lr_mode:
* "Integral": Integrates over each bin using trapezoidal rule
* "Point by Point": Preserves spectrum points within each bin
Process:
1. Skip if no low-resolution calculations needed
2. For each instrument:
- Bin stellar spectrum to instrument's wavelength grid
- Store result in self.stellar_spectrum[instrument.name]
Parameters
----------
instrument_LR : list
List of low-resolution instrument objects, each with:
- lr_data.ldata_LR_micron: wavelength centers (microns)
- lr_data.step_LR_micron: wavelength bin half-widths (microns)
- lr_data.npix: number of pixels
- lr_data.lbottom, lr_data.ltop: wavelength bin edges (nanometers)
Notes
-----
- Only applies to emission spectroscopy (not needed for transmission)
- Stores binned spectra in self.stellar_spectrum[instrument.name]
- For use_hr_linelists_for_lr=True, uses consistent binning with atmosphere
"""
# Skip if no low-resolution calculations needed
if not self.resolution_obj.low_resolution():
return
if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
# Use HR stellar spectrum and bin with ps.bin_spectrum (same as atmosphere binning)
for instrument in instrument_LR:
self.stellar_spectrum[instrument.name] = ps.bin_spectrum(
instrument.lr_data.ldata_LR_micron,
self.star_wavelength_cm * 1e4,
self.star_spectrum,
half_widths=instrument.lr_data.step_LR_micron,
gaps=None
)
else:
# Use LR stellar spectrum and bin according to stellar_spectrum_type_lr_mode
stellar_spc_reference = self.stellar_spectrum["model_PRT_LR"]
stellar_wavelength_nm = self.wl_lr_nm
for instrument in instrument_LR:
binned_spectrum = []
# Bin to each wavelength bin of the instrument
for i in range(instrument.lr_data.npix):
# Find stellar spectrum points within this instrument bin
mask = (
(stellar_wavelength_nm > instrument.lr_data.lbottom[i]) &
(stellar_wavelength_nm <= instrument.lr_data.ltop[i])
)
# Apply binning method
if self.stellar_spectrum_type_lr_mode == "Integral":
# Integrate stellar spectrum over bin using trapezoidal rule
flux_value = trapezoid(
stellar_spc_reference[mask],
stellar_wavelength_nm[mask]
)
else: # Point by Point
# Use stellar spectrum points directly (no integration)
flux_value = stellar_spc_reference[mask]
binned_spectrum.append(flux_value)
# Store binned stellar spectrum for this instrument
self.stellar_spectrum[instrument.name] = binned_spectrum
[docs]
def read_opacities(
self, table_output_file=None, path_target="", stellar_spectrum_type_lr="Phoenix",
stellar_spectrum_type_hr="Phoenix", min_hwhm_hr=1.5, instrument_LR=None,
stellar_spectrum_type_lr_mode="Point by Point", use_hr_linelists_for_lr=False
):
"""
Initialize opacity tables and stellar spectra for radiative transfer.
This method orchestrates the setup of petitRADTRANS Radtrans objects for both
high and low resolution calculations. It delegates to specialized methods for
LR and HR initialization, loads opacity data, and prepares stellar spectra
for emission calculations.
The method handles three resolution scenarios:
1. Low resolution only: Uses correlated-k opacity tables
2. High resolution only: Uses line-by-line opacity mode
3. Mixed mode: HR linelists binned to LR resolution (use_hr_linelists_for_lr)
Parameters
----------
table_output_file : str, optional
Path to output file for error messages
path_target : str, default=""
Path to target-specific data directory
stellar_spectrum_type_lr : str, default="Phoenix"
Type of stellar spectrum for LR ("Planck", "Phoenix", or "Integral")
stellar_spectrum_type_hr : str, default="Phoenix"
Type of stellar spectrum for HR ("Planck" or path to spectrum file)
min_hwhm_hr : float, default=1.5
Minimum pixel velocity width for high resolution (km/s)
instrument_LR : list, optional
List of instruments for low resolution calculations
stellar_spectrum_type_lr_mode : str, default="Point by Point"
Type of stellar spectrum for LR ("Point by Point" or "Integral")
use_hr_linelists_for_lr : bool, default=False
Use high resolution line lists for low resolution calculations
Warnings
--------
If opacity initialization fails, the error is caught internally,
a diagnostic message is written to table_output_file (if provided),
and the process exits via ``exit()``. No exception is raised to the caller.
Notes
-----
- Stellar spectra only loaded for emission mode (not transmission)
- Sets attributes: atmosphere_lr, atmosphere_hr, stellar_spectrum (dict with
"model_PRT_LR" and "model_PRT_HR" keys), wl_lr_nm, wl_hr_nm
"""
self.use_hr_linelists_for_lr = use_hr_linelists_for_lr
# Store stellar spectrum type configurations
self.stellar_spectrum_type_lr = stellar_spectrum_type_lr
self.stellar_spectrum_type_hr = stellar_spectrum_type_hr
self.stellar_spectrum_type_lr_mode = stellar_spectrum_type_lr_mode
self.stellar_spectrum = {}
self.hwhm_hr = min_hwhm_hr
try:
# Initialize LOW RESOLUTION radiative transfer
if self.resolution_obj.low_resolution():
self._read_opacities_low_resolution(
path_target=path_target
)
# Initialize HIGH RESOLUTION radiative transfer
if self.resolution_obj.high_resolution():
self._read_opacities_high_resolution(
path_target=path_target
)
# Bin stellar spectrum to instrument resolution (for emission mode)
if self.rad_mode != "Transmission" and instrument_LR is not None:
self._stellar_spectrum_instrument_LR(instrument_LR=instrument_LR)
except Exception as e:
# Handle errors in opacity initialization
print("Error initializing atmospheric opacities: " + str(e))
print(format_exc())
# Write error message to output file if specified
if table_output_file is not None:
with open(table_output_file, "w") as f:
f.write("Generate Atmosphere: Some molecules cannot be retrieved")
exit()
[docs]
def calc_model(self, temperature, mass_fraction, vmr, MMW, dict_calc_model,
HR=True):
"""
Compute the transmission or emission spectrum for given atmospheric model.
This method calculates either transmission or emission spectra using
either petitRADTRANS or PyratBay radiative transfer codes, depending
on the configuration.
Parameters
----------
temperature : array_like
Temperature profile
mass_fraction : dict
Mass fractions of atmospheric species
vmr : dict
Volume mixing ratios of atmospheric species
MMW : array_like
Mean molecular weight profile
dict_calc_model : dict
Dictionary containing model parameters
HR : bool, default=True
Whether to use high resolution calculation
Returns
-------
wl_full_resolution : array_like
Wavelengths (nm)
model : array_like or None
Raw model output (transit radii or flux). None when using PyratBay.
depth_full_resolution : array_like
Final spectrum (transmission depth or emission)
"""
# Extract model parameters from dictionary
Pc = dict_calc_model["Pc"] # Cloud top pressure
rp = dict_calc_model["rp"] # Planet radius
k0 = dict_calc_model["k0"] # Haze opacity at reference wavelength
gamma = dict_calc_model["gamma"] # Haze opacity power law index
gravity = dict_calc_model["gravity"] # Surface gravity
haze_factor = dict_calc_model["haze_factor"] # Haze enhancement factor
cloud_fraction = dict_calc_model["cloud_fraction"] # Cloud coverage fraction
k_cond = dict_calc_model["k_cond"] # Condensation opacity
k_opac = dict_calc_model["k_opac"] # Additional opacity
lambda0_micron = dict_calc_model["lambda0_micron"] # Reference wavelength
P_ref = dict_calc_model["P_ref"] # Reference pressure
xi = dict_calc_model["xi"] # Cloud asymmetry parameter
std_radius_distribution = dict_calc_model["std_radius_distribution"] # Cloud particle size distribution
cloud_fsed = dict_calc_model["cloud_fsed"] # Cloud sedimentation efficiency
eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"] # Eddy diffusion coefficient
omega_scale_micron = dict_calc_model["omega_scale_micron"] # Cloud wavelength scale
# Initialize output variables
wl_full_resolution = None
depth_full_resolution = None
model = None
# Calculate spectrum using petitRADTRANS
if self.retrieval_model == 'petitRADTRANS':
# Setup additional cloud opacity function if enabled
additional_absorption_opacities_function = (
cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperature)
if self.additional_opac else None
)
# Select appropriate resolution model
rt_model = self.atmosphere_hr if HR else self.atmosphere_lr
if self.rad_mode == "Transmission":
# Calculate transmission spectrum
_, model, _ = rt_model.calculate_transit_radii(
temperatures=temperature,
mass_fractions=mass_fraction,
mean_molar_masses=MMW,
planet_radius=rp,
reference_pressure=P_ref,
reference_gravity=gravity,
opaque_cloud_top_pressure=Pc,
haze_factor=haze_factor,
cloud_fraction=cloud_fraction,
power_law_opacity_350nm=k0,
power_law_opacity_coefficient=gamma,
frequencies_to_wavelengths=True,
additional_absorption_opacities_function=additional_absorption_opacities_function,
cloud_f_sed=cloud_fsed,
eddy_diffusion_coefficients=eddy_diff_coeff,
cloud_particle_radius_distribution_std=std_radius_distribution,
)
# Convert transit radii to transmission depth
mol_modf = (
(model / phys_const.r_jup_mean) /
(self.target.stellar_radius * ConstantVariables.RATIO_RSUN_RJUP_MEAN)
) ** 2
depth_full_resolution = 1 - mol_modf # Transmission depth
else:
# Calculate emission spectrum
_, model, _ = rt_model.calculate_flux(
temperatures=temperature,
mass_fractions=mass_fraction,
mean_molar_masses=MMW,
planet_radius=rp,
reference_gravity=gravity,
opaque_cloud_top_pressure=Pc,
haze_factor=haze_factor,
cloud_fraction=cloud_fraction,
power_law_opacity_350nm=k0,
power_law_opacity_coefficient=gamma,
frequencies_to_wavelengths=True,
additional_absorption_opacities_function=additional_absorption_opacities_function,
cloud_f_sed=cloud_fsed,
eddy_diffusion_coefficients=eddy_diff_coeff,
cloud_particle_radius_distribution_std=std_radius_distribution,
)
stellar_spectrum = self.stellar_spectrum["model_PRT_HR"] if HR else self.stellar_spectrum["model_PRT_LR"]
# Calculate scaled flux ratio (planet/star)
# 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 / stellar_spectrum) *
(rp / (self.target.stellar_radius * phys_const.r_sun)) ** 2
)
depth_full_resolution = 1 + Fscaled # Relative flux
# Set wavelength np.array
wl_full_resolution = self.wl_hr_nm if HR else self.wl_lr_nm
# Calculate spectrum using PyratBay
elif self.retrieval_model == 'PyratBay':
pyrat = self.atmosphere_hr
# Update atmospheric model parameters
pyrat.atm.rplanet = rp
pyrat.atm.refpressure = np.log10(P_ref)
# Update cloud parameters if specified
if Pc is not None:
cloud = pyrat.opacity.models[pyrat.icloud]
cloud.pars[0] = np.log10(Pc)
# Update Rayleigh scattering parameters if specified
if k0 is not None:
rayleigh = pyrat.opacity.models[pyrat.irayleigh]
rayleigh.pars[0] = np.log10(k0)
if gamma is not None:
rayleigh = pyrat.opacity.models[pyrat.irayleigh]
rayleigh.pars[1] = gamma
# Run radiative transfer calculation
pyrat.run(temperature, vmr)
# Extract results and convert units
wl_full_resolution = 1.0 / np.flip(pyrat.spec.wn) * 1e7 # Convert wavenumber to wavelength (nm)
depth = np.flip(pyrat.spec.spectrum) # Transit depth
depth_full_resolution = 1.0 - depth # Transmission spectrum
return wl_full_resolution, model, depth_full_resolution
[docs]
def calc_model_contribution(self, temperature, mass_fraction, MMW,
dict_calc_model, HR=True):
"""
Calculate transmission contribution function for atmospheric layers.
This method computes the contribution of different atmospheric layers
to the total transmission spectrum, useful for understanding which
pressure levels dominate the observed signal.
Parameters
----------
temperature : array_like
Temperature profile
mass_fraction : dict
Mass fractions of atmospheric species
MMW : array_like
Mean molecular weight profile
dict_calc_model : dict
Dictionary containing model parameters
HR : bool, default=True
Whether to use high resolution calculation
Returns
-------
wl : array_like
Wavelengths (nm)
contribution : array_like
Transmission contribution function
"""
# Extract model parameters
Pc = dict_calc_model["Pc"]
rp = dict_calc_model["rp"]
k0 = dict_calc_model["k0"]
gamma = dict_calc_model["gamma"]
gravity = dict_calc_model["gravity"]
haze_factor = dict_calc_model["haze_factor"]
cloud_fraction = dict_calc_model["cloud_fraction"]
k_cond = dict_calc_model["k_cond"]
k_opac = dict_calc_model["k_opac"]
lambda0_micron = dict_calc_model["lambda0_micron"]
P_ref = dict_calc_model["P_ref"]
xi = dict_calc_model["xi"]
omega_scale_micron = dict_calc_model["omega_scale_micron"]
std_radius_distribution = dict_calc_model["std_radius_distribution"]
cloud_fsed = dict_calc_model["cloud_fsed"]
eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"]
# Select appropriate atmosphere model
atmo = self.atmosphere_hr if HR else self.atmosphere_lr
# Setup additional opacity function if enabled
additional_absorption_opacities_function = (
cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperature)
if self.additional_opac else None
)
# Calculate transit radii with contribution function
wl, transit_radii, additional = atmo.calculate_transit_radii(
temperatures=temperature,
mass_fractions=mass_fraction,
mean_molar_masses=MMW,
planet_radius=rp,
reference_pressure=P_ref,
reference_gravity=gravity,
opaque_cloud_top_pressure=Pc,
haze_factor=haze_factor,
cloud_fraction=cloud_fraction,
power_law_opacity_350nm=k0,
power_law_opacity_coefficient=gamma,
frequencies_to_wavelengths=True,
return_contribution=True, # Request contribution function
additional_absorption_opacities_function=additional_absorption_opacities_function,
cloud_f_sed=cloud_fsed,
eddy_diffusion_coefficients=eddy_diff_coeff,
cloud_particle_radius_distribution_std=std_radius_distribution,
)
# Convert wavelengths to nanometers
wl *= 1e7
return wl, additional["transmission_contribution"]
[docs]
def opacity_contribution(self, temperatures, mass_fractions, MMW,
dict_calc_model, HR=True):
"""
Generate opacity contribution plots for atmospheric species.
This method creates plots showing the relative contribution of different
opacity sources (molecular species, clouds, hazes) to the total spectrum.
Parameters
----------
temperatures : array_like
Temperature profile
mass_fractions : dict
Mass fractions of atmospheric species
MMW : array_like
Mean molecular weight profile
dict_calc_model : dict
Dictionary containing model parameters
HR : bool, default=True
Whether to use high resolution calculation
Returns
-------
opacity_contributions : matplotlib.figure.Figure or None
Figure object with opacity contribution plots, or None if failed
"""
opacity_contributions = None
try:
# Extract model parameters
Pc = dict_calc_model["Pc"]
rp = dict_calc_model["rp"]
k0 = dict_calc_model["k0"]
gamma = dict_calc_model["gamma"]
gravity = dict_calc_model["gravity"]
haze_factor = dict_calc_model["haze_factor"]
cloud_fraction = dict_calc_model["cloud_fraction"]
k_cond = dict_calc_model["k_cond"]
k_opac = dict_calc_model["k_opac"]
lambda0_micron = dict_calc_model["lambda0_micron"]
P_ref = dict_calc_model["P_ref"]
xi = dict_calc_model["xi"]
omega_scale_micron = dict_calc_model["omega_scale_micron"]
std_radius_distribution = dict_calc_model["std_radius_distribution"]
cloud_fsed = dict_calc_model["cloud_fsed"]
eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"]
# Select appropriate atmosphere model
rt_model = self.atmosphere_hr if HR else self.atmosphere_lr
# Setup additional opacity function if enabled
additional_absorption_opacities_function = (
cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperatures)
if self.additional_opac else None
)
# Generate opacity contribution plots based on radiative transfer mode
if self.rad_mode == "Transmission":
# Plot transmission opacity contributions
opacity_contributions = plot_opacity_contributions(
rt_model,
mode=self.rad_mode.lower(),
fill_below=False, # Use curves instead of filled areas
x_axis_scale='log', # Logarithmic pressure scale
temperatures=temperatures,
mass_fractions=mass_fractions,
mean_molar_masses=MMW,
reference_gravity=gravity,
planet_radius=rp,
reference_pressure=P_ref,
opaque_cloud_top_pressure=Pc,
haze_factor=haze_factor,
cloud_fraction=cloud_fraction,
power_law_opacity_350nm=k0,
power_law_opacity_coefficient=gamma,
additional_absorption_opacities_function=additional_absorption_opacities_function,
cloud_f_sed=cloud_fsed,
eddy_diffusion_coefficients=eddy_diff_coeff,
cloud_particle_radius_distribution_std=std_radius_distribution,
)
else:
# TODO adjust and insert again plot
print("Emission opacity contributions not yet fully implemented")
return None
# Plot emission opacity contributions
# opacity_contributions = plot_opacity_contributions(
# rt_model,
# mode=self.rad_mode.lower(),
# fill_below=False, # Use curves instead of filled areas
# x_axis_scale='log', # Logarithmic pressure scale
# temperatures=temperatures,
# mass_fractions=mass_fractions,
# mean_molar_masses=MMW,
# reference_gravity=gravity,
# planet_radius=rp,
# opaque_cloud_top_pressure=Pc,
# haze_factor=haze_factor,
# cloud_fraction=cloud_fraction,
# power_law_opacity_350nm=k0,
# power_law_opacity_coefficient=gamma,
# additional_absorption_opacities_function=additional_absorption_opacities_function,
# cloud_f_sed=cloud_fsed,
# eddy_diffusion_coefficients=eddy_diff_coeff,
# cloud_particle_radius_distribution_std=std_radius_distribution,
# )
# Close the plot to free memory
plt.close()
except Exception as e:
# Handle errors in opacity contribution calculation
print("Opacity contributions failed: " + str(e) + "\n" + format_exc())
return opacity_contributions