This script is used for applying bolometric corrections to astronomical isochrones. It loads isochrone data, applies bolometric corrections using pre-calculated tables, and saves the corrected data. The script handles file reading, manipulation based on metallicity, and data interpolation. It is primarily designed for astronomers and astrophysicists dealing with stellar evolutionary models and their observable properties.

from fidanka.bolometric.color import get_mags
from fidanka.bolometric.bctab import Bolometric
from fidanka.bolometric.color import get_interpolated_FeHTable
from fidanka.bolometric.load import load_bol_table
from fidanka.isochrone.MIST import read_iso, read_iso_metadata
 
import os
import re
import pandas as pd
from pathlib import Path
import pickle
 
from tqdm import tqdm
 
FEHPATTERN = re.compile(r"feh(p|m)(\d{3})")
 
def load_all_from_bol_dir(bolDir : str):
    signLookup = {'p': '', 'm': '-'}
    files = os.listdir(bolDir)
    files = [f for f in files if f.startswith('feh')]
    filePaths = map(lambda f: os.path.join(bolDir, f), files)
    lookupKeys = map(lambda f: re.findall(FEHPATTERN, f), files)
    lookupKeys = filter(lambda f: len(f) > 0, lookupKeys)
    lookupKeys = map(lambda f: f[0], filter(lambda keys: len(keys[0]) == 2, lookupKeys))
    lookupKeys = map(lambda key: f"{signLookup[key[0]]}{key[1][0]}.{key[1][1:]}", lookupKeys)
    lookupDict = {float(key): val for key, val in zip(lookupKeys, filePaths)}
    return lookupDict
 
def correct_isochrone(isoPath, bolLookupTable):
    iso = read_iso(isoPath)
    iso_metadata = read_iso_metadata(isoPath)
    FeH = iso_metadata['[Fe/H]']
    nearestAboveKey = min(filter(lambda key: key >= FeH, bolLookupTable.keys()))
    nearestBelowKey = max(filter(lambda key: key <= FeH, bolLookupTable.keys()))
 
    BC1 = Bolometric(bolLookupTable[nearestBelowKey], nearestBelowKey)
    BC2 = Bolometric(bolLookupTable[nearestAboveKey], nearestAboveKey)
 
    interpolatedBolTable = get_interpolated_FeHTable(
            BC1,
            BC2,
            0.0,
            3.1,
            FeH
            )
    mags = dict()
    for age, isoAtAge in tqdm(iso.items(), total=len(iso), leave=False):
        magTable = get_mags(10**isoAtAge['log_Teff'], isoAtAge['log_g'], isoAtAge['log_L'], interpolatedBolTable)
        magDf = pd.DataFrame(magTable)
        mags[age] = pd.concat([isoAtAge, magDf], axis=1)
 
    origIsoStem = Path(isoPath).stem
    newFileName = f"{origIsoStem}.bol"
    origFilePath = Path(isoPath)
    newFilePath = origFilePath.parent / newFileName
    with open(newFilePath, 'wb') as f:
        pickle.dump(mags, f)
 
 
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Bolometrically correct isochrones')
    parser.add_argument('isoDir', type=str, help='isochrone directory to look through')
    parser.add_argument('bolDir', type=str, help='bolometric correction file')
 
    args = parser.parse_args()
    bolLookupTable = load_all_from_bol_dir(args.bolDir)
 
    for isochrone in tqdm(Path(args.isoDir).rglob("isochrones.txt")):
        correct_isochrone(str(isochrone), bolLookupTable)
 
 
    # magTable = get_mags(10**iso['log_Teff'], iso['log_g'], iso['[Fe/H]'], interpolatedBolTable)