This Python code is designed for fitting isochrones to Color-Magnitude Diagrams (CMDs) of star clusters using photometric data. It involves loading and manipulating photometric data and isochrone models, fitting these models to observational data to estimate stellar population characteristics such as age and distance, and plotting the results for visualization. The script employs libraries such as numpy for numerical operations, matplotlib for plotting, scipy for optimization, and utilizes pickle for data serialization. It's an example of how computational techniques are applied in astrophysics for understanding stellar evolution and properties of star clusters.

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)