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') # #