This Python script is designed for fitting isochrones to color-magnitude diagrams (CMDs) utilizing photometry data and bolometric correction tables. It supports multiple methods for fitting, can accept various filters, and allows for the generation of plots to visualize the fitting process. The script relies on external libraries for loading isochrone data and performing the optimization process. It saves the results in a Pickle file format, including the best fit results and processed photometry data. This can be used in astrophysical research to model stellar populations based on observational data.

from fidanka.isofit.fit import optimize, MC_optimize
from isochrones import load_ISO_CMDs, interCMDatMag, interp_isochrone_age
 
import pickle as pkl
import os
import shutil
from tqdm import tqdm
import re
 
PHOT_ROOT = "../../photometry/HUGS/ngc2808"
ISO_ROOT = "../../outputs.denseAlpha.fixedLowmass"
 
 
def order_best_fit_result(optimizationResults):
    comparison = {'A': list(), 'E': list()}
    for pop, popD in optimizationResults.items():
        for Y, YD in popD.items():
            for a, aD in YD.items():
                comparison[pop].append((aD['opt']['fun'], aD['opt']['x'], pop, Y, a))
        comparison[pop] = sorted(comparison[pop], key=lambda x: x[0])
    return comparison
 
 
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()
 
    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("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]
 
 
    with open(args.fLines, 'rb') as f:
        fiducialSequences = pkl.load(f)
    fiducialSequences[0].name = "A"
    fiducialSequences[1].name = "E"
    fiducialLookup = dict()
    print(fiducialSequences)
    for i, fid in enumerate(fiducialSequences):
        fiducialLookup[fid.name] = i
 
    ISOs, FeHs = load_ISO_CMDs(ISO_ROOT)
    results = dict()
    popFid = None
    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]
                print(f"Working on {population} {Y} {a} {FeH}")
                if os.path.exists(f"testFigures/{population}_{Y}_{a}"):
                    shutil.rmtree(f"testFigures/{population}_{Y}_{a}")
                os.mkdir(f"testFigures/{population}_{Y}_{a}")
                results[population][Y][a] = optimize(
                        popFid.mean,
                        alphaISO,
                        (f1, f2, f3),
                        bolPaths,
                        FeH,
                        figDir=f"{population}_{Y}_{a}",)
 
 
    orderedOptimizationResults = order_best_fit_result(results)
    with open(args.output, 'wb') as f:
        pkl.dump({'bf': orderedOptimizationResults, 'r': results, 'phot': cleanedPhot}, f)