Isochrone Shift and Interpolation in Astronomy

This python file, named 'isochrones.py', is dedicated to manipulating and interpolating astronomical isochrones. It includes functions for shifting isochrones based on distance and extinction, loading stellar evolutionary tracks or isochrones from Color-Magnitude Diagram (CMD) pickle files found within a specified directory, and interpolating these isochrones for given stellar ages. The file applies computational physics and astrophysics techniques to model and analyze the properties and evolutionary states of stellar populations.

import numpy as np
import pickle as pkl
import pathlib
import re
 
from tqdm import tqdm
 
from scipy.interpolate import interp1d
 
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 load_ISO_CMDs(root):
    CMDs = list(map(lambda x: str(x), pathlib.Path(root).rglob("CMD.pkl")))
    extract = list(map(lambda x: re.findall(r"Pop(A|E)\+(\d\.\d+)\/alpha-(\d\.\d+)\/", x)[0], CMDs))
    pops = set(map(lambda x: x[0], extract))
    Ys = set(map(lambda x: x[1], extract))
    alphas = set(map(lambda x: x[2], extract))
 
    lookup = dict()
    for pop in tqdm(pops, leave=False):
        lookup[pop] = dict()
        for Y in tqdm(Ys, leave=False):
            lookup[pop][float(Y)] = dict()
            for alpha in tqdm(alphas, leave=False):
                if checkTup := (pop, Y, alpha) in extract:
                    extractID = extract.index((pop, Y, alpha))
                    with open(CMDs[extractID], 'rb') as f:
                        CMD = pkl.load(f)
                    lookup[pop][float(Y)][float(alpha)] = CMD
    return lookup
 
def interCMDatMag(color, mag, targetMag):
    f = interp1d(mag, color)
    return f(targetMag)
 
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