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()