This Python script dynamically organizes stellar evolution track data (.hist files) and relevant abundance files into a structured directory layout based on population, metallicity, and file content. It categorizes files into population and metallicity specific folders, and prepares input files for further analysis by parsing template files and incorporating specific abundance ratios. The script is intended to streamline the preparation of data for stellar evolutionary models, enhancing the analytical workflow.

import os
import shutil
from pathlib import Path
import argparse
import re
 
from pysep.opac.tops.parseChemFile import open_and_parse
 
def main(files, abun, force=False):
    populations = set([re.findall(r'pop([A-Z])', x)[0] for x in files])
    metals = set([re.findall(r'(0\.\d+)Y', x)[0] for x in files])
    abunFiles = os.listdir(abun)
    for pop in populations:
        for metal in metals:
            dirName = f'pop{pop}Y{metal.replace(".", "")}'
            if not os.path.exists(dirName):
                os.mkdir(dirName)
            else:
                if len(os.listdir(dirName)) == 0:
                    shutil.rmtree(dirName)
                    os.mkdir(dirName)
                else:
                    if not force:
                        print(f"{dirName} exists and is not empty! skipping!")
                        break
                    else:
                        shutil.rmtree(dirName)
                        os.mkdir(dirName)
            fileSet = filter(lambda x: f'pop{pop}' in x and f'{metal}Y' in x, files)
            newFileNames = list()
            for file in fileSet:
                fileMass = re.findall(r'DSEPTrk(0\.\d+)M', file)[0]
                newFileName = f'DSEPTrk{fileMass.replace(".", "")}M.hist'
                newFileNames.append(newFileName)
                shutil.copy(file, f"{dirName}/{newFileName}")
            os.mkdir(os.path.join(dirName, "eep"))
            os.mkdir(os.path.join(dirName, "iso"))
 
            newFileNames = sorted(newFileNames, key=lambda x: float(".{}".format(re.findall(r'DSEPTrk(\d+)M', x)[0][1:])))
 
 
            assert os.path.exists('template')
 
            popFile = filter(lambda x: f"pop{pop}" in x and f"Y+{metal}" in x, abunFiles)
            popFile = next(popFile)
            parsed = open_and_parse(os.path.join(abun, popFile))
 
            inputFiles = filter(lambda x: 'input' in x, os.listdir('template'))
            for iFile in inputFiles:
                if iFile.endswith('.isoin'):
                    newFileName = f"{dirName}input.isoin"
                    with open(os.path.join('template', iFile), 'r') as f:
                        contents = f.read().split('\n')
                    contents[1] = dirName
                    compLine = f"   {parsed['AbundanceRatio']['Y']:0.2f}  "\
                                f"{parsed['AbundanceRatio']['Z']:0.4E}   "\
                                f"{parsed['AbundanceRatio']['[Fe/H]']:0.2f}    "\
                                f"{parsed['AbundanceRatio']['[alpha/Fe]']:0.2f}"\
                                "      0.0"
                    contents[3] = compLine
                    contents[5] = os.path.abspath(dirName)
                    contents[6] = os.path.abspath(os.path.join(dirName, "eep"))
                    contents[7] = os.path.abspath(os.path.join(dirName, "iso"))
                    number = contents[11]
                    contents[11] = str(len(newFileNames))
                    templateFiles = contents[12:12+int(number)]
                    filteredContents = [x for i, x in enumerate(contents) if not x in templateFiles]
                    filteredContents[12:12] = newFileNames
                    joinedContents = '\n'.join(filteredContents)
                    with open(os.path.join(dirName, newFileName), 'w') as f:
                        f.write(joinedContents)
                else:
                    shutil.copy(os.path.join('template', iFile), os.path.join(dirName, iFile))
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="seperate hist files into directories")
    parser.add_argument("files", type=str, nargs='+', help="all files")
    parser.add_argument("-f", "--force", action="store_true", help="delete files if they exist")
    parser.add_argument("--abundir", type=str, help="path to abundance files")
 
    args = parser.parse_args()
    main(args.files, args.abundir, force=args.force)#, args.alphafe, args.Y, args.Z,