This Python script is designed for fitting isochrones to Color-Magnitude Diagrams (CMDs) obtained from stellar photometry. It involves loading photometric data, isochrones, and fiducial sequences; performing optimizations to match these isochrones with observed stellar populations in CMDs based on chi-squared minimization; and handling the inputs/outputs for this process. It uses libraries such as NumPy for numerical operations, Matplotlib for plotting, and SciPy for optimization tasks. The script also supports command-line arguments for flexibility in selecting between different methods and adjusting parameters for the fitting process.

from photometry import load_HUGS, CMD_ridge_line_array, extract_pop, verticalize_CMD
from isochrones import shift_isochrone, load_ISO_CMDs, interCMDatMag, interp_isochrone_age
from fiducial import fiducial_points, iter_cut, two_point
 
from functools import lru_cache
 
import pickle as pkl
import os
import numpy as np
from tqdm import tqdm
import warnings
 
# import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
 
from scipy.optimize import minimize
from scipy.interpolate import interp1d
from mplEasyAnimate import animation
 
PHOT_ROOT = "../../photometry/HUGS/ngc2808"
ISO_ROOT = "../../outputs"
 
 
def pfD(r, I):
    # point function distance
    # r is a vector (x,y) to a point
    # I is some function, I(x)
    return lambda m: np.sqrt((I(m) - r[0])**2 + (m - r[1])**2)
 
def get_ISO_CMD_Chi2(iso, fiducialLine, filters=("F606W", "F814W", "F606W"), distance=0, Av=0, plot=False, ANIM=None):
    isoFilter1Name = f"WFC3_UVIS_{filters[0]}_MAG"
    isoFilter2Name = f"WFC3_UVIS_{filters[1]}_MAG"
    isoFilter3Name = f"WFC3_UVIS_{filters[2]}_MAG"
 
    isoColor = iso[isoFilter1Name] - iso[isoFilter2Name]
    isoMag = iso[isoFilter3Name]
 
 
    isoAptMag, isoAptColor = shift_isochrone(isoColor, isoMag, distance, Av)
 
    f = interp1d(isoAptMag, isoAptColor, bounds_error=False, fill_value='extrapolate')
    diffs = fiducialLine[:, 0] - f(fiducialLine[:, 1])
    # fig, ax = plt.subplots(1,1,figsize=(10,7))
    # ax.plot(isoAptColor, isoAptMag)
    # ax.invert_yaxis()
    sortedFiducial = fiducialLine[fiducialLine[:,1].argsort()]
    # ax.plot(sortedFiducial[:,0], sortedFiducial[:,1],color='red')
    minDist = np.empty(shape=(sortedFiducial.shape[0]))
    minDist[:] = np.nan
    for idx, (point, d0) in enumerate(zip(sortedFiducial, diffs)):
        d = pfD(point, f)
        nearestPoint = minimize(d, 0, method='Nelder-Mead')
        if not nearestPoint.success:
            print("FAIL")
            mD = np.apply_along_axis(d, 0, sortedFiducial[:,1])
            fig2, ax2 = plt.subplots(1,1)
            ax2.plot(sortedFiducial[:,1], mD)
            fig2.show()
        else:
            minDist[idx] = d(nearestPoint.x[0])
 
 
        # ax.plot([point[0], f(nearestPoint.x[0])],[point[1], nearestPoint.x[0]],color='green')
    # plt.show()
    minDist = minDist[~np.isnan(minDist)]
    chi2 = np.sum(np.apply_along_axis(lambda x: x**2,0,minDist))/minDist.shape[0]
    # chi2 = sum(map(lambda x: x**2, diffs))/len(diffs)
    return chi2
 
def optimize(fiducial, isochrone, filters,plot=False, ANIM=None):
    age_d_E_opt = lambda iso, cmd, age, d, E: get_ISO_CMD_Chi2(
            interp_isochrone_age(iso, age),
            fiducial,
            distance=d,
            Av=E,
            filters=filters,
            plot=plot,
            ANIM=ANIM)
 
    optimized = minimize(
            lambda x: age_d_E_opt(isochrone, photometry, x[0], x[1], x[2]),
            [12, 10000, 0.06],
            bounds=[
                (5,20),
                (5000, None),
                (0,0.3)
                ],
            )
    return {'opt': optimized, 'iso': isochrone, 'fiducial': fiducial}
 
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")
 
    args = parser.parse_args()
 
    if args.rFilterOrder:
        f3 = args.filter2
    else:
        f3 = args.filter1
    f1 = args.filter1
    f2 = args.filter2
 
    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]
 
    ISOs = load_ISO_CMDs(ISO_ROOT)
    print(ISOs.keys())
    print(ISOs['A'].keys())
 
    with open(args.fLines, 'rb') as f:
        fiducialSequences = pkl.load(f)
    fiducialLookup = dict()
    for i, fid in enumerate(fiducialSequences):
        fiducialLookup[fid.name] = i
 
    results = dict()
    popFid = None
    for population, popISO in tqdm(ISOs.items(), total=len(ISOs), leave=False):
        results[population] = dict()
        assert population in fiducialLookup.keys()
        print(population)
    #     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):
    #             results[population][Y][a] = optimize(popFid.mean, alphaISO, (f1, f2, f3))
    #
    # orderedOptimizationResults = order_best_fit_result(results)
    # with open(args.output, 'wb') as f:
    #     pkl.dump({'bf': orderedOptimizationResults, 'r': results, 'phot': cleanedPhot}, f)