Python Script for Managing Stellar Composition Data

This Python script is designed for processing and transforming stellar track files into historical data records, managing a database of stellar models, and generating input files for isochrone calculations. It involves operations such as directory creation, file parsing, regex matching, and iterating through model databases to process stellar evolutionary tracks. This script is useful for astronomers and astrophysicists involved in stellar evolution and isochrone studies.

from pysep.io.trk.utils import trk_to_hist
from pysep.misc.runDB import model_DB
from pysep.dsep import load_model
from pysep.opac.tops import parseChemFile
 
import os
import argparse
import re
from pathlib import Path
import shutil
 
import re
 
from tqdm import tqdm
 
import warnings
 
def force_make_directory(path: str):
    if os.path.exists(path):
        shutil.rmtree(path)
    os.mkdir(path)
 
def find_input_abun_file(root, pop, Y):
    path = os.path.join(root, "comp")
    matchString = os.path.join(path, f"pop{pop}_ngc2808_Y+{Y}.dat")
    escapedMatch = re.escape(matchString)
 
    abunf = None
    for file in Path(path).rglob("*.dat"):
        if re.match(escapedMatch, str(file)):
            abunf = file
    return abunf
 
def write_iso_in(root, inputD, files, population, Y, fname):
    abunf = find_input_abun_file(inputD, population, Y)
    parsed = parseChemFile.open_and_parse(abunf)
    eep = os.path.join(root, "eep")
    force_make_directory(eep)
    iso = os.path.join(root, "iso")
    force_make_directory(iso)
    filesFormated = '\n'.join(files)
 
    Y = parsed['AbundanceRatio']['Y']
    Z = parsed['AbundanceRatio']['Z']
    FeH = parsed['AbundanceRatio']['[Fe/H]']
    aFe = parsed['AbundanceRatio']['[alpha/Fe]']
 
    out = f"""#version string, max 8 characters
example
#initial Y, initial Z, [Fe/H], [alpha/Fe], v/vcrit (space separated)
   {Y:0.2f}  {Z:0.4E}   {FeH:0.2f}        {aFe:0.2f}     0.0
#data directories: 1) history files, 2) eeps, 3) isochrones
{root}
{eep}
{iso}
#read history_columns
inputcolumns.list
#specify tracks
{len(files)}
{filesFormated}
#specify isochrones
isochrones.txt
min_max
log10
76
5.0
11.0"""
 
    with open(fname, 'w') as f:
        f.write(out)
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="generate the hist files from the track files already evolved")
    parser.add_argument("--NSCModelDB", help="path to NSCModelDB", default="../PresentationISOModels_NSC.db", type=str)
    parser.add_argument("--SCModelDB", help="path to SCModelDB", default="../PresentationISOModels_SC.db", type=str)
    parser.add_argument("--outputRoot", type=str, help="root to output hist files too", default="hstFiles")
    parser.add_argument("-Y", type=float, default=0.24, help="Helium Mass Fraction")
    parser.add_argument("--population", type=str.upper, default="A", help="population", choices=["A", "E"])
    parser.add_argument("--SaveDir", type=str, default="../ISOModelData_PopA_Y0.24_noZeros", help="path to find models on local computer")
    parser.add_argument("--isoinExt", help="extension to add to the end of iso in", type=str, default="")
 
    args = parser.parse_args()
 
    SCHistDir = os.path.join(args.outputRoot, "SC")
    NSCHistDir = os.path.join(args.outputRoot, "NSC")
 
    force_make_directory(args.outputRoot)
    os.mkdir(SCHistDir)
    os.mkdir(NSCHistDir)
 
 
    NSCDB = model_DB(args.NSCModelDB, args.SaveDir, mode="r")
    SCDB = model_DB(args.SCModelDB, args.SaveDir, mode="r")
 
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
 
        for SCmodel, NSCmodel in tqdm(zip(SCDB.iter_models(orderByMass=True), NSCDB.iter_models(orderByMass=True))):
            SClp = os.path.join(args.SaveDir, SCmodel.basepath)
            NSClp = os.path.join(args.SaveDir, NSCmodel.basepath)
            SClp = os.path.join(SClp, "model.dsep")
            NSClp = os.path.join(NSClp, "model.dsep")
 
            SClm, SCtmp = load_model(SClp)
            NSClm, NSCtmp = load_model(NSClp)
 
            SCFileName = f"DSEPTrk{SCmodel.premsMass}M-{args.Y}Y-pop{args.population}.hist"
            SCFilePath = os.path.join(SCHistDir, SCFileName)
 
            NSCFileName = f"DSEPTrk{NSCmodel.premsMass}M-{args.Y}Y-pop{args.population}.hist"
            NSCFilePath = os.path.join(NSCHistDir, NSCFileName)
 
            SCOkay = trk_to_hist(SClm.trk(), SCFilePath)
            NSCOkay = trk_to_hist(NSClm.trk(), NSCFilePath)
 
 
    match = re.compile(r"(DSEPTrk(\d\.\d+)M.+)")
 
 
    SCfiles = list(filter(lambda x: re.match(match, x), os.listdir(SCHistDir)))
    SCfiles = sorted(SCfiles, key=lambda x: float(re.findall(match, x)[0][1]))
    NSCfiles = list(filter(lambda x: re.match(match, x), os.listdir(NSCHistDir)))
    NSCfiles = sorted(NSCfiles, key=lambda x: float(re.findall(match, x)[0][1]))
 
    write_iso_in(SCHistDir,"../inputs",SCfiles,args.population, args.Y,f"SC{args.isoinExt}.isoin")
    write_iso_in(NSCHistDir,"../inputs",NSCfiles,args.population, args.Y,f"NSC{args.isoinExt}.isoin")