1""" Create a summary file to store e.g. all history files of a MESA grid in a nested dictionary."""
  2
  3import logging
  4import os
  5import sys
  6
  7import h5py
  8import hdfdict
  9import numpy as np
 10
 11from foam import functions_for_mesa as ffm
 12from foam import support_functions as sf
 13
 14logger = logging.getLogger("logger.mg")
 15
 16
 17################################################################################
 18def _make_nested_dict(list_keys, value):
 19    """
 20    Recursively make a nested dictionary with the keys at different levels set by list_keys,
 21    and the most nested level containing the value.
 22
 23    Parameters
 24    ----------
 25    list_keys: list of keys
 26        List of keys for the different levels of the nested dictionary
 27    value: any
 28        The value, list, dictionary... that is coupled to the key
 29
 30    Returns
 31    ----------
 32    dict: dict
 33        A nested dictionary
 34    """
 35
 36    if len(list_keys) == 1:
 37        return {list_keys[0]: value}
 38    return {list_keys[0]: _make_nested_dict(list_keys[1:], value)}
 39
 40
 41################################################################################
 42
 43
 44class GridSummary:
 45    """
 46    Class to create and contain a summary of a MESA grid in a nested dictionary.
 47    """
 48
 49    def __init__(self, grid_parameters=None, Zsun=0.014):
 50        """
 51        Summary of the information of the grid of stellar models.
 52
 53        Parameters
 54        ----------
 55        grid_path: string
 56            Path to the stellar model grid.
 57        Zsun: float
 58            Solar metallicity value. Default = 0.014.
 59        """
 60        self.grid_parameters = grid_parameters
 61        self.grid_data = None
 62        self.Zsun = Zsun
 63
 64    ############################################################################
 65    def create_summary_file(
 66        self,
 67        dir_path,
 68        columns=["star_age", "log_L", "log_R", "log_Teff", "log_g"],
 69        magnitudes=False,
 70        output_name="stellar_grid.h5",
 71        file_ending="hist",
 72        files_directory_name="history",
 73    ):
 74        """
 75        Create a hdf5 summary file containing certain columns of all history files in a MESA grid.
 76
 77        Parameters
 78        ----------
 79        dir_path: string
 80            The path to the grid directory.
 81        columns: list of strings
 82            Names of the columns to be kept in the summary file.
 83        magnitudes: boolean
 84            If True, keep the absolute magnitude columns in the summary file.
 85        output_name: String
 86            Name for the summary file.
 87        """
 88        if os.path.isfile(output_name):
 89            logger.warning(f"Output file exists already: {output_name}")
 90            sys.exit()
 91
 92        summary_data = {}
 93        for subdir, dirs, files in os.walk(dir_path):
 94            # Skip the directory if there are no directories in it, and it's not the directory type we're looking for.
 95            if (len(dirs) == 0) and not subdir.endswith(files_directory_name):
 96                continue
 97            for file in files:
 98                if file.endswith(file_ending) and subdir.endswith(files_directory_name):
 99                    # Read in the files and get the main parameters from the filename
100                    _, data = ffm.read_mesa_file(os.path.join(subdir, file))
101                    filename_params = sf.get_param_from_filename(file, self.grid_parameters)
102
103                    # Keep only the specified columns and the absolute magnitudes (if magnitudes==True)
104                    for key in list(data.keys()):
105                        if magnitudes == True and key.startswith("abs_mag_"):
106                            continue
107                        elif key not in columns:
108                            del data[key]
109
110                    # Create a nested dictionary with the order of levels the same as the grid_parameters
111                    # and the column data dictionary at the bottom level
112                    keys = [filename_params[x] for x in self.grid_parameters]
113                    for i in range(len(self.grid_parameters)):
114                        param_list = keys
115                        param = param_list.pop(0)
116
117                        if i == 0:
118                            if param not in summary_data:
119                                summary_data[param] = _make_nested_dict(param_list, data)
120                                break
121                            else:
122                                nested_summary = summary_data[param]
123
124                        elif i == len(self.grid_parameters) - 1:
125                            nested_summary[param] = data
126
127                        else:
128                            if param not in nested_summary:
129                                nested_summary[param] = _make_nested_dict(param_list, data)
130                                break
131                            else:
132                                nested_summary = nested_summary[param]
133
134        self.grid_data = summary_data
135        self._set_param_ranges()
136        grid_info = {"grid_data": summary_data, "grid_parameters": self.grid_parameters}
137
138        with h5py.File(output_name, "w") as hdf5_file:
139            hdfdict.dump(grid_info, output_name)
140
141    ############################################################################
142    def read_summary_file(self, grid_path):
143        """
144        Read in a stellar model grid in hdf5 format created with create_summary_file.
145
146        Parameters
147        ----------
148        grid_path: string
149            Path to the stellar model grid.
150        """
151        # Load in the grid.
152        grid = hdfdict.load(grid_path, lazy=False)
153        self.grid_data = grid["grid_data"]
154        self.grid_parameters = [x.decode(sys.stdout.encoding) for x in grid["grid_parameters"]]
155        self._set_param_ranges()
156
157    ############################################################################
158    def _set_param_ranges(self):
159        """
160        Set the attributes holding the ranges of all grid parameters
161        """
162        dictionary = self.grid_data
163        for parameter in self.grid_parameters:
164            dict_keys = sorted(list(dictionary.keys()))
165            values = np.array([x for x in dict_keys])
166            # Highest mass values have largest range in logDext, hence use highest key of each parameter
167            dictionary = dictionary[dict_keys[-1]]
168            setattr(self, f"{parameter}_array", values)