Source code for GUIBRUSHR.core.functions.statistics
"""
Statistics utility functions for GUIBRUSHR atmospheric retrieval.
This module contains functions for statistical analysis extracted from
general_functions.py:
- get_medians: Extract and process parameter statistics from MCMC chains.
- pre_plot: Prepare posterior chain statistics and labels for plotting.
- linear_solver: Solve a linear system using SVD.
- get_keys_by_value: Find the first key in a dict whose value contains a target.
- compute_transit_durations: Compute transit durations T14, T23, T12.
"""
import copy
import os
import pickle
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
from scipy import stats
from GUIBRUSHR.General_Constants.Classes.BestFitParametersAndResults import BestFitParametersAndResults
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
[docs]
def get_medians(
folder_retrieval_final: Union[str, Path],
arr_values,
remove_chain: bool,
corner_plot_or_simulation: bool,
best_fit_method: str,
n_samples: int = 500
) -> BestFitParametersAndResults:
"""
Extract and process parameter statistics from MCMC chains.
This function processes MCMC chain data to compute median values and other statistics
for fitted parameters. It handles chain rejection based on log-likelihood criteria and performs
burn-in removal.
Parameters:
folder_retrieval_final (Union[str, Path]): Path to the final retrieval folder
arr_values: Array containing parameter information
remove_chain (bool): Whether to remove chains based on log-likelihood criteria
corner_plot_or_simulation (bool): Whether to process data for corner plots
best_fit_method: Method used to extract best fit parameters
n_samples: Number of samples to use for random model extraction
Returns:
BestFitParametersAndResults: Object containing processed parameter statistics
"""
from scipy.stats import gaussian_kde
from scipy.signal import find_peaks
# Extract best-fit parameters
list_best_pars = arr_values[ConstantVariables.DICT_FILTER_TABLE["Fitted Params"]["index"]].split(",")
list_best_pars_labels = copy.copy(list_best_pars)
# Process molecule parameters
all_molec = ConstantVariables.ALL_MOLEC
for i, param in enumerate(list_best_pars):
name_par = param.replace("_", "")
if name_par in all_molec:
list_best_pars_labels[i] = name_par
# Extract fixed parameters
fix_params = arr_values[ConstantVariables.DICT_FILTER_TABLE["#Fixed Params"]["index"]].split(",")
lista_par_fix = []
value_fix = []
for elem in fix_params:
name, value = elem.split("=")
lista_par_fix.append(name.strip())
value_fix.append(float(value))
parsc = None
lhood = None
# IDL: chi2 (log-likelihood array, higher = better)
# Load chain data
pkl_fname = str(Path(folder_retrieval_final, "partial_prob_and_chain_burnin.pkl"))
try:
with open(pkl_fname, "rb") as f:
m = pickle.load(f)
if isinstance(m, dict):
parsc = m["parameters"]
# Backward compat: read both new key 'lhood' and old key 'chi2'
lhood = m.get("lhood", m.get("chi2")) # IDL: chi2
det_chain = m.get("det_chain")
else:
parsc = np.array(m[0])
lhood = np.array(m[1]) # IDL: chi2
except FileNotFoundError:
print(f"Chain data file not found at {pkl_fname}")
except Exception as e:
print(f"Error loading chain data: {str(e)}")
# Compute chain statistics on log-likelihood (IDL: chi2 / medchi2)
med_lhood = np.median(lhood) # IDL: medchi2
med_lhood_chains = np.median(lhood, axis=0) # IDL: medchi2_chains
err_lhood = stats.median_abs_deviation(lhood, axis=None) # IDL: err_chi2
# Identify bad chains (those with median log-likelihood far below overall median)
aa_chains = np.where(med_lhood_chains < med_lhood - 1.0 * err_lhood)
number_bad_chains = len(aa_chains[0])
n_chains_total = len(med_lhood_chains)
lhood_threshold = med_lhood - 1.0 * err_lhood
if number_bad_chains > 0:
action = "they will be removed" if remove_chain else "removal is OFF, so they are kept"
print(
f"Chain quality: {number_bad_chains}/{n_chains_total} chain(s) flagged as bad "
f"-> median log-likelihood below the threshold "
f"{lhood_threshold:.2f} (= median {med_lhood:.2f} - MAD {err_lhood:.2f}); {action}."
)
else:
print(
f"Chain quality: all {n_chains_total} chains good "
f"(every chain median log-likelihood above {lhood_threshold:.2f})."
)
# Select chains based on criteria
bb_chains = np.where(med_lhood_chains > med_lhood - 1.0 * err_lhood)[0] if remove_chain else np.where(med_lhood_chains > -np.inf)[0]
lhood = lhood[:, bb_chains.T] # IDL: chi2
# Determine burn-in index
burnndx = 0
for jj in range(lhood.shape[1]):
tmpndx = np.where(lhood[:, jj] > med_lhood)[0]
if len(tmpndx) > 0 and tmpndx[0] > burnndx:
burnndx = tmpndx[0]
# else:
# burnndx = 0
if burnndx > (lhood.shape[0] - 3):
burnndx = 0
n_steps_total = lhood.shape[0]
pct_burn = 100.0 * burnndx / n_steps_total if n_steps_total else 0.0
print(
f"Burn-in: discarding the first {burnndx} of {n_steps_total} steps "
f"({pct_burn:.1f}%); {n_steps_total - burnndx} steps kept per chain."
)
# Keep full chain (good chains, all steps) before burn-in removal for trace plot
parsc_full = parsc[:, :, bb_chains.T] # shape: (nfit, nsteps_total, nchains_good)
# Remove burn-in and reshape
parsc = parsc_full[:, burnndx:, :]
nchains_good = parsc.shape[2]
arr1 = parsc.reshape((parsc.shape[0], -1)).T
lhood = lhood[burnndx:, :] # IDL: chi2
lhood_monodimensional = lhood.reshape(-1).T # IDL: chi2_monodimensional
# Compute convergence statistics
from GUIBRUSHR.Retrieval.ExofastMCMC import exofast_gelmanrubin
converged, gelmanrubin, tz = exofast_gelmanrubin.exofast_gelmanrubin(parsc)
n_samples_post, n_params_post = arr1.shape
steps_per_chain = n_samples_post // nchains_good if nchains_good else 0
print(
f"Posterior sample: {n_samples_post} draws over {n_params_post} fitted parameters "
f"({nchains_good} chains x ~{steps_per_chain} post-burn-in steps each)."
)
# Convergence diagnostics. A chain set is declared converged only if every
# parameter has Gelman-Rubin R-hat < rhat_threshold AND effective sample
# size Tz > min_tz (same rule as exofast_gelmanrubin).
rhat_threshold = 1.01
min_tz = 1000
if converged == 1:
print(f"Convergence: PASSED - all R-hat < {rhat_threshold} and all Tz > {min_tz}.")
elif converged == 0:
print(f"Convergence: NOT reached - need all R-hat < {rhat_threshold} and all Tz > {min_tz}.")
else:
print("Convergence: not available (too few chains/steps to compute).")
if gelmanrubin is not None and tz is not None:
gr = np.atleast_1d(np.asarray(gelmanrubin, dtype=float))
tzv = np.atleast_1d(np.asarray(tz, dtype=float))
n_show = min(len(gr), len(list_best_pars))
print(f" Per-parameter mixing (R-hat target < {rhat_threshold}, Tz target > {min_tz}):")
print(f" {'parameter':<18}{'R-hat':>9}{'Tz':>12} status")
for i in range(n_show):
ok = (gr[i] < rhat_threshold) and (tzv[i] > min_tz)
print(f" {list_best_pars[i]:<18}{gr[i]:>9.4f}{tzv[i]:>12.1f} {'ok' if ok else 'FAIL'}")
worst_i = int(np.argmax(gr[:n_show]))
print(
f" Worst-mixing parameter: {list_best_pars[worst_i]} "
f"(R-hat={gr[worst_i]:.4f}, Tz={tzv[worst_i]:.1f}); "
f"min Tz over all params = {np.min(tzv):.1f}."
)
# Load and update likelihood
dir_f = folder_retrieval_final
max_likelihood_pkl_path = Path(dir_f, "max_likelihood.pkl")
try:
if os.path.exists(str(max_likelihood_pkl_path)):
with open(str(max_likelihood_pkl_path), "rb") as fo:
last_max_LH = pickle.load(fo)
print(f"Previous run max log-likelihood: {last_max_LH:.4f}")
else:
last_max_LH = np.nan
except Exception as e:
print(f"Error loading likelihood: {str(e)}")
last_max_LH = np.nan
current_max_likelihood = np.max(lhood_monodimensional)
try:
with open(str(Path(dir_f, "max_likelihood.pkl")), "wb") as fo:
pickle.dump(current_max_likelihood, fo)
except Exception as e:
print(f"Error saving likelihood: {str(e)}")
delta_lh = current_max_likelihood - last_max_LH
if np.isfinite(delta_lh):
print(
f"Current run max log-likelihood: {current_max_likelihood:.4f} "
f"(change vs previous run: {delta_lh:+.4f})."
)
else:
print(f"Current run max log-likelihood: {current_max_likelihood:.4f}.")
# Process eccentricity parameters for corner plots
if corner_plot_or_simulation:
h_ecc_arr = None
k_ecc_arr = None
# Handle h_ecc parameter
if "h_ecc" in list_best_pars:
h_ecc_arr = np.array(arr1[:, list_best_pars.index("h_ecc")])
elif "h_ecc" in lista_par_fix:
h_ecc = value_fix[lista_par_fix.index("h_ecc")]
h_ecc_arr = np.full(arr1.shape[0], h_ecc)
arr1 = np.insert(arr1, list_best_pars.index("k_ecc"), h_ecc_arr, axis=1)
list_best_pars.insert(list_best_pars.index("k_ecc"), "h_ecc")
# Handle k_ecc parameter
if "k_ecc" in list_best_pars:
k_ecc_arr = np.array(arr1[:, list_best_pars.index("k_ecc")])
elif "k_ecc" in lista_par_fix:
k_ecc = value_fix[lista_par_fix.index("k_ecc")]
k_ecc_arr = np.full(arr1.shape[0], k_ecc)
arr1 = np.insert(arr1, list_best_pars.index("h_ecc") + 1, k_ecc_arr, axis=1)
list_best_pars.insert(list_best_pars.index("h_ecc") + 1, "k_ecc")
# Convert to eccentricity and argument of periastron
if h_ecc_arr is not None and k_ecc_arr is not None:
arr1[:, list_best_pars.index("h_ecc")] = np.power(h_ecc_arr, 2) + np.power(k_ecc_arr, 2)
ecc_opi_par = np.arctan2(h_ecc_arr, k_ecc_arr)
ecc_opi_par[ecc_opi_par < 0] += 2 * np.pi
arr1[:, list_best_pars.index("k_ecc")] = ecc_opi_par
# Compute final statistics
if best_fit_method == "Median":
# With the median
best_fit_parameters = np.median(arr1, axis=0)
elif best_fit_method == "Max LH":
# With the max of likelihood
# MAP = maximum of the posterior
idx_MAP = np.argmax(lhood_monodimensional) # IDL: chi2_monodimensional
params_MAP = arr1[idx_MAP]
best_fit_parameters = np.array(params_MAP)
print("MAP estimate (maximum-a-posteriori value, per fitted parameter):")
for name, value in zip(list_best_pars, np.atleast_1d(params_MAP)):
print(f" {name:<18} = {value: .6e}")
elif best_fit_method == "KDE":
# # With KDE
n_params = arr1.shape[1]
best_fit_parameters = np.empty(n_params)
for i in range(n_params):
samples = arr1[:, i]
kde = gaussian_kde(samples)
x_grid = np.linspace(samples.min(), samples.max(), 1000)
posterior_density = kde(x_grid)
idx_MAP = np.argmax(posterior_density)
best_fit_parameters[i] = x_grid[idx_MAP]
print(f"Parameter {i}: MAP = {best_fit_parameters[i]:.4f}, Median = {np.median(samples):.4f}")
elif best_fit_method == "Max from param":
# With histogram maximum bin median
# For each parameter: find bin with most samples, then take median within that bin
n_params = arr1.shape[1]
best_fit_parameters = np.empty(n_params)
for i in range(n_params):
samples = arr1[:, i]
# Use Freedman-Diaconis rule for optimal bin width
# This automatically determines a good number of bins based on data statistics
counts, bin_edges = np.histogram(samples, bins='fd')
# Find the bin with maximum count
idx_max_bin = np.argmax(counts)
# Get samples within this bin
bin_left = bin_edges[idx_max_bin]
bin_right = bin_edges[idx_max_bin + 1]
samples_in_max_bin = samples[(samples >= bin_left) & (samples < bin_right)]
# Take median of samples in the maximum bin
if len(samples_in_max_bin) > 0:
best_fit_parameters[i] = np.median(samples_in_max_bin)
else:
# Fallback to overall median if no samples found (shouldn't happen)
best_fit_parameters[i] = np.median(samples)
print(f"Parameter {i}: Max bin median = {best_fit_parameters[i]:.4f}, "
f"Bin range = [{bin_left:.4f}, {bin_right:.4f}], "
f"N samples in bin = {len(samples_in_max_bin)}, "
f"Overall median = {np.median(samples):.4f}")
else:
n_total = arr1.shape[0]
sample_indices = np.random.choice(n_total, size=min(n_samples, n_total), replace=False)
best_fit_parameters = arr1[sample_indices, :]
#
# best_fit_parameters_obj =
best_fit_parameters_obj = BestFitParametersAndResults(
number_bad_chains,
last_max_LH,
current_max_likelihood,
lhood_monodimensional, # IDL: chi2_monodimensional
list_best_pars,
arr1,
best_fit_parameters,
value_fix,
lista_par_fix,
[], # lista_parametri
list_best_pars_labels,
burnndx=burnndx,
nchains=nchains_good,
arr_full=parsc_full,
)
return best_fit_parameters_obj
[docs]
def pre_plot(
folder_retrieval_final: Union[str, Path],
arr_values: np.ndarray,
remove_chain: [bool, int],
corner_plot: [bool, int],
best_fit_method : str
):
"""
Prepare posterior chain statistics and parameter labels for plotting.
Parameters:
folder_retrieval_final (Union[str, Path]): Path to the final retrieval folder
arr_values (np.ndarray): Array containing parameter information
remove_chain (bool): Whether to remove chains based on log-likelihood criteria
corner_plot (bool): Whether to process data for corner plots
best_fit_method (str): Method used to extract best fit parameters
Returns:
Tuple containing:
- int: Number of bad chains
- float: AIC value
- float: Last maximum likelihood
- float: Current maximum likelihood
- np.ndarray: Flattened log-likelihood values # IDL: chi2
- List[str]: Best-fit parameter names
- np.ndarray: Post-burn-in chain parameters (arr1)
- np.ndarray: Parameter medians
- List[float]: Fixed parameter values
- List[str]: Fixed parameter names
- str: Summary string
- str: Parameter names string
- List[str]: Parameter labels
- np.ndarray: Parameter errors
- int: Burn-in index (burnndx)
- int: Number of good chains
- np.ndarray: Full chain including burn-in, shape (nfit, nsteps_total, nchains)
"""
best_fit_parameters_obj = get_medians(folder_retrieval_final, arr_values, remove_chain, corner_plot, best_fit_method)
# Compute AIC
aic_var = 2 * len(best_fit_parameters_obj.list_best_pars) - 2 * best_fit_parameters_obj.current_max_likelihood
# Format summary strings
stringa = ""
error_list = np.zeros(len(best_fit_parameters_obj.list_best_pars))
for i, param in enumerate(best_fit_parameters_obj.list_best_pars):
error = np.std(best_fit_parameters_obj.arr1[:, i])
error_list[i] = error
sorted_arr = np.sort(best_fit_parameters_obj.arr1[:, i])
pos = int(len(sorted_arr) * 0.95)
upper_limits = sorted_arr[pos]
if corner_plot:
stringa += f"{param}: {best_fit_method}.={best_fit_parameters_obj.best_fit_parameters[i]:.6f}, E.+={error:.6f}, U.L.={upper_limits:.6f}\n"
else:
if best_fit_method == "From sampler":
value_str = np.median(best_fit_parameters_obj.best_fit_parameters[:, i])
pre_str = "Median "
else:
pre_str = ""
value_str = best_fit_parameters_obj.best_fit_parameters[i]
stringa += f"{pre_str}{best_fit_method} {param} = {value_str:.6f}\n"
# Create parameter labels for corner plots
stringa_par = ""
lista_parametri = []
if corner_plot:
for elem in best_fit_parameters_obj.list_best_pars_labels:
stringa_par += elem + ","
no_number_elem = re.sub(r'\d+', '', elem)
number = re.findall(r'\d+', elem)
if no_number_elem in ConstantVariables.LIST_MULTIPLE_PARAM:
string_form = ConstantVariables.DICT_LABELS[no_number_elem].replace("x", number[0])
else:
#TODO remove for compatibility
if elem == "csuo":
elem = "co_ratio"
string_form = ConstantVariables.DICT_LABELS[elem]
string_form = string_form.replace('r"', "", 1).replace('"', "")
lista_parametri.append(string_form)
stringa_par = stringa_par[:-1] # Remove trailing comma
print("Producing the plot...")
return (
best_fit_parameters_obj.number_bad_chains,
aic_var,
best_fit_parameters_obj.last_max_LH,
best_fit_parameters_obj.current_max_likelihood,
best_fit_parameters_obj.lhood_monodimensional, # IDL: chi2_monodimensional
best_fit_parameters_obj.list_best_pars,
best_fit_parameters_obj.arr1,
best_fit_parameters_obj.best_fit_parameters,
best_fit_parameters_obj.value_fix,
best_fit_parameters_obj.lista_par_fix,
stringa,
stringa_par,
lista_parametri,
error_list,
best_fit_parameters_obj.burnndx,
best_fit_parameters_obj.nchains,
best_fit_parameters_obj.arr_full,
)
[docs]
def linear_solver(
topca_l: np.ndarray,
res_l: np.ndarray
):
"""
Solve a linear system using Singular Value Decomposition (SVD).
This function solves the linear system Ax = b using SVD decomposition,
where A is the design matrix (res_l) and b is the target vector (topca_l).
The solution is computed in a numerically stable way using SVD.
Parameters:
topca_l (np.ndarray): Target vector b in the system Ax = b
res_l (np.ndarray): Design matrix A in the system Ax = b
Returns:
Results: The reconstructed array after solving the linear system
Error: Whether the linear system was solved successfully
"""
error = False
try:
# Perform SVD decomposition
uf, s, vh = np.linalg.svd(res_l, full_matrices=False)
# Compute intermediate matrices
uT = np.transpose(uf)
nD = np.diag(1 / s)
vHcT = np.transpose(vh.conj())
# Solve the system in steps:
# 1. U diag(s) Vh x = b <=> diag(s) Vh x = U.T b = c
c = np.dot(uT, topca_l)
# 2. diag(s) Vh x = c <=> Vh x = diag(1/s) c = w
w = np.dot(nD, c)
# 3. Vh x = w <=> x = Vh.H w
x = np.dot(vHcT, w)
# Reconstruct the array
return np.dot(res_l, x), error
except np.linalg.LinAlgError as e:
print(f"SVD computation failed: {str(e)}")
error = True
return None, error
[docs]
def get_keys_by_value(
dictionary: Dict[str, Any],
target_value: Any
) -> Optional[str]:
"""
Find the first key in a dictionary whose value contains the target value.
This function searches through the dictionary values to find the first key
whose value contains the target value. It's useful for reverse lookups in
dictionaries where values are collections.
Parameters:
dictionary (Dict[str, Any]): Dictionary to search through
target_value (Any): Value to search for in the dictionary values
Returns:
Optional[str]: The first key whose value contains target_value,
or None if no match is found
"""
try:
return next(
(key for key, value in dictionary.items() if target_value in value),
None
)
except TypeError:
# Handle case where values are not iterable
return None
[docs]
def compute_transit_durations(period, a_Rs, k, inc_deg, ecc=0.0, omega_planet=np.pi / 2):
"""
Compute transit durations T14, T23, T12 (=T34) for circular and eccentric orbits.
Parameters
----------
period : float
Orbital period in days.
a_Rs : float
Semi-major axis in units of stellar radii (a / R_star).
k : float
Planet-to-star radius ratio (Rp / Rs).
inc_deg : float
Orbital inclination in degrees.
ecc : float, optional
Orbital eccentricity (default 0 for circular).
omega_planet : float, optional
Argument of periastron of the PLANET in radians (default pi/2). The Winn (2010)
transit formulae use the STAR's argument of periastron, so this value is
converted internally to the stellar one (omega_star = omega_planet - pi).
Returns
-------
dict with keys:
'T14', 'T23', 'T12' : durations in days
'T14_hours', 'T23_hours', 'T12_hours' : durations in hours
'b' : impact parameter
'ecc_factor' : eccentricity correction factor (1.0 if circular)
'grazing' : True if grazing transit (T23 = 0)
'limph_T14', 'limph_T12' : phase limits = T / (2*P)
References
----------
Winn (2010), "Transits and Occultations", arXiv:1001.2010, Eq. 14-16.
Kipping (2010), MNRAS 407, 301, arXiv:1004.3819.
"""
inc = np.radians(inc_deg)
sin_i = np.sin(inc)
# The input is the PLANET's argument of periastron; the Winn (2010) transit
# formulae below use the STAR's, which differs by pi.
omega_star = omega_planet - np.pi
# Impact parameter
# Eccentric: b = (a/Rs) * cos(i) * (1 - e^2) / (1 + e*sin(omega_star))
# Circular: b = (a/Rs) * cos(i)
if ecc > 0:
ecc_factor_b = (1.0 - ecc ** 2) / (1.0 + ecc * np.sin(omega_star))
else:
ecc_factor_b = 1.0
b = a_Rs * np.cos(inc) * ecc_factor_b
if b > (1.0 + k):
raise ValueError(
f"No transit: b = {b:.4f} > 1 + k = {1 + k:.4f}."
)
# Winn (2010) Eq. 14-15
arg_T14_sq = (1.0 + k) ** 2 - b ** 2
arg_T23_sq = (1.0 - k) ** 2 - b ** 2
if arg_T14_sq < 0:
raise ValueError(
f"No transit: (1+k)^2 - b^2 = {arg_T14_sq:.6f} < 0"
)
grazing = False
if arg_T23_sq < 0:
grazing = True
arg_T23_sq = 0.0
sin_T14 = min((1.0 / a_Rs) * np.sqrt(arg_T14_sq) / sin_i, 1.0)
sin_T23 = min((1.0 / a_Rs) * np.sqrt(arg_T23_sq) / sin_i, 1.0)
T14_circ = (period / np.pi) * np.arcsin(sin_T14)
T23_circ = (period / np.pi) * np.arcsin(sin_T23)
# Eccentricity correction (Winn 2010, Eq. 16)
# Psi = sqrt(1 - e^2) / (1 + e*sin(omega_star))
if ecc > 0:
ecc_factor = np.sqrt(1.0 - ecc ** 2) / (1.0 + ecc * np.sin(omega_star))
else:
ecc_factor = 1.0
T14 = T14_circ * ecc_factor
T23 = T23_circ * ecc_factor
T12 = (T14 - T23) / 2.0
# T14 and T23 are durations symmetric about mid-transit, so their phase
# "limit" is the half-width (centre-to-contact): duration / (2*period),
# used downstream as a ±limit around phase 0.
limph_T14 = T14 / (2.0 * period)
limph_T23 = T23 / (2.0 * period)
# T12 is the ingress duration (one-sided), NOT a duration centred on phase 0.
# Its phase extent equals limph_T14 - limph_T23 = T12 / period, so there is
# no extra /2 here (a /2 would yield half the ingress width).
limph_T12 = T12 / period
print(f"T14={T14 * 24:.2f}hours, T23={T23 * 24:.2f}hours, T12={T12 * 24:.2f}hours, limph_T14={limph_T14:.2f}, limph_T12={limph_T12:.2f}, limph_T23={limph_T23:.2f}")
return {
'T14': T14,
'T23': T23,
'T12': T12,
'T14_hours': T14 * 24.0,
'T23_hours': T23 * 24.0,
'T12_hours': T12 * 24.0,
'b': b,
'ecc_factor': ecc_factor,
'grazing': grazing,
'limph_T14': limph_T14,
'limph_T12': limph_T12,
'limph_T23': limph_T23,
}