This Python script automates the process of generating and evolving stellar models for a given range of stellar masses using pre-main sequence models, high temperature opacity tables, and chemical abundance files. It incorporates parallel processing for efficiency and allows for customization through command-line arguments.

from pysep.dm.filetypes import premsf_file
from pysep.dm.filetypes import opacf_file
 
from pysep.io.nml.control.defaults import solar, controlLow
from pysep.io.nml.physics.defaults import phys1, phys1Low
 
from pysep.opac.tops import open_and_parse
from pysep.opac.opal.defaults import GS98hz
 
from pysep.dsep import stellarModel as sm
 
from pysep.api.parallel import pStellarModel
 
 
import argparse
from typing import Generator
import re
import shutil
import os
 
def error_check(opac, prems, abun):
    A = (opac[1] == prems[1]) and (prems[1] == abun[1])
    B = (opac[2] == prems[2]) and (prems[2] == abun[2])
    return A and B
 
def mass_generator(mlow: float, mhigh: float) -> Generator[float, None, None]:
    """
    Generator to return mass with a varying incriment between masses in
    different ranges. Below 0.6 msolar masses incriment by 0.025, above that
    they incriment by 0.05.
 
    Parameters
    ----------
        mlow : float
            low end of mass range, used inclusivly
        mhigh : float
            high end of mass range, used inclusivly
 
    Yields
    ------
        rmass : float
            mass to investigate
    """
    nDigits = 6 # rounds are to quash floating point errors
    rmass = round(mlow, nDigits)
    while rmass <= mhigh:
        yield rmass
        if rmass < 0.6:
            rmass = round(rmass + 0.025, nDigits)
        else:
            rmass = round(rmass + 0.05, nDigits)
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate dsep output for all NGC_2808 abundance tables")
    parser.add_argument("rootPath", help="root path to store output", type=str)
    parser.add_argument("opacPath", help="path to high temperature opacity tables", type=str)
    parser.add_argument("premsPath", help="path to pre main sequence models", type=str)
    parser.add_argument("abunPath", help="path to chemical abundance files", type=str)
    parser.add_argument("-op", "--opacPattern", help='regular expression opac pattern', type=str, default=r".+_opac")
    parser.add_argument("-pp", "--premsPattern", help='regular expression prems pattern', type=str, default=r"^m.+\.prems")
    parser.add_argument("-ap", "--abunPattern", help='regular expression abun pattern', type=str, default=r'.+Y[+|-]\d\.\d+\.dat')
    parser.add_argument("--opal", help="use default opal opacity GS98hz", action="store_true", default=False)
    parser.add_argument("--kttau", type=int, default=5, help="KTTAU flag to pick atmosphere model")
    parser.add_argument("--outdir", type=str, help="output directory", default="data")
 
    args = parser.parse_args()
 
    massRange = (0.3, 0.9)
    masses = mass_generator(massRange[0], massRange[1])
 
    opacDirCont = os.listdir(args.opacPath)
    premsDirCont = os.listdir(args.premsPath)
    abunDirCont = os.listdir(args.abunPath)
 
    opacRegx = args.opacPattern
    premsRegx = args.premsPattern
    abunRegx = args.abunPattern
 
    opacFiles = filter(lambda x: re.match(opacRegx, x), opacDirCont)
    opacPaths = map(lambda x: os.path.join(os.path.abspath(args.opacPath), x), opacFiles)
    opacPathsM = map(lambda x: (x, float(re.search('(?<=Y[+|-])\d+\.\d+', x).group(0)), re.search("(?<=pop).", x).group(0)), opacPaths)
    opacPathsSorted = sorted(opacPathsM, key=lambda x: (x[2], x[1]))
 
    abunFiles = filter(lambda x: re.match(abunRegx, x), abunDirCont)
    abunPaths = map(lambda x: os.path.join(os.path.abspath(args.abunPath), x), abunFiles)
    abunPathsM = map(lambda x: (x, float(re.search('(?<=Y[+|-])\d+\.\d+', x).group(0)), re.search("(?<=pop).", x).group(0)), abunPaths)
    abunPathsSorted = sorted(abunPathsM, key=lambda x: (x[2], x[1]))
    models = list()
    for targetMass in [int(x*1000) for x in masses]:
        if 'm' in premsRegx:
            premsRegxS = premsRegx.replace('m', f"m{targetMass:0>4d}", 1)
 
        premsFiles = filter(lambda x: re.match(premsRegxS, x), premsDirCont)
        premsPaths = map(lambda x: os.path.join(os.path.abspath(args.premsPath), x), premsFiles)
        premsPathsM = map(lambda x: (x, float(re.search('(?<=Y[+|-])\d+\.\d+', x).group(0)), re.search("(?<=pop).", x).group(0)), premsPaths)
        premsPathsSorted = sorted(premsPathsM, key=lambda x: (x[2], x[1]))
 
        for opac, prems, abun in zip(opacPathsSorted, premsPathsSorted, abunPathsSorted):
            error_check(opac, prems, abun)
            pop = re.findall("(?:pop)(A|E)", abun[0])[0]
            parsed = open_and_parse(abun[0])
            X = parsed['AbundanceRatio']['X']
            Z = parsed['AbundanceRatio']['Z']
            Y = 1-(X+Z)
 
 
            if targetMass < 0.75:
                usePHYS = phys1Low.copy()
                useCTL = controlLow.copy()
            else:
                usePHYS = phys1.copy()
                useCTL = solar.copy()
 
            useCTL[0]['rsclx'] = X
            useCTL[0]['rsclz'] = Z
            useCTL[0]['descrip'] = f"NGC 2808 Pop{pop} Y+{Y:0.2f}, {targetMass/1000} solar masses"
            useCTL[1]['descrip'] = f"NGC 2808 Pop{pop} Y+{Y:0.2f}, {targetMass/1000} solar masses, evolve"
            useCTL[1]['endage'] = 12e9
 
            usePHYS['kttau'] = args.kttau
 
            outputName = f"{args.outdir}/DSEPRun_ngc2808_pop{opac[2]}_Y+{opac[1]:<0.2f}_m{targetMass:0>04d}"
 
            opacf = opacf_file(opac[0])
            premsf = premsf_file(prems[0])
 
            if not os.path.exists(outputName):
                os.mkdir(outputName)
            else:
                shutil.rmtree(outputName)
                os.mkdir(outputName)
            if args.opal:
                opac = GS98hz
            else:
                opac = opacf
            model = sm(outputName, useCTL, usePHYS, opac, premsf)
            models.append(model)
            print(f"Model {outputName} Defined")
 
    pModels = pStellarModel(models)
    pModels.pEvolve()
    pModels.pStash()