This Python script is designed for astrophysical research, particularly in the field of stellar cluster analysis. It includes functions to shift isochrones based on astronomical distances and extinction values, load and parse color-magnitude diagrams (CMDs) from pickle files located within a hierarchical directory structure, and interpolate CMD data to find corresponding colors for given magnitudes. The script employs libraries such as NumPy for mathematical operations, pickle for object serialization, pathlib for file system navigation, re for regular expressions, tqdm for progress indication, and scipy for interpolation. It is a tool for astronomers and astrophysicists who are working on analyzing stellar populations and require manipulation and analysis of isochrone and CMD data.

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)