Astronomical Isochrone Fitting Script

This Python script is designed for fitting theoretical isochrones to observed Color-Magnitude Diagrams (CMDs) from astronomical photometry data. It includes functionalities for loading stellar photometry, adjusting isochrones based on extinction and distance, calculating chi-squared statistics for the fit, and optimizing the fit parameters like age, distance, and extinction using scipy minimize function. The script supports command line arguments for specifying methods and population parameters, and it employs matplotlib for optional plotting of CMDs and the fit. It exemplifies the integration of astronomical data analysis techniques in Python, particularly for stellar population studies.

from photometry import load_HUGS, CMD_ridge_line_array
from isochrones import shift_isochrone, load_ISO_CMDs, interCMDatMag, interp_isochrone_age
 
from functools import lru_cache
 
import pickle as pkl
import os
import numpy as np
from tqdm import tqdm
 
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
 
from scipy.optimize import minimize
 
PHOT_ROOT = "../../photometry/HUGS/ngc2808"
ISO_ROOT = "../../outputs"
 
def get_ISO_CMD_Chi2(iso, fiducialLine, filters=("F606W", "F814W"), distance=0, Av=0, plot=False):
    isoFilter1Name = f"WFC3_UVIS_{filters[0]}_MAG"
    isoFilter2Name = f"WFC3_UVIS_{filters[1]}_MAG"
 
    isoColor = iso[isoFilter1Name] - iso[isoFilter2Name]
    isoMag = iso[isoFilter2Name]
 
    isoAptMag, isoAptColor = shift_isochrone(isoColor, isoMag, distance, Av)
 
    if plot:
        fig, ax = plt.subplots(1,1,figsize=(10,7))
        # ax.plot(isoAptColor, isoAptMag)
        ax.plot(fiducialLine[0, :], fiducialLine[1, :], color='green')
        ax.plot(isoAptColor, isoAptMag)
        ax.plot(isoColor, isoMag, color='red')
        ax.invert_yaxis()
        fig.savefig(f"Figures/CMD_d{distance}_E{Av}_Age{iso['log10_isochrone_age_yr'].iloc[0]}.png", bbox_inches="tight")
        plt.close(fig)
 
    diffs = list()
    for mag, color in zip(fiducialLine[1, :], fiducialLine[0, :]):
        if isoAptMag.min() <= mag <= isoAptMag.max():
            isoTargetColor = interCMDatMag(isoAptColor, isoAptMag, mag)
            diffs.append(color - isoTargetColor)
    chi2 = sum(map(lambda x: x**2, diffs))/len(diffs)
    return chi2
 
 
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="Method 1", action='store_true')
    method.add_argument("--meth2", help="Method 2", action='store_true')
    method.add_argument("--meth3", help="Method 3", action='store_true')
 
    parser.add_argument("-p", "--pop", choices=["A", "E"], type=str, help="population isochrones to fit", required=True)
    parser.add_argument("-Y", "--helium", type=float, choices=[0.24, 0.27, 0.30, 0.33, 0.36, 0.39], help="Helium mass fraction", required=True)
    parser.add_argument("-a", "--alpha", type=float, choices=[1.0, 1.45, 1.901, 2.35, 2.8], help="Mixing length parameter", required=True)
 
    args = parser.parse_args()
 
 
    CMDs = load_ISO_CMDs(ISO_ROOT)
    # for pop, CMDpop in CMDs.items():
    #     for Y, CMDY in CMDpop.items():
    #         fig, ax = plt.subplots(1,1,figsize=(10,7))
    #         for a, CMDa in CMDY.items():
    #             f1 = CMDa[10.272727272727272]["WFC3_UVIS_F606W_MAG"]
    #             f2 = CMDa[10.272727272727272]["WFC3_UVIS_F814W_MAG"]
    #             c = f1-f2
    #             ax.plot(c, f1, label=f"a = {a}")
    #         ax.invert_yaxis()
    #         ax.legend()
    #         fig.savefig(f"Figures/ISO_Pop{pop}_Y+{Y}.png", bbox_inches='tight')
    #         plt.close(fig)
    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...")
 
    cleaned = photometry[methID]
    fiducialLine = CMD_ridge_line_array(cleaned, "F606W", "F814W")
 
    isocmd = CMDs[args.pop][args.helium][args.alpha]
 
    age_d_E_opt = lambda iso, cmd, age, d, E: get_ISO_CMD_Chi2(
            interp_isochrone_age(iso, age),
            fiducialLine,
            distance=d,
            Av=E)
 
    optimized = minimize(
            lambda x: age_d_E_opt(isocmd, cleaned, x[0], x[1], x[2]),
            [13.5, 12000, 0.06]
            )
    print(optimized)
 
 
 
 
 
 
    # dGrid = np.linspace(12000, 14000, 20)
    # eGrid = np.logspace(-2, -1, 20)
    # aGrid = list(isocmd.keys())
    # chi2 = np.empty(shape=(eGrid.shape[0], dGrid.shape[0], 2, len(isocmd)))
    #
    # Av = 0.045
    # for eIDX, e in tqdm(enumerate(eGrid), total=len(eGrid)):
    #     for dIDX, d in tqdm(enumerate(dGrid), total=len(dGrid), leave=False):
    #         for idx, (age, iso) in enumerate(isocmd.items()):
    #             x = get_ISO_CMD_Chi2(iso, fiducialLine, distance=d, Av=e)
    #             chi2[eIDX, dIDX, 0, idx] = 10**age
    #             chi2[eIDX, dIDX, 1, idx] = x
    # # with open("data.pkl", 'wb') as f
    # #     pkl.dump({'chi2': chi2, 'iso': isocmd, 'phot': cleaned, 'dGrid': dGrid, 'eGrid': eGrid, 'aGrid': aGrid, 'fiducialLine': fiducialLine}, f)
    # minIDX = np.unravel_index(chi2[:, :, 1].argmin(), shape=chi2[:, :, 1].shape)
    # fig, axs = plt.subplots(1,2,figsize=(20,7))
    # for age, iso in isocmd.items():
    #     color = iso['WFC3_UVIS_F606W_MAG'] - iso['WFC3_UVIS_F814W_MAG']
    #     mag = iso['WFC3_UVIS_F606W_MAG']
    #     aptMag, aptColor = shift_isochrone(color, mag, dGrid[minIDX[1]], eGrid[minIDX[0]])
    #     axs[1].plot(aptColor, aptMag, color='blue', alpha=0.1)
    # axs[0].semilogx(chi2[minIDX[0], minIDX[1], 0], chi2[minIDX[0], minIDX[1], 1], 'o-')
    # axs[1].plot(fiducialLine[0,:], fiducialLine[1,:], color='red')
    # axs[1].scatter(cleaned["F606W"] - cleaned["F814W"], cleaned["F606W"], color='pink', alpha=0.1)
    # sISO = isocmd[aGrid[minIDX[2]]]
    # color = sISO['WFC3_UVIS_F606W_MAG'] - sISO['WFC3_UVIS_F814W_MAG']
    # mag = sISO['WFC3_UVIS_F606W_MAG']
    # aptMag, aptColor = shift_isochrone(color, mag, dGrid[minIDX[1]], eGrid[minIDX[0]])
    # axs[1].plot(aptColor, aptMag, color='blue')
    # axs[1].invert_yaxis()
    # axs[1].set_title(f"Age={(10**aGrid[minIDX[2]])/1e9:0.3f} Gyr, d={dGrid[minIDX[1]]}pc, E(B-V)={eGrid[minIDX[0]]}")
    # fig.savefig("Figures/MinChi2.png", bbox_inches='tight')
    #
    #