1"""Extract frequencies from a grid of GYRE output, and generate period spacing series."""
 2
 3import glob
 4import logging
 5import multiprocessing
 6from functools import partial
 7from pathlib import Path
 8
 9import pandas as pd
10
11from foam import support_functions as sf
12
13logger = logging.getLogger("logger.ffg")
14
15
16################################################################################
17def extract_frequency_grid(
18    gyre_files, output_file="pulsationGrid.hdf", parameters=["rot", "Z", "M", "logD", "aov", "fov", "Xc"], nr_cpu=None
19):
20    """
21    Extract frequencies from each globbed GYRE file and write them to 1 large file.
22
23    Parameters
24    ----------
25    gyre_files: string
26        String to glob to find all the relevant GYRE summary files.
27    output_file: string
28        Name (can include a path) for the file containing all the pulsation frequencies of the grid.
29    parameters: list of strings
30        List of parameters varied in the computed grid, so these are taken from the
31        name of the summary files, and included in the 1 file containing all the info of the whole grid.
32    nr_cpu: int
33        Number of worker processes to use in multiprocessing.
34        The default 'None' will use the number returned by os.cpu_count().
35    """
36    # make empty MultiProcessing listProxy
37    mp_list = multiprocessing.Manager().list()
38
39    # Glob all the files, then iteratively send them to a pool of processors
40    summary_files = glob.iglob(gyre_files)
41    with multiprocessing.Pool(nr_cpu) as p:
42        extract_func = partial(all_freqs_from_summary, parameters=parameters)
43        dictionaries = p.imap(extract_func, summary_files)
44        for new_row in dictionaries:
45            # Fill the listProxy with dictionaries for each read file
46            mp_list.append(new_row)
47
48        df = pd.DataFrame(data=list(mp_list))
49    # Sort the columns with frequencies by their radial order
50    column_list = list(df.columns[: len(parameters)])
51    column_list.extend(sorted(df.columns[len(parameters) :]))
52    df = df.reindex(column_list, axis=1)
53
54    # Generate the directory for the output file and write the file afterwards
55    Path(Path(output_file).parent).mkdir(parents=True, exist_ok=True)
56    df.to_hdf(path_or_buf=output_file, key="pulsation_grid", format="table", mode="w")
57
58
59################################################################################
60def all_freqs_from_summary(gyre_summary_file, parameters):
61    """
62    Extract model parameters and pulsation frequencies from a GYRE summary file
63
64    Parameters
65    ----------
66    gyre_summary_file: string
67        path to the GYRE summary file
68    parameters: list of strings
69        List of input parameters varied in the computed grid,
70        so these are read from the filename and included in returned line.
71
72    Returns
73    ----------
74    param_dict: dict
75        Dictionary containing all the model parameters and pulsation frequencies of the GYRE summary file.
76    """
77
78    _, data = sf.read_hdf5(gyre_summary_file)
79    param_dict = sf.get_param_from_filename(gyre_summary_file, parameters, values_as_float=True)
80
81    # Arrange increasing in radial order
82    for j in range(len(data["freq"]) - 1, -1, -1):
83        n_pg = data["n_pg"][j]
84        if abs(n_pg) < 10:
85            n_pg = f"{sf.sign(n_pg)}00{abs(n_pg)}"
86        elif abs(n_pg) < 100:
87            n_pg = f"{sf.sign(n_pg)}0{abs(n_pg)}"
88        param_dict.update({f"n_pg{n_pg}": data["freq"][j][0]})
89
90    return param_dict