"""
Parameter handling for atmospheric retrieval.
Manages parameter unpacking, indexing, and categorization for the MCMC
retrieval system. Extracted from ModelData to isolate parameter logic.
"""
import numpy as np
import pandas as pd
from pathlib import Path
from re import sub
from typing import List, Optional, Tuple, Any
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.ParamForModel import ParamForModel, RANGE_SENTINEL
from GUIBRUSHR.core.io.retrieval_io import read_parameters
from GUIBRUSHR.core.types import (
ChainSchema,
ParameterManagementResult,
ParamRegistry,
ReadDfParametersResult,
)
[docs]
def get_param_array_initial_value(array, index, single_value=True):
"""
Get initial parameter value from array at specified index.
Args:
array: Parameter array
index: Index to retrieve
single_value: If True, return single value, else return array
Returns:
Initial parameter value or None if not found
"""
if array[index] is None:
return None
if single_value:
return array[index].get_starting_value()
else:
return array[index].get_starting_value_arr()
[docs]
class ParameterHandler:
"""
Handles parameter indexing, unpacking, and categorization.
This class manages the parameter array and provides methods to:
- Look up parameter indices by name
- Extract parameter values from arrays
- Process and categorize parameters from CSV or manual model input
- Reconstruct full parameter arrays from MCMC chain values
Attributes:
params_list: Ordered list of all parameter names.
list_multiple_param: Names of parameters that can have multiple values.
initial_param_array: Array of ParamForModel objects (one per parameter slot).
"""
[docs]
def __init__(self, params_list: List[str], list_multiple_param: List[str]):
self.params_list = params_list
self.list_multiple_param = list_multiple_param
self.initial_param_array = [None for _ in self.params_list]
[docs]
def get_index(self, param_name: str) -> int:
"""
Find index of parameter with given name in the parameters list.
Handles both exact parameter names and parameter names
with numeric suffixes (e.g., 'jitter1', 'jitter2' -> 'jitter').
Args:
param_name: Name of parameter to find
Returns:
Index of parameter in params_list
"""
if param_name in self.params_list:
return self.params_list.index(param_name)
base_param_name = sub(r"\d+", "", param_name)
return int(self.params_list.index(base_param_name))
[docs]
def get_value(self, params, name: str, single_value: bool = True):
"""
Extract value of named parameter from params array.
Args:
params: Array of parameter objects
name: Parameter name to extract
single_value: If True, return single value, else return array
Returns:
Parameter value(s) or None if parameter not found
"""
index = self.get_index(name)
if params[index] is None:
return None
if single_value:
return params[index].get_retrieval_value()
return params[index].get_retrieval_value_arr()
[docs]
def parameter_management(
self,
name_for_retrieval,
in_bestpars_param,
scale_param,
constant_vmr_param,
range_min_param,
range_max_param,
molec_param,
molec_formula_param,
mass_param,
name_for_list_molec,
rayleigh_species_param,
value_during_retrieval,
sigma_prior_param,
isotope,
opacity_name_lr,
scale,
rayleigh_species,
line_species,
line_species_isotope,
line_species_complete_name_hr,
line_species_complete_name_lr,
list_bestpars,
list_bestpars_initial_value,
list_fixed,
list_condensed_molecules
):
"""
Manage parameter processing and categorization for the retrieval system.
Handles the logic of processing individual parameters and categorizing
them into appropriate lists based on their properties.
Args:
name_for_retrieval: Parameter name used in retrieval
in_bestpars_param: Whether parameter is in best parameters list
scale_param: Scale factor for parameter
constant_vmr_param: Whether VMR is constant
range_min_param: Minimum parameter range
range_max_param: Maximum parameter range
molec_param: Whether parameter is molecular
molec_formula_param: Molecular formula
mass_param: Molecular mass
name_for_list_molec: Name for molecule list
rayleigh_species_param: Whether parameter is Rayleigh species
value_during_retrieval: Parameter value during retrieval
sigma_prior_param: Sigma prior for parameter
isotope: Isotope information
opacity_name_lr: Opacity name for low resolution
scale: Scale factors list
rayleigh_species: Rayleigh species list
line_species: Line species list
line_species_isotope: Isotope list
line_species_complete_name_hr: Complete species names list HR
line_species_complete_name_lr: Complete species names list LR
list_bestpars: Best parameters list
list_bestpars_initial_value: Initial values list
list_fixed: Fixed parameters list
list_condensed_molecules: Condensed molecules list
Returns:
Tuple of updated parameter lists
"""
if name_for_retrieval not in self.params_list:
name = sub(r"\d+", "", name_for_retrieval)
else:
name = name_for_retrieval
if self.initial_param_array[self.get_index(name)] is None:
self.initial_param_array[self.get_index(name)] = ParamForModel(
name, in_bestpars_param, scale_param, constant_vmr_param,
range_min_param, range_max_param, molec_param, molec_formula_param,
mass_param, name_for_list_molec, isotope,
)
if name in self.list_multiple_param:
self.initial_param_array[self.get_index(name)].append_elem(
value_during_retrieval, sigma_prior_param
)
else:
self.initial_param_array[self.get_index(name)].update_starting_value(
value_during_retrieval, sigma_prior_param
)
if in_bestpars_param:
list_bestpars.append(name_for_retrieval)
list_bestpars_initial_value.append(value_during_retrieval)
scale.append(scale_param)
else:
list_fixed.append(name_for_retrieval)
if int(rayleigh_species_param) == 2:
rayleigh_species.append(name_for_retrieval)
if molec_param and name.replace("_", "") in ConstantVariables.ALL_MOLEC:
line_species.append(name_for_retrieval)
line_species_isotope.append(isotope)
line_species_complete_name_hr.append(name_for_list_molec)
line_species_complete_name_lr.append(opacity_name_lr)
elif int(rayleigh_species_param) == 1:
rayleigh_species.append(name_for_retrieval)
elif molec_param and name.replace("_", "") in ConstantVariables.ALL_MOLEC:
line_species.append(name_for_retrieval)
line_species_isotope.append(isotope)
line_species_complete_name_hr.append(name_for_list_molec)
line_species_complete_name_lr.append(opacity_name_lr)
elif name_for_retrieval in ConstantVariables.ALL_CONDENSED_MOLEC:
list_condensed_molecules.append(name_for_retrieval)
return (
scale,
rayleigh_species,
line_species,
line_species_isotope,
line_species_complete_name_hr,
line_species_complete_name_lr,
list_bestpars,
list_bestpars_initial_value,
list_fixed,
list_condensed_molecules
)
[docs]
def add_param_manual_model(
self,
name,
value_during_retrieval,
constant_VMR,
mass_elem,
molec_formula,
name_for_list_molec,
scale,
rayleigh,
rayleigh_species,
line_species,
line_species_isotope,
line_species_complete_name_hr,
line_species_complete_name_lr,
list_bestpars,
list_bestpars_initial_value,
isotope,
list_fixed,
list_condensed_molecules
):
"""
Add a parameter from manual model to the parameter management system.
Args:
name: Parameter name
value_during_retrieval: Parameter value for retrieval
constant_VMR: Whether VMR is constant
mass_elem: Molecular mass
molec_formula: Molecular formula
name_for_list_molec: Name for molecule list
scale: Scale factor list
rayleigh: Whether Rayleigh scattering applies
rayleigh_species: List of Rayleigh species
line_species: List of line species
line_species_isotope: List of isotopes
line_species_complete_name_hr: List of complete species names HR
line_species_complete_name_lr: List of complete species names LR
list_bestpars: List of best parameters
list_bestpars_initial_value: List of initial values
isotope: Isotope information
list_fixed: List of fixed parameters
list_condensed_molecules: List of condensed molecules
Returns:
Tuple of updated parameter lists
"""
in_best_pars = 0
scale_param = 0
sigma_prior_param = None
is_molec = mass_elem is not None
range_param = np.inf
opacity_name_lr = isotope
return self.parameter_management(
name, in_best_pars, scale_param, constant_VMR,
-range_param, range_param, is_molec, molec_formula,
mass_elem, name_for_list_molec, rayleigh,
value_during_retrieval, sigma_prior_param, isotope, opacity_name_lr,
scale, rayleigh_species, line_species, line_species_isotope,
line_species_complete_name_hr, line_species_complete_name_lr, list_bestpars,
list_bestpars_initial_value, list_fixed, list_condensed_molecules
)
[docs]
def read_df_parameters(self, path_df_parameters: str):
"""
Read the fitting-parameters table for the retrieval.
The path argument is treated as the *folder* containing the parameters
file (the filename component is ignored). The reader prefers
``df_parameters.yaml`` and falls back to legacy ``df_parameters.csv``,
lazily writing a YAML sibling on first read for forward migration.
Args:
path_df_parameters: Path to the parameters file (only its parent
directory is used). For backwards compatibility the signature
still accepts the legacy file path.
Returns:
Tuple containing all processed parameter information
"""
rayleigh_species = []
line_species = []
line_species_isotope = []
line_species_complete_name_hr = []
line_species_complete_name_lr = []
list_condensed_molecules = []
list_bestpars = []
list_bestpars_initial_value = []
list_fixed = []
scale = []
mass_vector = []
folder = Path(path_df_parameters).parent
df_parameters = read_parameters(folder)
df_parameters = df_parameters.where(pd.notnull(df_parameters), None)
for _, row in df_parameters.iterrows():
name_for_retrieval = row.get('name')
name_for_list_molec = row.get('name')
# Backward compatibility: the legacy per-night noise param "beta" was
# renamed to "beta_HR" (HR-only). Remap old retrieval rows
# (beta, beta0, beta1, ...) to beta_HR(0, 1, ...) at read time so they
# are interpreted as the high-resolution noise scale. beta_HR / beta_LR
# rows (base name != "beta") are left untouched.
if sub(r"\d+", "", str(name_for_retrieval)) == "beta":
name_for_retrieval = "beta_HR" + str(name_for_retrieval)[len("beta"):]
name_for_list_molec = name_for_retrieval
molec_param = int(row.get('molec'))
value_during_retrieval = float(row.get('value'))
scale_param = float(row.get('scale'))
range_min_param = float(row.get('range_min'))
range_max_param = float(row.get('range_max'))
rayleigh_species_param = int(row.get('rayleigh_species'))
in_bestpars_param = int(row.get('in_bestpars'))
mass_param = row.get('mass')
sigma_prior_param = row.get('sigma_prior')
molec_formula_param = row.get('molec_formula')
constant_vmr_param = row.get('constant_vmr')
isotope = row.get('isotope')
opacity_name_lr = row.get('opacity_name_lr') or isotope
if mass_param is None or np.isnan(mass_param):
mass_param = None
else:
mass_param = float(mass_param)
mass_vector.append(mass_param)
if sigma_prior_param is None or np.isnan(sigma_prior_param):
sigma_prior_param = None
else:
sigma_prior_param = float(sigma_prior_param)
if molec_param == 0:
molec_formula_param = None
elif molec_param == 1:
name_for_retrieval = molec_formula_param
(
scale, rayleigh_species, line_species, line_species_isotope,
line_species_complete_name_hr, line_species_complete_name_lr, list_bestpars,
list_bestpars_initial_value, list_fixed, list_condensed_molecules
) = self.parameter_management(
name_for_retrieval, in_bestpars_param,
scale_param, constant_vmr_param,
range_min_param, range_max_param,
molec_param, molec_formula_param,
mass_param, name_for_list_molec, rayleigh_species_param,
value_during_retrieval, sigma_prior_param,
isotope, opacity_name_lr, scale, rayleigh_species,
line_species, line_species_isotope,
line_species_complete_name_hr, line_species_complete_name_lr, list_bestpars,
list_bestpars_initial_value, list_fixed, list_condensed_molecules
)
scale = np.array(scale)
if len(list_condensed_molecules) == 0:
list_condensed_molecules = None
# Build per-parameter prior-bound vectors aligned 1:1 with ``scale``
# (i.e. only ``in_bestpars=1`` rows). Done as a separate post-pass on
# ``df_parameters`` rather than threaded through ``parameter_management``
# so the (already long) parameter_management signature stays untouched.
if "in_bestpars" in df_parameters.columns:
mask_bestpars = df_parameters["in_bestpars"].fillna(0).astype(int) == 1
range_min_vector = (
df_parameters.loc[mask_bestpars, "range_min"].astype(float).to_numpy()
)
range_max_vector = (
df_parameters.loc[mask_bestpars, "range_max"].astype(float).to_numpy()
)
# Unbounded-range sentinels: -999999 means -inf, 999999 means +inf
# (same convention as ParamForModel)
range_min_vector[range_min_vector == -RANGE_SENTINEL] = -np.inf
range_max_vector[range_max_vector == RANGE_SENTINEL] = np.inf
else:
range_min_vector = np.array([], dtype=float)
range_max_vector = np.array([], dtype=float)
# Return a NamedTuple so positional unpacking still works (legacy
# callers, pin tests) but downstream consumers (ModelSetup) can use
# named attribute access and let any future field reorder fail loud.
return ReadDfParametersResult(
mass_vector=mass_vector,
scale=scale,
rayleigh_species=rayleigh_species,
line_species=line_species,
line_species_isotope=line_species_isotope,
line_species_complete_name_hr=line_species_complete_name_hr,
line_species_complete_name_lr=line_species_complete_name_lr,
list_bestpars=list_bestpars,
list_bestpars_initial_value=list_bestpars_initial_value,
list_fixed=list_fixed,
list_condensed_molecules=list_condensed_molecules,
range_min_vector=range_min_vector,
range_max_vector=range_max_vector,
)
# ------------------------------------------------------------------
# Phase-1 typed accessors (additive; existing tuple API unchanged).
# ------------------------------------------------------------------
[docs]
def build_typed_result(
self,
scale,
rayleigh_species,
line_species,
line_species_isotope,
line_species_complete_name_hr,
line_species_complete_name_lr,
list_bestpars,
list_bestpars_initial_value,
list_fixed,
list_condensed_molecules,
mass_vector=None,
) -> "ParameterManagementResult":
"""Wrap legacy lists into a ``ParameterManagementResult`` dataclass.
Used during Phase 1 to expose the new typed view without altering
the read_df_parameters / parameter_management call signatures.
Builds a ``ParamRegistry`` from ``self.initial_param_array`` and a
``ChainSchema`` (with ``multi_expansion``) from ``list_bestpars``.
"""
registry = ParamRegistry(self.params_list)
registry.populate_from_list(self.initial_param_array)
multi_expansion: dict[str, list[str]] = {}
for bp in list_bestpars:
if bp not in self.params_list:
base = sub(r"\d+", "", bp)
if base in self.list_multiple_param:
multi_expansion.setdefault(base, []).append(bp)
multi_expansion_frozen = {
k: tuple(v) for k, v in multi_expansion.items()
}
scale_arr = (
scale if isinstance(scale, np.ndarray)
else np.asarray(scale, dtype=float)
)
schema = ChainSchema(
names=tuple(list_bestpars),
initial_values=np.asarray(list_bestpars_initial_value, dtype=float),
scale_vector=scale_arr,
multi_expansion=multi_expansion_frozen,
)
mass_arr = (
None if mass_vector is None
else np.asarray(mass_vector, dtype=float)
)
return ParameterManagementResult(
registry=registry,
schema=schema,
rayleigh_species=list(rayleigh_species),
line_species=list(line_species),
line_species_isotope=list(line_species_isotope),
line_species_complete_name_hr=list(line_species_complete_name_hr),
line_species_complete_name_lr=list(line_species_complete_name_lr),
list_fixed=list(list_fixed),
list_condensed_molecules=(
None if list_condensed_molecules is None
else list(list_condensed_molecules)
),
mass_vector=mass_arr,
)
[docs]
def read_df_parameters_typed(
self, path_df_parameters: str
) -> "ParameterManagementResult":
"""Typed wrapper around ``read_df_parameters``.
Returns a ``ParameterManagementResult`` containing both the new
``ParamRegistry`` / ``ChainSchema`` view and the legacy lists for
Phase-1 callers. Internally calls ``read_df_parameters`` so the
canonical build logic is not duplicated.
"""
(mass_vector, scale, rayleigh_species, line_species,
line_species_isotope, line_species_complete_name_hr,
line_species_complete_name_lr, list_bestpars,
list_bestpars_initial_value, list_fixed,
list_condensed_molecules) = self.read_df_parameters(path_df_parameters)
return self.build_typed_result(
scale=scale,
rayleigh_species=rayleigh_species,
line_species=line_species,
line_species_isotope=line_species_isotope,
line_species_complete_name_hr=line_species_complete_name_hr,
line_species_complete_name_lr=line_species_complete_name_lr,
list_bestpars=list_bestpars,
list_bestpars_initial_value=list_bestpars_initial_value,
list_fixed=list_fixed,
list_condensed_molecules=list_condensed_molecules,
mass_vector=mass_vector,
)
[docs]
def create_param_full(self, newpars_chain, bestpars_data):
"""
Create full parameter array from chain parameter values.
Reconstructs ``initial_param_array`` (in place) by mapping each
entry of ``newpars_chain`` to its canonical slot via the explicit
chain-name -> base-name routing table built from
``bestpars_data.list_bestpars`` and ``self.list_multiple_param``.
The routing table eliminates the substring scan that the legacy
implementation used to attribute consecutive chain values to a
multi-parameter slot. ``"f_rot"``-style base names that happen to
be substrings of other base names can no longer collide.
Args:
newpars_chain: Array of parameter values from MCMC chain.
bestpars_data: Bestpars object exposing ``list_bestpars`` and
``nfit`` (length of ``newpars_chain``).
Returns:
``self.initial_param_array`` with ``value_in_retrieval`` /
``value_arr_in_retrieval`` populated for every active slot.
"""
param_full = self.initial_param_array
list_bestpars = bestpars_data.list_bestpars
# Routing table: chain index -> (base param name, is_multi)
# Built once per call from list_bestpars (canonical chain order)
# and self.list_multiple_param (the static set of multi parameters).
chain_to_base: List[Tuple[str, bool]] = []
for chain_name in list_bestpars:
if chain_name in self.params_list:
chain_to_base.append((chain_name, False))
else:
base = sub(r"\d+", "", chain_name)
if base not in self.list_multiple_param:
raise ValueError(
f"create_param_full: chain entry {chain_name!r} stripped "
f"to {base!r}, which is not in list_multiple_param. "
"Multi-parameter naming convention violated."
)
chain_to_base.append((base, True))
# Group consecutive chain indices by their base name (multi-params
# produce N consecutive entries that all map back to one slot).
multi_buckets: dict[str, list] = {}
for k, (base, is_multi) in enumerate(chain_to_base):
if is_multi:
multi_buckets.setdefault(base, []).append(newpars_chain[k])
# Walk the canonical slot order and assign values. We still iterate
# by canonical index so that any drift between canonical slot order
# and chain order is loud rather than silent.
chain_param_index = 0
assigned_multi: set[str] = set()
for i in range(len(param_full)):
slot = param_full[i]
if slot is None or not slot.status:
continue
if slot.starting_value_variable is not None:
base, is_multi = chain_to_base[chain_param_index]
assert not is_multi and slot.name == base, (
f"create_param_full index misalignment at canonical slot {i}: "
f"param_full[i].name={slot.name!r} but chain entry "
f"{chain_param_index} maps to base {base!r} (is_multi={is_multi})"
)
slot.value_in_retrieval = newpars_chain[chain_param_index]
chain_param_index += 1
else:
if slot.name in assigned_multi:
continue
values = multi_buckets.get(slot.name)
if values is None:
raise ValueError(
f"create_param_full: multi slot {slot.name!r} has no chain "
"entries. Inconsistent list_bestpars vs initial_param_array."
)
slot.value_arr_in_retrieval = list(values)
assigned_multi.add(slot.name)
chain_param_index += len(values)
return param_full