This Python script is designed for astrophysical analysis, specifically for adjusting isochrones based on distance and extinction, loading stellar population models with varying characteristics (such as population type, helium content, and alpha enrichment) from pickle files, and interpolating these models by magnitude or age. It uses libraries such as numpy for numerical operations, scipy for interpolation, and tqdm for progress bars during data processing.

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