Source code for GUIBRUSHR.Retrieval.ModelCalculation.Classes.Atmosphere

"""
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