"""
Retrieval class for astronomical data processing and model calculations.
This module contains the Retrieval class which handles the initialization
and configuration of retrieval parameters for astronomical observations.
"""
import numpy as np
# Allowed values for ``init_mode`` (chain-init scatter strategy after MAP).
# Order matters only for human-facing GUI dropdowns; semantic comparisons
# use set membership.
_VALID_INIT_MODES = ("isotropic", "diagonal", "correlated")
# Allowed values for ``mcmc_sampler``. "DE-MC" is the legacy byte-identical
# sampler; "Snooker" selects DE-MCzs (ter Braak and Vrugt 2008). The canonical
# (mixed-case) labels are what the rest of the pipeline compares against.
_VALID_SAMPLERS = ("DE-MC", "Snooker")
# Case-insensitive lookup canonical label. The config loader lowercases string
# values (ModelSetup._gi_str), so "snooker"/"de-mc" must map back to the
# canonical labels rather than being rejected.
_SAMPLER_CANONICAL = {s.lower(): s for s in _VALID_SAMPLERS}
[docs]
class Retrieval:
"""
A class to handle retrieval parameters and configuration for astronomical data processing.
This class manages various parameters related to atmospheric retrieval models,
including chemistry settings, temperature formatting, telluric removal methods,
and output configurations.
"""
[docs]
def __init__(
self,
scale,
chemistry=None,
ecc_opi=None,
tell_rm_method=None,
model_reprocessing=None,
format_temperature=None,
table_output_file=None,
maxsteps=None,
order_sel=None,
id_process=None,
path_results=None,
retrieval_model=None,
manual_model_obj=None,
total_nights=0,
total_hr_instruments=0,
standardize_PCA="From pkl",
mask_phase_enabled=False,
mask_night_index=None,
min_phase=None,
max_phase=None,
mask_inside_range=False,
transit_limit_mode="T14",
nthin=1,
init_scatter_factor=2.0,
epsilon_scale_divisor=100.0,
use_map_optimizer=False,
map_maxiter_global=40,
map_popsize=10,
map_tol_global=1.0e-2,
map_maxiter_local=150,
init_mode="isotropic",
mcmc_sampler="DE-MC",
range_min_vector=None,
range_max_vector=None,
gelman_rubin_threshold=1.01,
maxnpass=6,
save_chain_debug=True,
):
"""
Initialize the Retrieval object with configuration parameters.
Args:
scale: Scale vector parameters for the retrieval process
chemistry: Chemistry configuration for atmospheric modeling
ecc_opi: Eccentricity and orbital parameters information
tell_rm_method: Method for telluric removal
model_reprocessing: Model reprocessing configuration
format_temperature: Temperature formatting settings
table_output_file: Output file path for tabulated results
maxsteps: Maximum number of steps for the retrieval algorithm
order_sel: Order selection parameters
id_process: Process identifier
path_results: Path where results will be stored
retrieval_model: The retrieval model to be used
manual_model_obj: Manual model object containing pre-configured parameters
standardize_PCA: PCA standardization mode (default: "From pkl")
total_nights: Total number of observation nights (default: 0)
total_hr_instruments: Total number of HR instruments used (default: 0)
mask_phase_enabled: Flag indicating whether to apply a mask to the phase data (default: False)
mask_night_index: Index of the night to be masked (default: None)
min_phase: Minimum phase value to be masked (default: None)
max_phase: Maximum phase value to be masked (default: None)
mask_inside_range: Flag indicating whether to mask data inside the specified phase range (default: False)
nthin: Thinning factor — proposals per saved MCMC step (default: 1, no thinning)
init_scatter_factor: Multiplicative factor for initial chain scatter
(default: 2.0, G-R overdispersed starts)
epsilon_scale_divisor: Divisor in the uniform DE-MCMC epsilon perturbation
epsilon = (U-0.5)*2*scale/epsilon_scale_divisor (ter Braak 2006 §2.3;
IDL/EXOFAST uses 100). Must be > 0. Default 100.0.
use_map_optimizer: When True, run a bounded MAP optimization before
DE-MCMC and use the optimum as list_bestpars_initial_value
(default: False, keep the user-provided starting point).
map_maxiter_global: Max iterations for the differential evolution
stage of the MAP optimizer (default: 40).
map_popsize: Population size multiplier for differential evolution
(population = popsize * nfit) (default: 10).
map_tol_global: Relative tolerance for differential evolution
early-stop (default: 1e-2).
map_maxiter_local: Max iterations for the Nelder-Mead polish stage
of the MAP optimizer; set to 0 to skip the polish (default: 150).
gelman_rubin_threshold: Upper bound on R-hat used by ``exofast_gelmanrubin``
to declare chain convergence (default 1.01; canonical Gelman-Rubin
cutoff). Must be > 1.0. Raise (e.g. 1.05, 1.1) for faster but less
strict stopping on hard retrievals; lower (e.g. 1.005) to demand
tighter mixing.
maxnpass: Number of consecutive Gelman-Rubin checks that must all pass
before convergence is declared and the run stops (default 6, the
canonical EXOFAST/Ford 2006 value). A single failed check resets the
counter. Must be an integer >= 1. Lower (e.g. 2-3) for faster
stopping; raise to demand more sustained mixing.
save_chain_debug: When True (default), the MCMC driver writes the
per-step JSONL stream to ``{path_results}/debug/retrieval_debug.jsonl``
so the GUI "Analyse Chains" tool can plot acceptance / proposal
diagnostics. When False, no debug folder is created — useful for
long production runs where the JSONL would balloon to GBs.
Note: The MAP optimizer reuses the MCMC seed via SeedSequence.spawn(),
so reproducibility is controlled by the existing 'seed' field. There
is intentionally no separate map_seed.
"""
# If a manual model object is provided, extract all parameters from it
# This allows for easy parameter transfer from existing model configurations
if manual_model_obj is not None:
model_reprocessing = manual_model_obj.model_reprocessing
standardize_PCA = manual_model_obj.standardize_PCA
format_temperature = manual_model_obj.format_temperature
chemistry = manual_model_obj.chemistry
ecc_opi = manual_model_obj.ecc_opi
tell_rm_method = manual_model_obj.tell_rm_method
table_output_file = manual_model_obj.table_output_file
order_sel = manual_model_obj.order_sel
maxsteps = manual_model_obj.maxsteps
id_process = manual_model_obj.id_process
path_results = manual_model_obj.path_results
retrieval_model = manual_model_obj.retrieval_model
# Assign all parameters to instance variables
# Model processing and formatting settings
self.model_reprocessing = model_reprocessing
self.standardize_PCA = standardize_PCA
# Normalize old PT profile names to new ones for backward compatibility
_pt_fallback = {"isot": "Isothermal", "madhu": "Madhusudhan", "guillot": "Guillot", "personalized": "DoubleIsothermal"}
self.format_temperature = _pt_fallback.get(format_temperature, format_temperature)
# Atmospheric and orbital parameters
self.chemistry = chemistry
self.eccentricity = ecc_opi # Note: renamed from ecc_opi for clarity
# Processing methods and configurations
self.tell_rm_method = tell_rm_method
self.table_output_file = table_output_file
self.scale_vector_params = scale
self.order_sel = order_sel
# Algorithm and process control parameters
self.maxsteps = maxsteps
self.id_process = id_process
# Output and model settings
self.path_results = path_results
self.retrieval_model = retrieval_model
# Observation statistics
self.total_nights = total_nights
self.total_hr_instruments = total_hr_instruments
self.save_chain_debug = bool(save_chain_debug)
self.mask_phase_enabled = mask_phase_enabled
self.mask_night_index = mask_night_index
self.min_phase = min_phase
self.max_phase = max_phase
self.mask_inside_range = mask_inside_range
self.transit_limit_mode = transit_limit_mode
self.nthin = nthin
self.init_scatter_factor = init_scatter_factor
self.epsilon_scale_divisor = epsilon_scale_divisor
self.use_map_optimizer = use_map_optimizer
self.map_maxiter_global = map_maxiter_global
self.map_popsize = map_popsize
self.map_tol_global = map_tol_global
self.map_maxiter_local = map_maxiter_local
# Gelman-Rubin R-hat threshold for convergence (canonical default 1.01).
# Validated here so an invalid YAML entry fails loud rather than silently
# disabling the stop criterion.
gelman_rubin_threshold = float(gelman_rubin_threshold)
if gelman_rubin_threshold <= 1.0:
raise ValueError(
f"gelman_rubin_threshold must be > 1.0 (got {gelman_rubin_threshold}). "
f"R-hat is bounded below by 1 by construction; values <= 1.0 can never "
f"be reached and would prevent convergence from ever being declared."
)
self.gelman_rubin_threshold = gelman_rubin_threshold
# Number of consecutive G-R passes required to stop (canonical default 6).
# Validated here so an invalid YAML entry fails loud rather than silently
# changing the stop criterion.
maxnpass = int(maxnpass)
if maxnpass < 1:
raise ValueError(
f"maxnpass must be an integer >= 1 (got {maxnpass}). It is the number "
f"of consecutive Gelman-Rubin checks required to declare convergence."
)
self.maxnpass = maxnpass
# Chain-initialisation strategy after MAP (default ``"isotropic"`` =
# legacy byte-identical behaviour). Validated at parse time so a typo
# in df_general_info.yaml fails loud instead of silently picking the
# legacy path.
if init_mode not in _VALID_INIT_MODES:
raise ValueError(
f"init_mode={init_mode!r} not in {_VALID_INIT_MODES}. "
f"Set df_general_info.init_mode to one of these strings."
)
self.init_mode = init_mode
# MCMC sampler selection. Default ``"DE-MC"`` is byte-identical to the
# legacy DE-MC path; ``"Snooker"`` enables DE-MCzs. Validated at parse
# time so a typo in df_general_info.yaml fails loud instead of silently
# falling back to DE-MC.
# Canonicalise case-insensitively so the lowercased config value and the
# mixed-case GUI label both resolve to the canonical label that the
# worker (`sampler == "Snooker"`) and exofast loop compare against.
sampler_key = str(mcmc_sampler).strip().lower()
if sampler_key not in _SAMPLER_CANONICAL:
raise ValueError(
f"mcmc_sampler={mcmc_sampler!r} not in {_VALID_SAMPLERS} "
f"(case-insensitive). Set df_general_info.mcmc_sampler to one "
f"of these strings."
)
self.mcmc_sampler = _SAMPLER_CANONICAL[sampler_key]
# Per-parameter prior bounds, aligned 1:1 with ``scale``. Required by
# map_optimizer.compute_hessian_at_map() for the σ_p cap on orphan
# dimensions; safe to leave empty when init_mode == "isotropic".
self.range_min_vector_params = (
np.asarray(range_min_vector, dtype=float)
if range_min_vector is not None
else np.array([], dtype=float)
)
self.range_max_vector_params = (
np.asarray(range_max_vector, dtype=float)
if range_max_vector is not None
else np.array([], dtype=float)
)