Source code for GUIBRUSHR.Retrieval.main

#!/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




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)





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





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}")





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, chi2_values) from the MCMC sampling
    """
    # Run MONTECARLO sampling using exofast_demc algorithm
    from GUIBRUSHR.Retrieval.ExofastMCMC import exofast_demc
    parameters, chi2_values = exofast_demc.likelihood(model_obj, moresteps=moresteps)
    return parameters, chi2_values





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")





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)
    
    # 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()