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)