A Python script designed for creating a grid of stellar models based on different parameters such as metallicity (Z), alpha-enhancement (A), and stellar mass (M). The script utilizes the PYSEP library for stellar physics calculations, handling input ranges for the desired parameters, and generating pre-main sequence models. It supports parallel processing to evolve the models and aims to build a grid of models to interpolate effective temperature (Teff) and luminosity (Lum) based on composition and mass. Suitable for astrophysical research and study in stellar evolution.

from pysep.api import stellarModel as sm
from pysep.api.parallel import pStellarModel
from pysep.misc.runDB.dbmanage import model_DB
 
from pysep.dm import filetypes as pyFts
 
from pysep.io.nml.control.defaults import solar
from pysep.io.nml.control.defaults import controlLow
from pysep.io.nml.control.defaults import controlHigh
 
from pysep.io.nml.physics.defaults import phys1
from pysep.io.nml.physics.defaults import phys1Low
from pysep.io.nml.physics.defaults import phys1High
 
from pysep.opac.opal.defaults import GS98hz
 
from pysep.opac.tops import open_and_parse
 
from pysep.newpoly.popFileToPolyModel import generate_prems_models
 
from pysep.prems import defaults as premsD
 
import argparse
import re
import os
import shutil
 
from tqdm import tqdm
import numpy as np
 
def get_grid_points(Zarange, Aarange, Marange):
    # TODO Use proper X inference from BBN for a given Z
    Zvalues = np.arange(*Zarange)
    Avalues = np.arange(*Aarange)
    Mvalues = np.arange(*Marange)
 
    elements = np.empty(shape=(Zvalues.shape[0]*Avalues.shape[0]*Mvalues.shape[0], 3))
    # TODO do idx better
    idx = 0
    for iz, z in enumerate(Zvalues):
        for ia, a in enumerate(Avalues):
            for im, m in enumerate(Mvalues):
                elements[idx] = np.array([z, a, m])
                idx += 1
    return elements
 
def build_prems(mass, abun, initpath):
    path = generate_prems_models(abun, mass, initpath, compName="TEMP", DDAGE=1000, fnameT="TEMP_%e%m.gs98")
    premsf = pyFts.premsf_file(os.path.join(initpath, path[mass]))
    return premsf
 
def build_model(X,Z,A,M,outDir,abun,premsPath):
    prems = build_prems(M, abun, premsPath)
 
    if M < 0.7:
        useCTL = controlLow.copy()
        usePHYS = phys1Low.copy()
    elif 0.7 <= M < 2:
        useCTL = solar.copy()
        usePHYS = phys1.copy()
    else:
        useCTL = controlHigh.copy()
        usePHYS = phys1High.copy()
 
    useCTL[0]['descrip'] = f"Evolution of {M} star"
    useCTL[1]['descrip'] = "Goal: Build grid of models to interpolate Teff Lum"
 
    useCTL['lbnout'] = True
    useCTL['nbn'] = 1
 
    usePHYS.add_key("ATIME", [0.001, 0.002, 0.5, 0.02, 0.3, 0.0015, 0.1, 0.02, 0.4, 0.02, 0.02])
 
    useCTL[1]['nmodls'] = 3
 
    useCTL[0]['cmixla'] = A
    useCTL[1]['cmixla'] = A
 
    tmpModel = sm(outDir,useCTL, usePHYS, GS98hz, prems)
    tmpModel.update_composition(X=X, Z=Z, ID=0)
 
    return tmpModel
 
 
def main(args):
    # gridPoints should be iterable of three elemenet iterable [Z, alpha, M]...
    gridPoints = get_grid_points(args.Zarange, args.Aarange, args.Marange)
    models = list()
    abun = open_and_parse(args.chem)
    for gsZ, gsA, gsM in tqdm(gridPoints):
        model = build_model(args.X, gsZ, gsA, gsM, args.outDir, abun, args.premsPath)
        models.append(model)
 
    DB = model_DB("TeffLumModelGrid.db", args.outDir, mode='wr')
    nModels = DB.enroll_models_to_dir(models, pbar=True)
 
 
    pModels = pStellarModel(nModels)
    pModels.pEvolve()
    pModels.pStash(data=True)
 
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run a grid of models to "
                                                 "generate interpolation "
                                                 "coefficients for getting teff"
                                                 " and lum from composition "
                                                 "and mass")
    parser.add_argument("Zarange", type=float, nargs=3, help="Z arange, high, low, step")
    parser.add_argument("Aarange", type=float, nargs=3, help="Alpha arange, high, low, step")
    parser.add_argument("Marange", type=float, nargs=3, help="M arange, high, low, step")
    parser.add_argument("chem", type=str, help="path to chemical abundance file")
    parser.add_argument("--X", type=float, help="X to use", default=0.735)
    parser.add_argument("--outDir", type=str, help="path to save models too", default="models/")
    parser.add_argument("--premsPath", type=str, help="path to save premain sequence models too", default="premsTMP/")
 
    args = parser.parse_args()
 
    main(args)