1"""Helpful functions in general. Reading HDF5, processing strings, manipulating dataFrames..."""
  2
  3import logging
  4import re
  5from pathlib import Path
  6
  7import h5py
  8import pandas as pd
  9
 10logger = logging.getLogger("logger.sf")
 11
 12
 13################################################################################
 14def split_line(line, sep):
 15    """
 16    Splits a string in 2 parts.
 17
 18    Parameters
 19    ----------
 20    line: string
 21        String to split in 2.
 22    sep: string
 23        Separator where the string has to be split around.
 24
 25    Returns
 26    ----------
 27    head: string
 28        Part 1 of the string before the separator.
 29    tail: string
 30        Part 2 of the string after the separator.
 31    """
 32    head, sep_, tail = line.partition(sep)
 33    assert sep_ == sep
 34    return head, tail
 35
 36
 37################################################################################
 38def substring(line, sep_first, sep_second):
 39    """
 40    Get part of a string between 2 specified separators.
 41    If second separator is not found, return everything after first separator.
 42
 43    Parameters
 44    ----------
 45    line: string
 46        String to get substring from.
 47    sep_first: string
 48        First separator after which the returned substring should start.
 49    sep_first: string
 50        Second separator at which the returned substring should end.
 51
 52    Returns
 53    ----------
 54    head: string
 55        Part of the string between the 2 separators.
 56    """
 57    head, tail = split_line(line, sep=sep_first)
 58    if sep_second not in tail:
 59        return tail
 60    head, tail = split_line(tail, sep=sep_second)
 61    return head
 62
 63
 64################################################################################
 65def get_param_from_filename(file_path, parameters, values_as_float=False):
 66    """
 67    Get parameters from filename
 68
 69    Parameters
 70    ----------
 71    file_path : string
 72        Full path to the file
 73    parameters: list of string
 74        Names of parameters to extract from filename
 75
 76    Returns
 77    ----------
 78    param_dict: dict
 79        Keys are strings describing the parameter, values are strings giving corresponding parameter values
 80    """
 81
 82    param_dict = {}
 83    for parameter in parameters:
 84        try:
 85            p = substring(Path(file_path).stem, parameter, "_")
 86            if values_as_float:
 87                p = float(p)
 88            param_dict[parameter] = p
 89        except:
 90            logger.warning(
 91                f"In get_param_from_filename: parameter \"{parameter}\" not found in '{file_path}', value not added"
 92            )
 93
 94    return param_dict
 95
 96
 97################################################################################
 98def read_hdf5(filename):
 99    """
100    Read a HDF5-format file (e.g. GYRE)
101
102    Parameters
103    ----------
104    filename : string
105        Input file
106
107    Returns
108    ----------
109    attributes: dict
110        Dictionary containing the attributes of the file.
111    data: dict
112        Dictionary containing the data from the file as numpy arrays.
113    """
114    # Open the file
115    with h5py.File(filename, "r") as file:
116        # Read attributes
117        attributes = dict(zip(file.attrs.keys(), file.attrs.values()))
118        # Read datasets
119        data = {}
120        for k in file.keys():
121            data[k] = file[k][...]
122    return attributes, data
123
124
125################################################################################
126def sign(x):
127    """
128    Returns the sign of a number as a string
129
130    Parameters
131    ----------
132    x: float or int
133
134    Returns
135    ----------
136    s: string
137        A string representing the sign of the number
138    """
139    if abs(x) == x:
140        return "+"
141    else:
142        return "-"
143
144
145################################################################################
146def get_subgrid_dataframe(file_to_read, fixed_params=None):
147    """
148    Read a hdf5 file containing the grid information as a pandas dataframe.
149    Parameters can be fixed to certain values to filter out entries with other values of that parameter.
150
151    Parameters
152    ----------
153    file_to_read: string
154        path to the file to read
155    fixed_params: dict
156        keys are parameters to fix to the value specified in the dictionary
157
158    Returns
159    ----------
160    df: pandas DataFrame
161    """
162    df = pd.read_hdf(file_to_read)
163
164    if fixed_params is not None:
165        for param in fixed_params.keys():
166            indices_to_drop = df[df[param] != fixed_params[param]].index
167            df.drop(indices_to_drop, inplace=True)
168        df.reset_index(drop=True, inplace=True)
169
170    return df
171
172
173################################################################################
174def add_surface_to_puls_grid(
175    grid_frequencies,
176    grid_surface,
177    output_name=" grid_surface+freq.hdf",
178    grid_parameters=["Z", "M", "logD", "aov", "fov", "Xc"],
179):
180    """
181    Combine the output files with the frequencies and surface info of the grid in one new file,
182    only keeping models that have entries in both the grid files.
183
184    Parameters
185    ----------
186    grid_frequencies: string
187        Path to the file containing the model input parameters and corresponding frequencies of the model.
188    grid_surface: string
189        Path to the file containing the model input parameters and corresponding surface info of the model.
190    output_name: string
191        Name of the generated file containing the combined info.
192    grid_parameters: list of string
193        List of the model parameters to use for matching the entries in the freq/surface file.
194    """
195    freq_df = pd.read_hdf(grid_frequencies)
196    surface_df = pd.read_hdf(grid_surface)
197    # Merge with surface info first, freq info second. Only keeping rows that both dataFrames have in common based on the 'on' columns.
198    df_merged = pd.merge(surface_df, freq_df, how="inner", on=grid_parameters)
199
200    _ = df_merged.pop("age")  # Don't add the age in the combined file
201
202    # take the column with rotation and place it as the first column, and its error as second column
203    col = df_merged.pop("rot")
204    df_merged.insert(0, col.name, col)
205    col = df_merged.pop("rot_err")
206    df_merged.insert(1, col.name, col)
207    # write the merged dataFrame to a new file
208    df_merged.to_hdf(path_or_buf=output_name, key="pulsation_surface_grid", format="table", mode="w")