#!/usr/bin/env python3
"""
Main script for running atmospheric retrieval using MCMC sampling.
This script serves as the entry point for performing atmospheric retrieval
analysis on exoplanet spectra. It initializes a ModelData object with the
provided parameters and runs Monte Carlo Markov Chain (MCMC) sampling using
the exofast_demc algorithm.
Usage:
python main.py <path_params> <path_df> <table_output_file> <id_process> <path_default> <moresteps>
Args:
path_params: Path to parameter configuration file
path_df: Path to dataframe file containing observational data
table_output_file: Path for output table file
id_process: Process ID for parallel execution
path_default: Default working directory path
moresteps: Boolean string ("True"/"False") to continue from previous run
Author: GUIBRUSHR Team
License: [License information if available]
"""
import os
import time
import sys
# Configure threading for numerical libraries
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1" # export MKL_NUM_THREADS=6
os.environ["VECLIB_MAXIMUM_THREADS"] = "1" # export VECLIB_MAXIMUM_THREADS=4
os.environ["NUMEXPR_NUM_THREADS"] = "1" # export NUMEXPR_NUM_THREADS=6
[docs]
def parse_command_line_arguments():
"""
Parse and validate command line arguments.
Returns
-------
tuple
A tuple containing (path_params, path_df, table_output_file,
id_process, path_default, moresteps_bool)
Raises
------
SystemExit
If incorrect number of arguments provided
"""
if len(sys.argv) != 7:
print(f"Usage: {sys.argv[0]} <path_params> <path_df> <table_output_file> "
f"<id_process> <path_default> <moresteps>")
sys.exit(1)
# Extract command line arguments
path_params = sys.argv[1]
path_df = sys.argv[2]
table_output_file = sys.argv[3]
id_process = sys.argv[4]
path_default = sys.argv[5]
moresteps_str = sys.argv[6]
# Convert moresteps string to boolean
moresteps_bool = moresteps_str == "True"
return (path_params, path_df, table_output_file,
id_process, path_default, moresteps_bool)
[docs]
def create_model_object(path_params, path_df, id_process, table_output_file, path_default):
"""
Create and configure the ModelData object for retrieval.
This function initializes the ModelData object with all necessary
parameters for atmospheric retrieval analysis.
Parameters
----------
path_params : str
Path to parameter configuration file
path_df : str
Path to dataframe file containing observational data
id_process : str
Process ID for parallel execution
table_output_file : str
Path for output table file
path_default : str
Default working directory path
Returns
-------
ModelData.ModelData
Configured ModelData object ready for retrieval
"""
from GUIBRUSHR.Retrieval.ModelCalculation import ModelData
model_obj = ModelData.ModelData(
path_params=path_params,
path_df=path_df,
id_process=id_process,
table_output_file=table_output_file,
model_type="Retrieval",
lbl_sampling_hr=None,
lbl_sampling_lr=None,
range_min=None,
range_max=None,
nlayers=None,
manual_model_obj=None,
load_new_opacities=True,
path_default=path_default,
)
return model_obj
[docs]
def setup_output_directory(model_obj):
"""
Create the output directory for retrieval results.
Parameters
----------
model_obj : ModelData.ModelData
The model object containing path information
"""
output_dir = model_obj.retrieval_data.path_results
os.system(f"mkdir -p {output_dir}")
[docs]
def run_map_warm_start(model_obj, moresteps):
"""
Run MAP optimization before MCMC, if enabled.
When model_obj.retrieval_data.use_map_optimizer is True, runs a bounded
MAP search and overwrites list_bestpars_initial_value so that chains
start around the posterior mode. When False, this is a no-op.
Skipped entirely on resume (moresteps=True): a resumed run loads the
previous chains from partial_prob_and_chain_burnin.pkl and continues from
each chain's last position, so the first-step chain initialization is bypassed.
The MAP point and Hessian-derived scale artifacts only ever feed that
initialization, so running MAP on resume would just discard the result -
wasted compute. The chains instead carry on from their last best-fit
positions, and sampling keeps using the starting run's scale_vector_params.
"""
if moresteps:
return
if not getattr(model_obj.retrieval_data, "use_map_optimizer", False):
return
import numpy as np
from GUIBRUSHR.Retrieval.ExofastMCMC.map_optimizer import find_map_start
rd = model_obj.retrieval_data
# Derive a statistically independent child seed from the MCMC SeedSequence
# so MAP and MCMC share one user-facing seed without any RNG correlation.
# spawn() advances the parent counter; subsequent MCMC spawn() calls
# therefore receive non-colliding child seeds (see ModelData.parallel_chain).
child_seq = model_obj.random_obj.seed_seq.spawn(1)[0]
map_seed = int(child_seq.generate_state(1, dtype=np.uint32)[0])
# Reuse the MCMC core count (multiplier_cores from the GUI) as the MAP
# pool size. ncores == 1 transparently keeps the historical serial path.
n_workers = int(getattr(model_obj.bestpars_data, "ncores", 1) or 1)
find_map_start(
model_obj,
maxiter_global=getattr(rd, "map_maxiter_global", 40),
popsize=getattr(rd, "map_popsize", 10),
tol_global=getattr(rd, "map_tol_global", 1.0e-2),
maxiter_local=getattr(rd, "map_maxiter_local", 150),
seed=map_seed,
n_workers=n_workers,
init_mode=getattr(rd, "init_mode", "isotropic"),
)
[docs]
def run_mcmc_sampling(model_obj, moresteps):
"""
Execute Monte Carlo Markov Chain sampling using exofast_demc.
Parameters
----------
model_obj : ModelData.ModelData
The configured model object
moresteps : bool
Whether to continue from a previous run
Returns
-------
tuple
A tuple containing (parameters, lhood_values) from the MCMC sampling # IDL: chi2_values
"""
# Run MONTECARLO sampling using exofast_demc algorithm
from GUIBRUSHR.Retrieval.ExofastMCMC import exofast_demc
parameters, lhood_values = exofast_demc.likelihood(model_obj, moresteps=moresteps) # IDL: chi2_values
return parameters, lhood_values
[docs]
def cleanup_and_report_timing(model_obj, start_time):
"""
Clean up resources and report execution timing.
Parameters
----------
model_obj : ModelData.ModelData
The model object to clean up
start_time : float
The start time from time.time()
"""
# Clean up model object to free memory
del model_obj
# Calculate and report execution time
end_time = time.time()
serial_time = end_time - start_time
print(f"Serial took {serial_time:.1f} seconds")
[docs]
def main():
"""
Main function orchestrating the atmospheric retrieval process.
This function coordinates the entire retrieval workflow:
1. Set up environment paths
2. Parse command line arguments
3. Create and configure the model object
4. Run MCMC sampling
5. Clean up and report results
"""
# Setup paths from command line arguments
path_default = sys.argv[5]
sys.path.append(path_default)
os.chdir(path_default)
# Parse command line arguments
(path_params, path_df, table_output_file,
id_process, path_default, moresteps) = parse_command_line_arguments()
# Record start time for performance measurement
start_time = time.time()
# Create and configure the model object
model_obj = create_model_object(
path_params, path_df, id_process, table_output_file, path_default
)
# Set up output directory
setup_output_directory(model_obj)
# Optional MAP warm-start (only runs if GUI flag use_map_optimizer is set,
# and never on resume - see run_map_warm_start docstring)
run_map_warm_start(model_obj, moresteps)
# Run MCMC sampling
_, _ = run_mcmc_sampling(model_obj, moresteps)
# Clean up and report timing
cleanup_and_report_timing(model_obj, start_time)
if __name__ == "__main__":
main()