This script is designed for loading and processing data from dsep3 track (trk) and isochrone (iso) files related to stellar evolution models. It includes functionalities for parsing metadata and data from these files, handling file paths and user arguments for operational flexibility, and optionally saving processed data as pickles. The script combines various Python libraries, such as pandas and argparse, to manipulate and organize data efficiently for further analysis or visualization in astrophysical research.

import re
import pandas as pd
import numpy as np
import os
import itertools
import argparse
import pickle
 
def load_iso(path):
    header = ["Age", "Log_T", "Log_g", "Log_L", "Log_R", "Y_core", "Z_core", "(Z/X)_surf", "L_H", "L_He", "M_He_core", "M_CO_core"]
    iso = pd.read_csv(path, delim_whitespace=True, names=header, comment='#')
    with open(path, 'r') as f:
        metadata_str = f.readline()
    metadata_str = metadata_str.strip('#').lstrip().rstrip().split(' ')[:3]
    metadata = dict()
    for i, data_elem in enumerate(metadata_str):
        keyvalue = data_elem.split('=')
        metadata[keyvalue[0]] = float(keyvalue[1])
    return iso, metadata
 
def get_trk_metadata(line):
    mass_search = re.compile('(?<=Total mass =  ).*?(?=\s)')
    mixing_length_search = re.compile('(?<=Mixing length = ).*?(?=\s)')
    EOS_search = re.compile('(?<=EOS = ).*?(?=\s)')
    Atm_search = re.compile('(?<=Atm = ).*?(?=\s)')
    Low_T_opacities_search = re.compile('(?<=Low T opacities = ).*?(?=\s)')
 
    mass = mass_search.search(line)
    mixing_length = mixing_length_search.search(line)
    EOS = EOS_search.search(line)
    Atm = Atm_search.search(line)
    Low_T_opacities = Low_T_opacities_search.search(line)
    meta = {'Mass':float(mass.group()), 'Mixing Length': float(mixing_length.group()),
            'EOS': EOS.group(), 'Atm':Atm.group(), 'Low T Opacities': Low_T_opacities.group()}
    return meta
 
def load_trk(path):
    header = [ "Model_#", "shells", "AGE", "log_L", "log_R", "log_g",
              "log_Teff", "Mconv_core", "Mconv_env", "Rconv_env",
              "M_He_core", "Xenv", "Zenv", "L_ppI", "L_ppII", "L_ppIII",
              "L_CNO", "L_triple-alpha", "L_He-C", "L_gravity", "L_neutrinos_old",
              "L_%_Grav_eng", "L_Itot", "C_log_T", "C_log_RHO", "C_log_P",
              "C_BETA", "C_ETA", "C_X", "C_Z", "C_H" "C_shell_midpoint",
              "C_H_shell_mass", "C_T_at_base_of_cz", "C_rho_at_base_of_cz","CA_He3",
              "CA_C12", "CA_C13", "CA_N14", "CA_N15", "CA_O16",
              "CA_O17", "CA_O18", "SA_He3", "SA_C12", "SA_C13", "SA_N14",
              "SA_N15", "SA_O16", "SA_O17","SA_O18", "N_pp", "N_pep", "N_hep",
              "N_Be7", "N_B8", "N_N13", "N_O15", "N_F17", "Cl37_flux",
              "Ga71_flux"]
    trk = pd.DataFrame(columns=header)
    with open(path, 'r') as trk_file:
        all_lines = trk_file.readlines()
    metadata = get_trk_metadata(all_lines[4])
    lines = [x.lstrip().rstrip().split() for x in all_lines[14:]]
    merged_lines = [x+y+z+w+q for x, y, z, w, q in zip(lines[::6], lines[1::6], lines[2::6], lines[3::6], lines[4::6], lines[5::6])]
    numeric_lines = [[float(y) for y in x] for x in merged_lines]
    for index, row in enumerate(numeric_lines):
        trk.loc[index] = row
    return trk, metadata
 
def load_trk_models(path):
    trks = list()
    metas = list()
    for file in os.listdir(path):
        if file.endswith('.trk'):
            print(file)
            trk, meta = load_trk(os.path.join(path, file))
            trks.append(trk)
            metas.append(meta)
    trks, metas = sort_based_on_key(trks, metas, key='Mass')
    return trks, metas
 
def load_iso_models(path):
    isos = list()
    metas = list()
    for file in os.listdir(path):
        if file.endswith('.iso'):
            iso, metadata = load_iso(os.path.join(path, file))
            isos.append(iso)
            metas.append(metadata)
    isos, metas = sort_based_on_key(isos, metas, key='M')
    return isos, metas
 
def load(path):
    functions = {'track': load_trk, 'iso': load_iso}
    suffix = args.path.split('.')[-1]
    data, metadata = functions[suffix](args.path)
    return data, metadata
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Load data from a dsep3 trk or iso file")
    parser.add_argument("path", help="path to file", type=str)
    parser.add_argument("-o", "--output", help="path to save as pickle", type=str)
 
    args = parser.parse_args()
 
    data, metadata = load(args.path)
 
    if args.output:
        data_package = {"data": data, "metadata": metadata}
        pickle.dump(data_package, open(args.output, "wb"))
    else:
        print("========= METADATA ===========")
        for key in metadata:
            print(f"{key}: {metadata[key]}")
 
        print("=========== DATA =============")
        for index, row in data.iterrows():
            print(row)