This Python script is part of a larger project for fitting isochrones to color-magnitude diagrams (CMDs) of star clusters. It imports functionality from a custom package fidanka, handles astronomical data processing with numpy and file operations, uses pickle for serialization, and relies on argparse for command line argument parsing. The script supports selecting different fitting methods, applies bolometric corrections, loads photometric and isochrone data, performs optimization to match observed and model CMDs, and finally saves the results in a Pickle file.

from fidanka.isofit.fit import optimize
from fidanka.isochrone.isochrone import load_ISO_CMDs
from fidanka.isofit.fit import order_best_fit_result
from fidanka.misc.logging import LoggerManager
 
import pickle as pkl
import numpy as np
import os
import shutil
from tqdm import tqdm
import re
 
PHOT_ROOT = "../../photometry/HUGS/ngc2808"
ISO_ROOT = "../../outputs.denseAlpha.fixedLowmass"
# ISO_ROOT = "../../outputs.smallTest"
 
 
def format_iso_nd_array(isochrones):
    outDict = dict()
    for key, iso in isochrones.items():
        Teff = 10**iso["log_Teff"].values
        logg = iso["log_g"].values
        logL = iso["log_L"].values
        EEPs = iso["EEP"].values
        outDict[key] = np.array([EEPs, Teff, logg, logL]).T
    return outDict
 
 
if __name__ == "__main__":
    import argparse
 
    parser = argparse.ArgumentParser(description="Fit an isochrones to a CMD")
    method = parser.add_mutually_exclusive_group(required=True)
    method.add_argument("--meth1", help="Use HUGS Method 1", action='store_true')
    method.add_argument("--meth2", help="Use HUGS Method 2", action='store_true')
    method.add_argument("--meth3", help="Use HUGS Method 3", action='store_true')
 
    parser.add_argument("--fLines", help="Fiducial piclke file", default="Fiducial.pkl")
 
    parser.add_argument("--filter1", help="Magnitude filter", default="F606W")
    parser.add_argument("--filter2", help="Second Color Filter", default="F814W")
    parser.add_argument("--rFilterOrder", action="store_true", help="flip which filter in the color is the magnitude axis")
 
    parser.add_argument("--verbosePlotting", action="store_true", help="Write out a lot of plots")
    parser.add_argument("-o", "--output", type=str, default="OptResult.pkl", help="Where to save Pickle File")
    parser.add_argument("--bTabs", type=str, help="Root paths to the bolometric  correction tables")
 
    args = parser.parse_args()
    LoggerManager.config_logger("isoFit_fidanka.log")
 
    if args.rFilterOrder:
        f3 = args.filter2
    else:
        f3 = args.filter1
    f1 = args.filter1
    f2 = args.filter2
 
    assert os.path.exists(args.bTabs), "Bolometric Correction Table Path does not exist"
    bolFilenames = list(filter(lambda x: re.search(r"feh[mp]\d+", x), os.listdir(args.bTabs)))
    bolPaths = list(map(lambda x: os.path.join(args.bTabs, x), bolFilenames))
 
    with open(os.path.join(PHOT_ROOT, "photometry.pkl"), 'rb') as f:
        photometry = pkl.load(f)
 
    if args.meth1:
        methID = 1
    elif args.meth2:
        methID = 2
    elif args.meth3:
        methID = 3
    else:
        raise RuntimeError("Something Horrid has gone wrong! You should re-evalute the choices which lead you here...")
    cleanedPhot = photometry[methID]
 
 
    print("Loading Fiducial Sequences....")
    with open(args.fLines, 'rb') as f:
        fiducialSequences = pkl.load(f)
    fiducialSequences[0].name = "A"
    fiducialSequences[1].name = "E"
    fiducialLookup = dict()
    for i, fid in enumerate(fiducialSequences):
        fiducialLookup[fid.name] = i
    print("Fiducial Sequences Loaded")
 
    print("Loading Isochrones....")
    ISOs, FeHs = load_ISO_CMDs(ISO_ROOT)
    print("Isochrones Loaded")
    results = dict()
    popFid = None
 
    print("Starting Optimization...")
    for population, popISO in tqdm(ISOs.items(), total=len(ISOs), leave=False):
        results[population] = dict()
        assert population in fiducialLookup.keys()
        if population == "A":
            popFid = fiducialSequences[fiducialLookup['A']]
        if population == "E":
            popFid = fiducialSequences[fiducialLookup['E']]
        for Y, YISO in tqdm(popISO.items(), total=len(popISO), leave=False):
            results[population][Y] = dict()
            for a, alphaISO in tqdm(YISO.items(), total=len(YISO), leave=False):
                FeH = FeHs[population][Y][a]
                iso = format_iso_nd_array(alphaISO)
                results[population][Y][a] = optimize(
                        popFid.mean,
                        iso,
                        'WFC3',
                        FeH,
                        filters=(f1,f2,f3)
                        )
 
    print("Optimization Complete")
    print("Saving Results...")
    orderedOptimizationResults = order_best_fit_result(results)
    with open(args.output, 'wb') as f:
        pkl.dump({'bf': orderedOptimizationResults, 'r': results, 'phot': cleanedPhot}, f)
    print("Results Saved")