Chi-Square Plotting for Astrophysical Data Analysis

The file contains Python code for plotting and analyzing astrophysical data, particularly focused on fitting stellar population models using chi-square minimization techniques. It performs various data manipulations, optimizes model parameters to best fit observational data, and visualizes the comparison between different stellar populations. Key operations include interpolating isochrone ages, shifting isochrone data based on distance and extinction, and plotting the results to compare model predictions with observational data.

snippet.python

import matplotlib.pyplot as plt
import pickle as pkl
import numpy as np
from scipy.optimize import minimize
from scipy.interpolate import interp1d
from functools import lru_cache, wraps
from scipy.spatial.distance import cdist
import hdbscan
from scipy.optimize import curve_fit
from sklearn.cluster import KMeans

snippet.python

ageKey = 10.09090909090909

snippet.python

def shift_isochrone(color, magnitude, distance : float, extinction : float):
    mu = 5*np.log10(distance) - 5 + extinction
    aptMag = mu + magnitude
    aptCol = 3.2*extinction + color
    return aptMag, aptCol
 
def interp_isochrone_age(iso, targetAge):
    logTargetAgeYr= np.log10(targetAge*1e9)
    ageKeys = list(iso.keys())
    distance = [(x-logTargetAgeYr, x) for x in ageKeys]
    below = sorted(filter(lambda x: x[0] <=0, distance), key=lambda x: abs(x[0]))
    above = sorted(filter(lambda x: x[0] > 0, distance), key=lambda x: x[0])
    isoBelow = iso[below[0][1]]
    isoAbove = iso[above[0][1]]
    age1 = isoBelow['log10_isochrone_age_yr'].iloc[0]
    age2 = isoAbove['log10_isochrone_age_yr'].iloc[0]
 
 
    def linearinterpolate(x, other, age1, age2):
        newIso = ((other[x.name] - x)/(age2-age1)) * (logTargetAgeYr - age1) + x
        return newIso
 
    interpolated = isoBelow.apply(lambda x: linearinterpolate(x, isoAbove, age1, age2))
    return interpolated, (isoBelow, isoAbove)
 
def iso_at_params(iso, age, d, E, f1, f2, f3):
    interpolated, _= interp_isochrone_age(iso, age)
    f1 = interpolated[f"WFC3_UVIS_{f1}_MAG"]
    f2 = interpolated[f"WFC3_UVIS_{f2}_MAG"]
    f3 = interpolated[f"WFC3_UVIS_{f3}_MAG"]
    return shift_isochrone(f1-f2, f3, d, E)

snippet.python

with open("../Scripts/OptResults.pkl", "rb") as f:
    optimizationResults = pkl.load(f)

snippet.python

with plt.style.context(pubStyle):
    fig, ax = plt.subplots(1,1,figsize=(10,7))
    filters = ("F275W", "F814W", "F814W")
 
    bestFitResultsA = optimizationResults['bf']['A'][0]
    bestFitResultsB = optimizationResults['bf']['E'][0]
 
    bestFitOptimizationResultA = optimizationResults['r']['A'][bestFitResultsA[3]][bestFitResultsA[4]]
    bestFitOptimizationResultB = optimizationResults['r']['E'][bestFitResultsB[3]][bestFitResultsB[4]]
 
    bestFitPhotF1 = optimizationResults['phot'][filters[0]]
    bestFitPhotF2 = optimizationResults['phot'][filters[1]]
    bestFitPhotF3 = optimizationResults['phot'][filters[2]]
 
    bestFitPhotColor = bestFitPhotF1-bestFitPhotF2
    bestFitPhotMag = bestFitPhotF3
 
    bestFitAgeA = bestFitOptimizationResultA['opt']['x'][0]
    bestFitDistanceA = bestFitOptimizationResultA['opt']['x'][1]
    bestFitExtinctionA = bestFitOptimizationResultA['opt']['x'][2]
 
    bestFitAgeB = bestFitOptimizationResultB['opt']['x'][0]
    bestFitDistanceB = bestFitOptimizationResultB['opt']['x'][1]
    bestFitExtinctionB = bestFitOptimizationResultB['opt']['x'][2]
 
    c = 0
 
    bestFitIsoAptMagA, bestFitIsoAptColorA = iso_at_params(bestFitOptimizationResultA['iso'], bestFitAgeA, bestFitDistanceA-c, bestFitExtinctionA, *filters)
    bestFitIsoAptMagB, bestFitIsoAptColorB = iso_at_params(bestFitOptimizationResultB['iso'], bestFitAgeB, bestFitDistanceB, bestFitExtinctionB, *filters)
 
    # ax.scatter(bestFitPhotColor, bestFitPhotMag, s=1, alpha=0.01, color='yellow')
    ax.plot(bestFitOptimizationResultA['fiducial'][:, 0], bestFitOptimizationResultA['fiducial'][:, 1], color='black', marker='o', linestyle='', markersize=2)
    ax.plot(bestFitOptimizationResultB['fiducial'][:, 0], bestFitOptimizationResultB['fiducial'][:, 1], color='orange', marker='o', linestyle='', markersize=2)
    ax.plot(bestFitIsoAptColorA, bestFitIsoAptMagA, color='blue', alpha=0.5, label='Population A')
    ax.plot(bestFitIsoAptColorB, bestFitIsoAptMagB, color='red', alpha=0.5, label='Population E')
    ax.set_xlabel(f"{filters[0]} - {filters[1]}", fontsize=27)
    ax.set_ylabel(f"{filters[2]}", fontsize=27)
    ax.invert_yaxis()
 
    ax.annotate(fr"$\mu_{{A}} = {5*np.log10(bestFitDistanceA-c)-5:0.2f}$", (7, 14), fontsize=10)
    ax.annotate(fr"Age$_{{A}} = {bestFitAgeA:0.2f}$ Gyr", (7, 14.5), fontsize=10)
    ax.annotate(fr"$E(B-V)_{{A}} = {bestFitExtinctionA:0.2f}$", (7, 15), fontsize=10)
    ax.annotate(fr"Y$_{{A}} = {bestFitResultsA[3]}$", (7, 15.5), fontsize=10)
    ax.annotate(fr"$\alpha_{{ML,A}} = {bestFitResultsA[4]}$", (7, 16), fontsize=10)
    ax.annotate(fr"$\chi^{{2}}_{{\nu,A}} = {bestFitOptimizationResultA['opt']['fun']:0.4f}$",(7,16.5),fontsize=10)
 
    ax.annotate(fr"$\mu_{{E}} = {5*np.log10(bestFitDistanceB)-5:0.2f}$", (7, 18), fontsize=10)
    ax.annotate(fr"Age$_{{E}} = {bestFitAgeB:0.2f}$ Gyr", (7, 18.5), fontsize=10)
    ax.annotate(fr"$E(B-V)_{{E}} = {bestFitExtinctionB:0.2f}$", (7, 19), fontsize=10)
    ax.annotate(fr"Y$_{{E}} = {bestFitResultsB[3]}$", (7, 19.5), fontsize=10)
    ax.annotate(fr"$\alpha_{{ML,E}} = {bestFitResultsB[4]}$", (7, 20), fontsize=10)
    ax.annotate(fr"$\chi^{{2}}_{{\nu,E}} = {bestFitOptimizationResultB['opt']['fun']:0.4f}$",(7,20.5),fontsize=10)
 
    ax.legend(frameon=False)
 
    plt.savefig("Figures/DualPopFit.pdf")

snippet.python

bestFitOptimizationResultA['opt']

  message: CONVERGENCE: NORM_OF_PROJECTED_GRADIENT_<=_PGTOL
  success: True
   status: 0
      fun: 0.003141037634200801
        x: [ 1.768e+01  1.000e+04  7.489e-02]
      nit: 17
      jac: [ 2.482e-07 -6.485e-06  1.587e-06]
     nfev: 84
     njev: 21
 hess_inv: <3x3 LbfgsInvHessProduct with dtype=float64>

snippet.python

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

snippet.python

bestFitPhotF1 = optimizationResults['phot']["F275W"]
bestFitPhotF2 = optimizationResults['phot']["F814W"]
 
comp = order_best_fit_result(optimizationResults['r'])
AShiftParams = comp['A'][0]
EShiftParams = comp['E'][0]
 
bestFitPhotColor = bestFitPhotF1-bestFitPhotF2
bestFitPhotMag = bestFitPhotF2
testA = optimizationResults['r']['A'][AShiftParams[3]][AShiftParams[4]]
testB = optimizationResults['r']['E'][0.24][1.901]
fig, ax = plt.subplots(1,1,figsize=(10,7))
ax.scatter(bestFitPhotColor, bestFitPhotMag, s=1, c='yellow', alpha=0.1)
ax.plot(testA['fiducial'][:,0], testA['fiducial'][:,1], 'o', color='green')
ax.plot(testB['fiducial'][:,0], testB['fiducial'][:,1], 'o', color='red')
 
iso = testA['iso']
isoAAtAge, _ = interp_isochrone_age(iso, AShiftParams[1][0])
 
isoAColor = isoAAtAge['WFC3_UVIS_F275W_MAG'] - isoAAtAge["WFC3_UVIS_F814W_MAG"]
isoAMag = isoAAtAge["WFC3_UVIS_F814W_MAG"]
isoAAptMag, isoAAptColor = shift_isochrone(isoAColor, isoAMag, AShiftParams[1][1], AShiftParams[1][2])
 
ax.plot(isoAAptColor, isoAAptMag, color='tab:cyan')
ax.invert_yaxis()

14174741629.268076
16297508346.206469

snippet.python

comp['A'][0]

(0.00998234825756846,
 array([1.41747416e+01, 1.05306450e+04, 1.47434696e-01]),
 'A',
 0.33,
 1.45)

snippet.python