1"""A few helpful functions to process MESA output."""
  2
  3import glob
  4import multiprocessing
  5from functools import partial
  6from pathlib import Path
  7
  8import h5py
  9import numpy as np
 10import pandas as pd
 11
 12from foam import support_functions as sf
 13
 14
 15################################################################################
 16def read_mesa_file(file_path, index_col=None):
 17    """
 18    Read in a mesa profile or history file and return 2 dictionaries.
 19    The first with the header info, and the second with the data.
 20    If the format is hdf5, assumes the attributes are the MESA header.
 21    If the format is not an hdf5 file, assumes the default MESA output format.
 22    This format is an ascii file delimited by whitespace with the following structure:
 23    line 1: header names
 24    line 2: header data
 25    line 3: blank
 26    line 4: main data names
 27    line >4:main data values
 28
 29    Parameters
 30    ----------
 31    file_path: String
 32        The path to the MESA profile or history file to read in.
 33    index_col: Hashable
 34        Column to use as row labels.
 35
 36    Returns
 37    ----------
 38    header: dict
 39        A dictionary holding the header info of the MESA file.
 40    data: dict
 41        A dictionary holding the data columns of the MESA file as numpy arrays.
 42    """
 43    if h5py.is_hdf5(file_path):
 44        return sf.read_hdf5(file_path)
 45
 46    else:  # assumes the default MESA output format
 47        header_df = pd.read_table(file_path, sep="\s+", nrows=1, header=1)
 48        data_df = pd.read_table(file_path, sep="\s+", skiprows=3, header=1, index_col=index_col)
 49
 50        header = {}
 51        for k in header_df.keys():
 52            header.update({k: header_df[k].to_numpy()[0]})
 53        data = {}
 54        for k in data_df.keys():
 55            data.update({k: data_df[k].to_numpy()})
 56
 57        return header, data
 58
 59
 60################################################################################
 61def calculate_number_densities(hist_file):
 62    """
 63    Calculate surface number densities for all isotopes in the MESA grid.
 64    All isotopes in the used nuclear network need to be be written in the history file,
 65    otherwise this will function give wrong numbers.
 66
 67    Parameters
 68    ----------
 69    hist_file: string
 70        The path to the MESA history file.
 71
 72    Returns
 73    ----------
 74    number_densities: dict
 75        Column keys specify the element (surf_X_per_N_tot), values are number densities of that element.
 76    """
 77    _, data = read_mesa_file(hist_file)
 78    element_list = {}
 79    number_densities = {}
 80    inverse_average_atomic_mass = np.zeros(len(data[list(data.keys())[0]]))
 81    for column_name in data.keys():
 82        if "_per_Mass_tot" in column_name:
 83            element_list.update({column_name: data[column_name]})
 84            inverse_average_atomic_mass += data[column_name]
 85
 86    average_atomic_mass = inverse_average_atomic_mass ** (-1)
 87    for key in element_list.keys():
 88        number_densities.update({key.replace("_per_Mass_tot", "_per_N_tot"): element_list[key] * average_atomic_mass})
 89
 90    return number_densities
 91
 92
 93################################################################################
 94def extract_surface_grid(
 95    mesa_profiles,
 96    output_file="surfaceGrid.hdf",
 97    parameters=["Z", "M", "logD", "aov", "fov", "Xc"],
 98    nr_cpu=None,
 99    additional_observables=None,
100):
101    """
102    Extract 'logTeff', 'logL', 'logg', 'age', and extra requested info for each globbed MESA profile and write them to 1 large file.
103
104    Parameters
105    ----------
106    mesa_profiles: string
107        String to glob to find all the relevant MESA profiles.
108    output_file: string
109        Name (or path) for the file containing all the pulsation frequencies of the grid.
110    parameters: list of strings
111        List of parameters varied in the computed grid, so these are taken from the
112        name of the profile files, and included in the output file containing the info of the whole grid.
113    nr_cpu: int
114        Number of worker processes to use in multiprocessing. The default 'None' will use the number returned by os.cpu_count().
115    additional_observables: list of strings
116        List of observables to add to the surface grid. Must correspond to mesa profile header-item names.
117    """
118    # Make list of extra observables requested by the user
119    if additional_observables is None:
120        additional_observables = []
121    extras_to_be_extracted = [x for x in additional_observables if x not in ["logTeff", "logL", "logg", "age"]]
122
123    extract_func = partial(info_from_profiles, parameters=parameters, extra_header_items=extras_to_be_extracted)
124    # Glob all the files, then iteratively send them to a pool of processors
125    profiles = glob.iglob(mesa_profiles)
126    with multiprocessing.Pool(nr_cpu) as p:
127        surface = p.imap(extract_func, profiles)
128
129        # Generate the directory for the output file and write the file afterwards
130        Path(Path(output_file).parent).mkdir(parents=True, exist_ok=True)
131        # make a new list, so 'parameters' is not extended before passing it on to 'info_from_profiles'
132        header_parameters = list(parameters)
133        header_parameters.extend(["logTeff", "logL", "logg", "age"])
134        # Add the extra observables requested by the user
135        header_parameters.extend(extras_to_be_extracted)
136
137        # Make list of lists, put it in a dataframe, and write to a file
138        data = []
139        for line in surface:
140            data.append(line)
141
142    df = pd.DataFrame(data=data, columns=header_parameters)
143    df.to_hdf(path_or_buf=output_file, key="surface_grid", format="table", mode="w")
144
145
146################################################################################
147def info_from_profiles(mesa_profile, parameters, extra_header_items):
148    """
149    Extract 'logTeff', 'logL', 'logg', 'age', and extra requested info from a MESA profile and the model parameters from its filename.
150
151    Parameters
152    ----------
153    mesa_profile: string
154        path to the MESA profile
155    parameters: list of strings
156        List of input parameters varied in the computed grid, so these are read from the filename and included in returned line.
157    extra_header_items: list of strings
158        List of extra observables to add to the surface grid. Must correspond to mesa profile header-item names.
159
160    Returns
161    ----------
162    line: string
163        Line containing all the model- and surface parameters of the MESA profile.
164    """
165    param_dict = sf.get_param_from_filename(mesa_profile, parameters, values_as_float=True)
166    prof_header, prof_data = read_mesa_file(mesa_profile)
167
168    logL = np.log10(float(prof_header["photosphere_L"]))
169    logTeff = np.log10(float(prof_header["Teff"]))
170    logg = prof_data["log_g"][0]
171    age = int(float(prof_header["star_age"]))
172
173    line = []
174    for p in parameters:
175        line.append(param_dict[p])
176    line.extend([logTeff, logL, logg, age])
177
178    for obs in extra_header_items:  # Add the extra observables requested by the user
179        item = float(prof_header[obs])
180        line.append(item)
181
182    return line