Python Script for Stellar Model Evolution Simulation

This Python script leverages the pysep library to run simulations on the evolution of stellar models based on varying mass ranges. It allows users to set parameters such as mass range and composition, control physical conditions, and output high-resolution models. The script uses multiple modules for data handling, model generation, and running simulations in parallel, aimed at identifying fully convective boundaries in stars with ATOMIC opacities.

from pysep.api import stellarModel as sm
from pysep.api.parallel import pStellarModel
 
from pysep.dm.filetypes import opacf_file
from pysep.dm.filetypes import premsf_file
 
from pysep.io.nml.control.defaults import solar
from pysep.io.nml.physics.defaults import phys1
 
from pysep.io.nml.control.defaults import controlLow
from pysep.io.nml.physics.defaults import phys1Low
 
from pysep.opac.tops import open_and_parse
 
from pysep.newpoly.popFileToPolyModel import generate_prems_models
 
import argparse
import re
import os
import shutil
 
from tqdm import tqdm
import numpy as np
 
 
def get_prems_model_path(mass, searchPath):
    files = os.listdir(searchPath)
    formatedMass = f"{int(mass*100000):0>5}"
    premsFileFilter = filter(lambda x: re.search(f"{formatedMass}", x), files)
    premsFile = next(premsFileFilter)
    premsFilePath = os.path.join(searchPath, premsFile)
    return premsFilePath
 
def set_high_resolution(pnml):
    hr_pnml = pnml.copy()
    hr_pnml['kttau'] = 5
    hr_pnml['lenvg'] = True
    hr_pnml['atmstp'] = 0.002
    hr_pnml['envstp'] = 0.002
    hr_pnml['tridt'] = 1.0e-4
    hr_pnml['tridl'] = 8.0e-3
    hr_pnml['atmerr'] = 3.0e-5
    hr_pnml['atmmax'] = 0.05
    hr_pnml['atmmin'] = 0.015
    hr_pnml['atmbeg'] = 0.015
    hr_pnml['enverr'] = 3.0e-5
    hr_pnml['envmax'] = 0.05
    hr_pnml['envmin'] = 0.015
    hr_pnml['envbeg'] = 0.015
    hr_pnml['htoler'][0][0] = 6.0e-7
    hr_pnml['htoler'][0][1] = 4.5e-7
    hr_pnml['htoler'][0][2] = 3.0e-7
    hr_pnml['htoler'][0][3] = 9.0e-7
    hr_pnml['htoler'][0][4] = 3.0e-7
    hr_pnml['htoler'][1][4] = 2.5e-8
    hr_pnml['shelltol'][1] = 0.0575
    hr_pnml['shelltol'][7] = 0.01
    hr_pnml['shelltol'][8] = 0.01
    hr_pnml['shelltol'][9] = 0.01
    hr_pnml['shelltol'][10] = 0.01
    hr_pnml['niter2'] = 40
    hr_pnml['sstandard'] = [0.997543, 1.02913, 1.03704, 0.965517, 1.47273,
                            0.478916, 1.12766, 1.01083, 0.876543, 1.075,
                            0.625, 1.21795]
    return hr_pnml
 
 
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run models to identify fully convective boundary")
    parser.add_argument("mlow", help="low mass end", type=float)
    parser.add_argument("mhigh", help="high mass end", type=float)
    parser.add_argument("prems", help="path to prems directory", type=str)
    parser.add_argument("opac", help="high temperature opacity file", type=str)
    parser.add_argument("abun", help="path to abundance file", type=str)
    parser.add_argument("-o", "--output", type=str, help="output dir", default="data")
    parser.add_argument("-Z", help="rescale Z", default=None, type=float)
    parser.add_argument("--dm", help="mass step", default=0.001, type=float)
    parser.add_argument("--chem", help="path to chem file", type=str)
 
    args = parser.parse_args()
    masses = np.arange(args.mlow, args.mhigh, args.dm)
 
 
    opac = opacf_file(args.opac)
    abun = open_and_parse(args.abun)
    parsed = open_and_parse(args.chem)
 
    X = abun['AbundanceRatio']['X']
    if not args.Z:
        Znew = abun['AbundanceRatio']['Z']
    else:
        Znew = args.Z
    Ys = 1 - (abun['AbundanceRatio']['X'] + abun['AbundanceRatio']['Z'])
    Ybbn = 0.244
    dYdZ = (Ys - Ybbn)/abun['AbundanceRatio']['Z']
    Ynew = Ybbn + Znew*dYdZ
    Xnew = 1-(Ynew+Znew)
 
    X = Xnew
    Z = Znew
    if os.path.exists(args.output):
        shutil.rmtree(args.output)
    os.mkdir(args.output)
    if os.path.exists("premsTMP"):
        shutil.rmtree("premsTMP")
    os.mkdir("premsTMP")
 
    models = list()
 
 
    for mass in tqdm(map(lambda x: round(x, 4), masses), total=masses.shape[0]):
        # prems = premsf_file(get_prems_model_path(mass, args.prems))
        path = generate_prems_models(parsed, mass, "premsTMP/", compName="test", DDAGE=1000, fnameT="temp_m%e%m.prems")
        prems = premsf_file(os.path.join('premsTMP/', path[mass]))
 
        if mass < 0.7:
            useCTL = controlLow.copy()
            usePHYS = phys1Low.copy()
        else:
            useCTL = solar.copy()
            usePHYS = phys1.copy()
 
        usePHYS = set_high_resolution(usePHYS)
 
        useCTL[0]['descrip'] = f"Evolution of {mass} solar mass star"
        useCTL[1]['descrip'] = "Purpose: Identify fully convective boundary with ATOMIC opacities"
 
        saveDirName = f"{args.output}/M{mass}_model"
        os.mkdir(saveDirName)
 
        tmpModel = sm(saveDirName, useCTL, usePHYS, opac, prems)
        tmpModel.update_composition(X=X, Z=Z)
        models.append(tmpModel)
 
    pModels = pStellarModel(models)
    pModels.pEvolve()
    pModels.pStash(data=True)