"""
I/O utility functions for GUIBRUSHR atmospheric retrieval.
This module contains functions for file I/O and data access extracted from
general_functions.py:
- make_pyratbay_config: Create a pyratbay configuration file.
- create_fits: Create and save a FITS file.
- create_path_night: Construct observation data file path.
- get_csv_value: Retrieve value from a DataFrame.
- get_condensed_line_list: Retrieve condensed molecule info from pRT.
- get_line_lists: Retrieve line lists from molecular opacity data.
- read_fits: Read FITS file with optional slicing.
- get_depth_filename: Return depth FITS filename by mode.
- read_depth_fits: Read depth FITS with legacy fallback.
- get_stellar_model: Load stellar spectrum from pickle.
- generate_star_model_hr: Generate stellar model with convolution.
- download_star_spectrum: Download PHOENIX stellar spectrum.
"""
import os
import pickle
import traceback
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union
from astropy.io import fits
import numpy as np
from pandas import read_csv, DataFrame
from petitRADTRANS.config import petitradtrans_config_parser
from scipy.interpolate import splrep
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.core.functions.spectroscopy import calculate_mass_molecule, spectral_convolution_ptr
[docs]
def make_pyratbay_config(
output_folder: str,
target: Any,
pressure: Any,
species: Any,
wavelength: Any,
resolution: Any,
params
) -> str:
"""
Create a pyratbay configuration file based on guibrush configuration.
Parameters:
output_folder (str): Folder where to store the outputs.
target: Target object containing stellar and planetary parameters.
pressure: Pressure configuration object.
species: Species configuration object.
wavelength: Wavelength configuration object.
resolution: Resolution configuration object.
params: List of input parameters.
Returns:
str: Path to the created configuration file.
"""
try:
param_values = {
param.name: param.get_starting_value()
for param in params
if param is not None
}
log_file = f'{output_folder}/giano_b_retrieval.log'
cfg_file = f'{output_folder}/giano_b_retrieval.cfg'
# Build configuration text
cfg_text = (
"[pyrat]\n\nrunmode = spectrum\n"
f"logfile = {log_file}\n"
"rt_path = transit\n\n"
)
# System parameters
logp_ref = target.p0
cfg_text += (
"# System parameters\n"
f"rstar = {target.stellar_radius} rsun\n"
f"mstar = {target.stellar_mass} msun\n"
f"rplanet = {target.radius} rjup\n"
f"mplanet = {target.mass} mjup\n"
f"refpressure = {logp_ref} bar\n\n"
)
# Wavelength sampling
cfg_text += (
"# Spectral sampling\n"
f"wllow = {wavelength.min_wl_hr} um\n"
f"wlhigh = {wavelength.max_wl_hr} um\n"
f"wn_thinning = {resolution.lbl_high_res}\n\n"
)
# Atmospheric model
p_max = 10 ** pressure.range_max_pressures
p_min = 10 ** pressure.range_min_pressures
cfg_text += (
"# Atmospheric model\n"
f"nlayers = {pressure.n_layers}\n"
f"ptop = {p_min} bar\n"
f"pbottom = {p_max} bar\n\n"
)
# Chemistry configuration
atm_species = ' '.join(list(species.composition))
atm_vmr = []
for mol, vmr in species.composition.items():
if mol in ['H2', 'He']:
atm_vmr.append(str(vmr))
else:
atm_vmr.append('1e-10')
mass_fraction = ' '.join(atm_vmr)
cfg_text += (
"chemistry = uniform\n"
f"species = {atm_species}\n"
f"uniform = {mass_fraction}\n"
"radmodel = hydro_m\n\n"
)
# Temperature model
cfg_text += (
"tmodel = isothermal\n"
"tpars = 1000.0\n\n"
)
# Opacities configuration
cfg_text += "# Opacities\n"
lbl_models = ''
prt_path = petitradtrans_config_parser.get_input_data_path()
for i, mol in enumerate(species.line_species):
opacity_file = '\n ' + str(
Path(prt_path, "opacities", "lines", "line_by_line", mol,
species.line_species_isotope[i],
species.line_species_complete_name_hr[i] + ".xsec.petitRADTRANS.h5"))
lbl_models += opacity_file
if len(species.line_species) > 0:
cfg_text += (
f'extfile ={lbl_models}\n\n'
'tmin = 300\n'
'tmax = 3000\n'
'tstep = 150\n\n'
)
# CIA configuration
root = '{ROOT}/pyratbay/data/CIA/'
cia_dict = {
'H2-H2': root + 'CIA_Borysow_H2H2_0060-7000K_0.6-500um.dat',
'H2-He': root + 'CIA_Borysow_H2He_0050-7000K_0.5-031um.dat',
}
cia_models = ''
for pair in species.continum_opacity:
cia_models += f'\n {cia_dict[pair]}'
if len(species.continum_opacity) > 0:
cfg_text += f'csfile ={cia_models}\n\n'
# Rayleigh scattering configuration
rayleigh = None
if len(species.rayleigh_species) > 0:
rayleigh = [
f'\n dalgarno_{spec}'
for spec in species.rayleigh_species
]
ray_params = ''
if "k0" in param_values:
rayleigh.append('\n lecavelier')
ray_params = f"\nrpars = {param_values['k0']} {param_values['gamma']}"
if len(rayleigh) > 0:
rayleigh_models = ''.join(rayleigh)
cfg_text += f'rayleigh ={rayleigh_models}{ray_params}\n\n'
# Cloud configuration
if 'Pc' in param_values:
cfg_text += f'clouds = deck\ncpars = {param_values["Pc"]}\n\n'
# Write configuration to file
with open(cfg_file, 'w') as f:
f.write(cfg_text)
return cfg_file
except IOError as e:
print(f"Failed to write configuration file: {str(e)}")
[docs]
def create_fits(data: np.ndarray, pathf: Union[str, Path]) -> None:
"""
Create and save a FITS file from the provided data.
This function creates a primary HDU (Header/Data Unit) using the input data,
encapsulates it in an HDUList, and writes the FITS file to the specified path.
If a file with the same name exists, it will be overwritten.
Parameters:
data (np.ndarray): The data array to be stored in the FITS file.
pathf (Union[str, Path]): The destination file path for the FITS file.
"""
try:
if not isinstance(data, np.ndarray):
print("Input data must be a numpy array")
# Create a primary HDU from the input data
hdu = fits.PrimaryHDU(data)
# Create an HDUList to hold the primary HDU
hdul = fits.HDUList([hdu])
# Write the HDUList to the file, overwriting any existing file
hdul.writeto(pathf, overwrite=True)
except IOError as e:
print(f"Failed to write FITS file: {str(e)}")
except Exception as e:
print(f"Invalid data format: {str(e)}")
[docs]
def create_path_night(
pathfolder: Union[str, Path],
target: str,
rad_mode: str,
instrument: str
) -> str:
"""
Construct and return the file path for a specific night's observation data.
The path is built by concatenating the following components:
- `pathfolder`: The base directory containing observation data.
- `target`: The target object or observation name.
- A constant subdirectory defined by `ConstantVariables.HR_INSTRUMENTS`, representing high-resolution instruments.
- `rad_mode`: The radiative mode, which specifies the type of observation.
This can be either 'emission' or 'transmission'.
- `instrument`: The instrument identifier.
Parameters:
pathfolder (Union[str, Path]): The base directory path.
target (str): The target object or observation name.
rad_mode (str): The radiative mode; valid options are 'emission' or 'transmission'.
instrument (str): The instrument identifier.
Returns:
str: The fully constructed file path as a string.
"""
# Validate rad_mode
if rad_mode not in ConstantVariables.LIST_RAD_MODE:
print(f"Invalid rad_mode: {rad_mode}. Must be 'Emission' or 'Transmission'")
# Construct the full path using pathlib.Path
full_path = Path(pathfolder, target, ConstantVariables.HR_INSTRUMENTS, rad_mode, instrument)
# Convert the Path object to a string and return it
return str(full_path)
[docs]
def get_csv_value(
df,
search_value: Any,
search_column: str = "Variable",
target_column: str = "Value"
) -> Optional[Any]:
"""
Retrieve a value from a general-info container by key.
This helper accepts three input shapes (the legacy DataFrame mode is preserved
so the 100+ existing call sites keep working without edits):
1. ``pd.DataFrame`` with a 2-column ``Variable``/``Value`` legacy shape.
The ``search_column``/``target_column`` arguments override the defaults.
2. ``dict`` (typed general-info, as returned by
:func:`GUIBRUSHR.core.io.retrieval_io.read_general_info`). Equivalent to
``data.get(search_value)``. ``search_column``/``target_column`` are ignored.
3. ``str`` or ``pathlib.Path`` pointing to a retrieval folder. The function
calls :func:`read_general_info` internally and looks up the key.
``search_column``/``target_column`` are ignored.
Parameters:
df: Container to query. ``DataFrame``, ``dict``, ``str``, or ``Path``.
search_value (Any): The key (or value in ``search_column``) to look up.
search_column (str, optional): Only used in DataFrame mode. Defaults to "Variable".
target_column (str, optional): Only used in DataFrame mode. Defaults to "Value".
Returns:
Optional[Any]: The matching value (typed when input is dict/folder; raw cell
when input is DataFrame), or ``None`` if no match is found.
"""
# Folder mode: lazy import to avoid a circular dependency
if isinstance(df, (str, Path)):
from GUIBRUSHR.core.io.retrieval_io import read_general_info, RetrievalIOError
try:
data = read_general_info(df)
except RetrievalIOError:
return None
return data.get(search_value)
# Dict mode (typed general-info)
if isinstance(df, dict):
return df.get(search_value)
# Legacy DataFrame mode — original behavior preserved
if search_column not in df.columns:
print(f"Search column '{search_column}' not found in DataFrame")
if target_column not in df.columns:
print(f"Target column '{target_column}' not found in DataFrame")
matching_rows = df[df[search_column] == search_value]
return None if matching_rows.empty else matching_rows[target_column].values[0]
[docs]
def get_condensed_line_list():
"""
Retrieve condensed molecule information from petitRADTRANS opacity data.
This function scans the condensed molecule opacity files in the petitRADTRANS
installation and extracts information about available condensed species
including their chemical formulas, names, masses, and visualization names.
Returns:
tuple: A tuple containing:
- np.ndarray: Array of condensed molecule chemical formulas
- np.ndarray: Array of condensed molecule file names
- np.ndarray: Array of calculated molecular masses (in u)
- np.ndarray: Array of visualization names for display
Note:
If no condensed molecules are found, returns arrays with ["None"]
and mass [-999] as placeholders.
"""
# Load required paths from ConstantVariables
petitradtrans = ConstantVariables.path_petitradtrans
path_default = ConstantVariables.path_default
# Read the periodic table CSV file (assumed to return a pandas DataFrame)
periodic_table_path = Path(path_default, "GUIBRUSHR", "Files", "Molecules", "periodic_table.csv")
periodic_table = read_csv(str(periodic_table_path))
# Define the path containing line-by-line opacity data for molecules
clouds_path = Path(petitradtrans, "opacities", "continuum", "clouds")
# Initialize containers for molecules, isotopes, opacity_line_lists_HR, and computed masses
all_condensed_molec_from_yaml = ConstantVariables.ALL_CONDENSED_MOLEC
condensed_molecs = []
condensed_names = []
condensed_masses = []
condensed_visualization_name = []
#
try:
# Retrieve a sorted list of molecule directories
condensed_names = [cloud.name.removesuffix(".cotable.petitRADTRANS.h5") for cloud in sorted(clouds_path.glob('**/*.h5'))]
for i, condensed_name in enumerate(condensed_names):
name_split = condensed_name.split("-")
chemical_formula = []
for nn in name_split:
if "NatAbund" in nn:
break
else:
chemical_formula.append(nn)
condensed_molecs.append("".join(chemical_formula))
condensed_masses.append(calculate_mass_molecule(chemical_formula, periodic_table))
if len(condensed_names) == 0:
condensed_visualization_name = ["None"]
condensed_molecs = ["None"]
condensed_names = ["None"]
condensed_masses = [-999]
else:
condensed_visualization_name = [all_condensed_molec_from_yaml[key] for key in condensed_names]
except Exception as e:
print(e, "\nNo condensed molecules present in opacities")
# Return the sorted molecules array and the dictionaries for isotopes, opacity_line_lists_HR, and masses
return np.array(condensed_molecs), np.array(condensed_names), np.array(condensed_masses), np.array(condensed_visualization_name)
[docs]
def get_line_lists():
"""
Retrieve line lists, isotopes, opacity_line_list identifiers, and compute isotopic masses
from molecular opacity data used by petitRADTRANS.
This function performs the following steps:
1. Reads the periodic table data from a CSV file to obtain numbers of protons
and neutrons for each element.
2. Scans the 'line_by_line' directory (within the opacities folder) for molecule directories.
3. For each molecule, lists its available isotopes (each isotope is represented by a directory).
4. Parses each isotope's name (which encodes its elemental composition) to
determine the elemental symbols and their corresponding mass numbers,
compute the total mass of the isotope using predefined atomic masses
(proton: 1.007276 u, neutron: 1.008665 u, electron: 0.000548579909 u),
and adjust the electron count based on the number of '+' symbols.
5. For each isotope, retrieves the available opacity_line_list file names (by removing the
".xsec.petitRADTRANS.h5" suffix).
Returns:
tuple: A tuple containing:
- np.ndarray: Sorted array of molecule names.
- dict: Dictionary mapping each molecule to an array of its isotope names.
- dict: Dictionary mapping each isotope to an array of HR opacity_line_list identifiers.
- dict: Dictionary mapping each isotope to an array of LR opacity_line_list identifiers.
- dict: Dictionary mapping each isotope to its computed total mass (in atomic mass units).
"""
# Load required paths from ConstantVariables
petitradtrans = ConstantVariables.path_petitradtrans
path_default = ConstantVariables.path_default
# Read the periodic table CSV file (assumed to return a pandas DataFrame)
periodic_table_path = Path(path_default, "GUIBRUSHR", "Files", "Molecules", "periodic_table.csv")
periodic_table = read_csv(str(periodic_table_path))
# Define the path containing line-by-line opacity data for molecules
line_lists_path = Path(petitradtrans, "opacities", "lines", "line_by_line")
# Initialize containers for molecules, isotopes, opacity_line_lists_HR, and computed masses
molecs = []
isotopes = dict()
opacity_line_lists_HR = dict()
opacity_line_lists_LR = dict()
masses_isotopes = dict()
try:
# Retrieve a sorted list of molecule directories
molecs = sorted([d.name for d in line_lists_path.iterdir() if d.is_dir()])
# Loop over each molecule directory
for molec in molecs:
# Path to the current molecule's isotope directories
path_isotope = line_lists_path / molec
# List available isotopes for this molecule (each as a directory)
isotopes[molec] = np.array([d.name for d in path_isotope.iterdir() if d.is_dir()])
# Process each isotope for the current molecule
for isotope in isotopes[molec]:
# The isotope name encodes its elemental composition separated by "-"
elements = isotope.split("-")
# Store the computed mass for the isotope
masses_isotopes[isotope] = calculate_mass_molecule(elements, periodic_table)
# Retrieve opacity_line_list file names (strip the specific suffix) from the isotope directory
path_opacity_line_list = path_isotope / isotope
opacity_line_lists_HR[isotope] = np.array([
d.name.removesuffix(".xsec.petitRADTRANS.h5")
for d in path_opacity_line_list.iterdir() if not d.is_dir()
])
if len(opacity_line_lists_HR[isotope]) == 0:
opacity_line_lists_HR[isotope] = np.array(["None"])
# LR
path_opacity_line_list_LR = Path(str(path_opacity_line_list).replace("line_by_line", "correlated_k"))
if path_opacity_line_list_LR.exists():
opacity_line_lists_LR[isotope] = np.array([
d.name.removesuffix(".ktable.petitRADTRANS.h5")
for d in path_opacity_line_list_LR.iterdir() if not d.is_dir()
])
if len(opacity_line_lists_LR[isotope]) == 0:
opacity_line_lists_LR[isotope] = ["None"]
else:
opacity_line_lists_LR[isotope] = ["None"]
except Exception as e:
print(str(traceback.format_exc()), "\nNo molecules present in opacities")
# Return the sorted molecules array and the dictionaries for isotopes, opacity_line_lists_HR, and masses
return np.array(molecs), isotopes, opacity_line_lists_HR, opacity_line_lists_LR, masses_isotopes
[docs]
def read_fits(
path_fits: Union[str, Path],
transpose: bool = False,
intransit: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Read data from a FITS file and optionally apply slicing and transposition.
Parameters:
path_fits (Union[str, Path]): Path to the FITS file
transpose (bool, optional): Whether to transpose the data array. Defaults to False.
intransit (Optional[np.ndarray], optional): Indices to slice the data along the second axis.
If provided, data is sliced as data[:, intransit, :].
Defaults to None.
Returns:
np.ndarray: The processed data array from the FITS file
"""
try:
with fits.open(path_fits) as hdul:
data = hdul[0].data
if intransit is not None:
data = data[:, intransit, :]
if transpose:
data = np.transpose(data)
return data
except FileNotFoundError:
print(f"FITS file not found at {path_fits}")
except Exception as e:
print(f"Error reading FITS file: {str(e)}")
[docs]
def get_depth_filename(rad_mode: str) -> str:
"""Return the depth FITS filename based on the radiative transfer mode."""
if rad_mode == "Emission":
return "emission_depth.fits"
return "transmission_depth.fits"
[docs]
def read_depth_fits(folder: Union[str, Path], rad_mode: str, **kwargs) -> np.ndarray:
"""Read the depth FITS file with fallback to legacy transit_depth.fits."""
folder = Path(folder)
primary = folder / get_depth_filename(rad_mode)
fallback = folder / "transit_depth.fits"
if primary.exists():
return read_fits(str(primary), **kwargs)
if fallback.exists():
print(f"{primary.name} not found, falling back to transit_depth.fits")
return read_fits(str(fallback), **kwargs)
return read_fits(str(primary), **kwargs)
[docs]
def get_stellar_model(path_pkl, min_wl_cm, max_wl_cm):
"""Load a stellar spectrum from a pickle file and optionally trim to a wavelength range (cm)."""
with open(path_pkl, "rb") as f:
stellar_spectrum = pickle.load(f)
wavelength = stellar_spectrum["wavelength"]
spectrum = stellar_spectrum["spectrum"]
if min_wl_cm is not None and max_wl_cm is not None:
mask = (wavelength >= min_wl_cm) & (wavelength <= max_wl_cm)
wavelength = wavelength[mask]
spectrum = spectrum[mask]
return wavelength, np.pi * spectrum
[docs]
def generate_star_model_hr(path_pkl, pixel_dv, min_wl=None, max_wl=None):
"""
Generate a stellar spectrum model from a pickled file.
This function loads a stellar spectrum from a pickle file, convolves it
with the instrumental resolution, and creates a spline interpolation model.
Parameters
----------
path_pkl : str or Path
Path to the pickle file containing stellar spectrum data
pixel_dv : float
Pixel velocity width for resolution calculation
min_wl : float
Minimum wavelength for interpolation model in cm
max_wl : float
Maximum wavelength for interpolation model in cm
Returns
-------
tuple
Spline representation of the stellar spectrum model
"""
# Load stellar spectrum data from pickle file
wavelength, spectrum = get_stellar_model(path_pkl, min_wl, max_wl)
# Convolve spectrum with instrumental resolution
spc = spectral_convolution_ptr(wavelength, spectrum, pixel_dv)
# Create spline interpolation model
spline_model = splrep(wavelength, spc)
return wavelength * 1e7, spectrum, spline_model
[docs]
def download_star_spectrum(figure_tk, target_obj, folder_stellar_spc):
"""
Download and save a PHOENIX stellar spectrum if not already present.
Args:
figure_tk: Parent Tkinter widget for the stellar parameter dialog.
target_obj: Target object with stellar_effective_temperature attribute.
folder_stellar_spc (str): Directory for storing the stellar spectrum file.
Returns:
str or None: Path to the stellar spectrum pickle file, or None if cancelled.
"""
# Download stellar spectrum if missing
path_phoenixf = str(Path(folder_stellar_spc, "star_spectrum_default.pkl"))
if not os.path.exists(path_phoenixf):
print("Stellar spectrum not found. Requesting stellar parameters...")
from GUIBRUSHR.GUI.Input_Output_Panels.Input_Panels.TabPanels.FrameGenerationHRData.StellarDialog import StellarDialog # subprocess-lazy-import
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import ask_stellar_params
temperature, gravity, met = ask_stellar_params(
parent=figure_tk,
stellar_teff=target_obj.stellar_effective_temperature
)
if temperature is not None:
print(f"Downloading stellar spectrum: T_eff={temperature}K, log(g)={gravity}, [Fe/H]={met}")
# Create directory if it doesn't exist
os.makedirs(folder_stellar_spc, exist_ok=True)
# Download stellar spectrum files
file_wlen_phoenix = (
f"wget -P {folder_stellar_spc} "
f"ftp://phoenix.astro.physik.uni-goettingen.de/HiResFITS/"
f"PHOENIX-ACES-AGSS-COND-2011/Z{met}/"
f"lte{temperature}-{gravity}0-0.0.PHOENIX-ACES-AGSS-COND-2011-HiRes.fits"
)
file_spectrum_phoenix = (
f"wget -P {folder_stellar_spc} "
f"ftp://phoenix.astro.physik.uni-goettingen.de/HiResFITS/"
f"WAVE_PHOENIX-ACES-AGSS-COND-2011.fits"
)
print("Downloading spectrum files...")
os.system(file_wlen_phoenix)
os.system(file_spectrum_phoenix)
# Process downloaded files
path_spc = str(Path(
folder_stellar_spc,
f"lte{temperature}-{gravity}0-0.0.PHOENIX-ACES-AGSS-COND-2011-HiRes.fits"
))
path_wlen = str(Path(folder_stellar_spc, "WAVE_PHOENIX-ACES-AGSS-COND-2011.fits"))
path_pkl = str(Path(folder_stellar_spc, "star_spectrum_default.pkl"))
print("Processing and saving stellar spectrum...")
# Load and process spectrum data
spc = fits.open(path_spc)[0].data
wlen = fits.open(path_wlen)[0].data
star_model = {
"wavelength": wlen * 1e-8, # angstrom to cm
"spectrum": spc / np.pi
}
# Save processed spectrum
with open(path_pkl, "wb") as file:
pickle.dump(star_model, file)
print("Stellar spectrum saved successfully.")
return path_phoenixf
else:
print("Stellar parameter dialog cancelled. Aborting synthetic generation.")
return None
else:
print("Using existing stellar spectrum.")
return path_phoenixf