1"""Selection of the theoretical pulsation patterns that best match the observations,
  2   with the possibility to optimise the rotation rate in the process through rescaling."""
  3
  4import logging
  5import multiprocessing
  6import sys
  7from functools import partial
  8from pathlib import Path
  9
 10import astropy.units as u
 11import matplotlib.pyplot as plt
 12import numpy as np
 13import pandas as pd
 14from lmfit import Minimizer, Parameters
 15
 16logger = logging.getLogger("logger.bop")
 17
 18
 19################################################################################
 20def generate_spacing_series(periods, errors=None):
 21    """
 22    Generate the period spacing series (delta P = p_(n+1) - p_n )
 23
 24    Parameters
 25    ----------
 26    periods: list of floats
 27        Periods in units of days
 28     errors (optional): list of floats
 29        Errors on periods in units of days
 30    Returns
 31    ----------
 32    observed_spacings, observed_spacings_errors: tuple of lists of floats
 33        period spacing series (delta P values) and its errors (if supplied) in units of seconds
 34    """
 35    spacings = []
 36    if errors is None:
 37        spacings_errors = None
 38        for n, period_n in enumerate(periods[:-1]):
 39            spacings.append(abs(period_n - periods[n + 1]) * 86400.0)
 40    else:
 41        spacings_errors = []
 42        for n, period_n in enumerate(periods[:-1]):
 43            spacings.append(abs(period_n - periods[n + 1]) * 86400.0)
 44            spacings_errors.append(np.sqrt(errors[n] ** 2 + errors[n + 1] ** 2) * 86400.0)
 45
 46    return spacings, spacings_errors
 47
 48
 49################################################################################
 50def construct_theoretical_puls_pattern(
 51    pulsation_grid_file,
 52    observations_file,
 53    method_build_series,
 54    pattern_starting_pulsation=[],
 55    which_observable="period",
 56    output_file="theoretical_frequency_patterns.hdf",
 57    asymptotic_object=None,
 58    estimated_rotation=None,
 59    grid_parameters=None,
 60    nr_cpu=None,
 61):
 62    """
 63    Construct the theoretical frequency pattern for each model in the grid, which correspond to the observed pattern.
 64    (Each theoretical model is a row in 'pulsation_grid_file'.)
 65    The rotation rate will be scaled and optimised for each theoretical pattern individually.
 66    This optimisation will not be performed if asymptotic_object=None.
 67
 68    Parameters
 69    ----------
 70    pulsation_grid_file: string
 71        path to file containing input parameters of the models, and the pulsation frequencies of those models
 72        (as generated by function 'extract_frequency_grid' in 'functions_for_gyre').
 73    observations_file: string
 74        Path to the tsv file with observations, with a column for each observable and each set of errors.
 75        Column names specify the observable, and "_err" suffix denotes that it's the error.
 76    method_build_series: string
 77        way to generate the theoretical frequency pattern from each model to match the observed pattern. Options are:
 78            provided-pulsation: build pattern from the provided pulsation    (function 'puls_series_from_given_puls')
 79            highest-frequency: build pattern from the observed highest frequency    (function 'puls_series_from_given_puls')
 80            chisq-longest-sequence: build pattern based on longest, best matching sequence of pulsations (function 'chisq_longest_sequence')
 81    pattern_starting_pulsation: list of floats
 82        Only needed if you set method_build_series=provided-pulsation
 83        Value of the pulsation to start building the pattern from, one for each separated part of the pattern.
 84        The unit of this value needs to be the same as the observable set through which_observable.
 85    which_observable: string
 86        Observable used in the theoretical pattern construction, options are 'frequency' or 'period'.
 87    output_file: string
 88        Name (can include a path) for the file containing all the pulsation frequencies of the grid.
 89    asymptotic_object: asymptotic (see 'gmode_rotation_scaling')
 90        Object to calculate g-mode period spacing patterns in the asymptotic regime using the TAR.
 91    estimated_rotation: float
 92        Estimation of the rotation rate of the star, used as initial value in the optimisation problem.
 93    grid_parameters: list of string
 94        List of the parameters in the theoretical grid.
 95    nr_cpu: int
 96        Number of worker processes to use in multiprocessing. The default 'None' will use the number returned by os.cpu_count().
 97    """
 98    # Read in the files with observed and theoretical frequencies as pandas DataFrames
 99    obs_dataframe = pd.read_table(observations_file, sep="\s+", header=0, index_col="index")
100    theory_dataframe = pd.read_hdf(pulsation_grid_file)
101
102    obs = np.asarray(obs_dataframe[which_observable])
103    obs_err = np.asarray(obs_dataframe[f"{which_observable}_err"])
104
105    # if frequency was filled in as 0, it indicates an interruption in the pattern
106    missing_puls = np.where(obs == 0)[0]
107    # remove values indicating interruptions in the pattern
108    obs_without_missing = obs[obs != 0]
109    # remove values indicating interruptions in the pattern
110    obs_err_without_missing = obs_err[obs_err != 0]
111    # Adjust indices for removed 0-values of missing frequencies
112    missing_puls = [missing_puls[i] - i for i in range(len(missing_puls))]
113
114    # split into different parts of the interrupted pattern
115    obs_pattern_parts = np.split(obs_without_missing, missing_puls)
116    obs_err_pattern_parts = np.split(obs_err_without_missing, missing_puls)
117
118    # partial function fixes all parameters of the function except for 1 that is iterated over in the multiprocessing pool.
119    theory_pattern_func = partial(
120        theoretical_pattern_from_dfrow,
121        obs=obs,
122        obs_pattern_parts=obs_pattern_parts,
123        obs_err_pattern_parts=obs_err_pattern_parts,
124        which_observable=which_observable,
125        method_build_series=method_build_series,
126        pattern_starting_pulsation=pattern_starting_pulsation,
127        grid_parameters=grid_parameters,
128        asymptotic_object=asymptotic_object,
129        estimated_rotation=estimated_rotation,
130        plot_rotation_optimisation=False,
131    )
132
133    # Make the output file directory
134    Path(Path(output_file).parent).mkdir(parents=True, exist_ok=True)
135    # Send the rows of the dataframe iteratively to a pool of processors to get the theoretical pattern for each model
136    with multiprocessing.Pool(nr_cpu) as p:
137        freqs = p.imap(theory_pattern_func, theory_dataframe.iterrows())
138        header_parameters = ["rot", "rot_err"] + grid_parameters
139
140        for i in range(1, obs_dataframe.shape[0] + 1):
141            if i - 1 in np.where(obs_dataframe.index == "f_missing")[0]:
142                f = f"{which_observable}_missing"
143            else:
144                f = f"{which_observable}{i}"
145            header_parameters.append(f.strip())
146
147        data = []
148        for line in freqs:
149            data.append(line)
150
151    df = pd.DataFrame(data=data, columns=header_parameters)
152    df.to_hdf(path_or_buf=output_file, key="selected_puls_grid", format="table", mode="w")
153
154
155# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
156def theoretical_pattern_from_dfrow(
157    summary_grid_row,
158    obs,
159    obs_pattern_parts,
160    obs_err_pattern_parts,
161    which_observable,
162    method_build_series,
163    pattern_starting_pulsation=[],
164    asymptotic_object=None,
165    estimated_rotation=None,
166    plot_rotation_optimisation=False,
167    grid_parameters=None,
168):
169    """
170    Extract model parameters and a theoretical pulsation pattern from a row of the dataFrame that contains all model parameters and pulsation frequencies.
171
172    Parameters
173    ----------
174    summary_grid_row: tuple, made of (int, pandas series)
175        tuple returned from pandas.iterrows(), first tuple entry is the row index of the pandas dataFrame
176        second tuple entry is a pandas series, containing a row from the pandas dataFrame.
177        (This row holds model parameters and pulsation frequencies.)
178    obs: numpy array, dtype=float
179        Array of observed frequencies or periods. (Ordered increasing in frequency.)
180    obs_pattern_parts : list of numpy array, dtype=float
181        Holds a numpy array (with the observed frequencies or periods) per split off part
182        of the observed pattern if it is interrupted.
183        The list contains only one array if the observed pattern is uninterrupted.
184    obs_err_pattern_parts : list of numpy array, dtype=float
185        The errors on the data in obs_pattern_parts.
186    which_observable: string
187        Which observables are used in the pattern building, options are 'frequency' or 'period'.
188    method_build_series: string
189        way to generate the theoretical frequency pattern from each model
190        to match the observed pattern. Options are:
191            provided-pulsation: build pattern from the provided pulsation    (function 'puls_series_from_given_puls')
192            highest-frequency: build pattern from the observed highest frequency    (function 'puls_series_from_given_puls')
193            chisq-longest-sequence: build pattern based on longest, best matching sequence of pulsations    (function 'chisq_longest_sequence')
194    pattern_starting_pulsation: list of floats
195        Only needed if you set method_build_series=provided-pulsation
196        Value of the pulsation to start building the pattern from, one for each separated part of the pattern.
197        The unit of this value needs to be the same as the observable set through which_observable.
198    asymptotic_object: asymptotic (see 'gmode_rotation_scaling')
199        Object to calculate g-mode period spacing patterns in the asymptotic regime using the TAR.
200    estimated_rotation: float
201        Estimation of the rotation rate of the star, used as initial value in the optimisation problem.
202    grid_parameters: list of string
203        List of the parameters in the theoretical grid.
204
205    Returns
206    ----------
207    list_out: list of float
208        The input parameters and pulsation frequencies of the theoretical pattern
209        (or periods, depending on 'which_observable').
210    """
211    # all keys containing n_pg (these are all the radial orders)
212    freqs = np.asarray(summary_grid_row[1].filter(like="n_pg"))
213    # array with radial orders
214    orders = np.asarray([int(o.replace("n_pg", "")) for o in summary_grid_row[1].filter(like="n_pg").index])
215    orders = orders[~np.isnan(freqs)]
216    # remove all entries that are NaN in the numpy array (for when the models have a different amount of computed modes)
217    freqs = freqs[~np.isnan(freqs)]
218
219    # Check if pattern_starting_pulsation has enough entries to not truncate other parts in the zip function.
220    if len(obs_pattern_parts) != len(pattern_starting_pulsation):
221        if method_build_series == "provided-pulsation":
222            sys.exit(
223                logger.error(
224                    "Amount of pulsations specified to build patterns from is not equal to the amount of split-off parts in the pattern."
225                )
226            )
227        # Content of pattern_starting_pulsation doesn't matter if it's not used to build the pattern.
228        else:
229            # We only care about the length if the method doesn't use specified pulsations.
230            pattern_starting_pulsation = [None] * len(obs_pattern_parts)
231
232    # In this case, rescaling nor optimisation will happen
233    if asymptotic_object is None:
234        residual = rescale_rotation_and_select_theoretical_pattern(
235            None,
236            asymptotic_object,
237            estimated_rotation,
238            freqs,
239            orders,
240            obs,
241            obs_pattern_parts,
242            obs_err_pattern_parts,
243            which_observable,
244            method_build_series,
245            pattern_starting_pulsation,
246        )
247
248        list_out = [estimated_rotation, 0]
249        for parameter in grid_parameters:
250            list_out.append(summary_grid_row[1][parameter])
251
252        selected_pulsations = obs + residual
253        list_out.extend(selected_pulsations)
254
255    else:
256        # Optimise the rotation rate and get the pulsations at that rotation rate
257        params = Parameters()
258        params.add("rotation", value=estimated_rotation, min=1e-5)
259        # Fit rotation to observed pattern with the default leastsq algorithm
260        optimise_rotation = Minimizer(
261            rescale_rotation_and_select_theoretical_pattern,
262            params,
263            fcn_args=(
264                asymptotic_object,
265                estimated_rotation,
266                freqs,
267                orders,
268                obs,
269                obs_pattern_parts,
270                obs_err_pattern_parts,
271                which_observable,
272                method_build_series,
273                pattern_starting_pulsation,
274            ),
275        )
276
277        result_minimizer = optimise_rotation.minimize()
278
279        # Search from second initial value, which is as separated from the first solution as the first initial guess.
280        search_second_initial_value = 2 * result_minimizer.params["rotation"].value - estimated_rotation
281        params = Parameters()
282        params.add("rotation", value=search_second_initial_value, min=1e-5)
283        optimise_rotation = Minimizer(
284            rescale_rotation_and_select_theoretical_pattern,
285            params,
286            fcn_args=(
287                asymptotic_object,
288                estimated_rotation,
289                freqs,
290                orders,
291                obs,
292                obs_pattern_parts,
293                obs_err_pattern_parts,
294                which_observable,
295                method_build_series,
296                pattern_starting_pulsation,
297            ),
298        )
299
300        result_minimizer2 = optimise_rotation.minimize()
301        # Select the minimizer with the initial guess that resulted in the best fit.
302        if result_minimizer2.chisqr < result_minimizer.chisqr:
303            result_minimizer = result_minimizer2
304
305        optimised_pulsations = result_minimizer.residual + obs
306
307        if result_minimizer.message != "Fit succeeded.":
308            logger.warning(
309                f"""Fitting rotation did not succeed: {result_minimizer.message}
310                            for model {summary_grid_row[1].drop(summary_grid_row[1].filter(regex='n_pg').index).drop('rot').to_dict()} using method: {method_build_series}
311                            rotation found: {result_minimizer.params['rotation'].value} with error: {result_minimizer.params['rotation'].stderr}"""
312            )
313
314        if plot_rotation_optimisation:
315            fig1, ax1 = plt.subplots()
316            if which_observable == "frequency":
317                obsperiod = 1 / obs
318                optimised_periods = 1 / optimised_pulsations
319            else:
320                obsperiod = obs
321                optimised_periods = optimised_pulsations
322
323            # Recombine observational errors in one array
324            combined_obs_err = []
325            for obs_err_part in obs_err_pattern_parts:
326                if len(combined_obs_err) > 0:
327                    combined_obs_err.extend([0])
328                combined_obs_err.extend(obs_err_part)
329            combined_obs_err = np.asarray(combined_obs_err)
330
331            spacings = generate_spacing_series(obsperiod, combined_obs_err)
332            ax1.errorbar(obsperiod[:-1], spacings[0], fmt="o", yerr=spacings[1], label="obs", color="blue", alpha=0.8)
333            ax1.plot(obsperiod[:-1], spacings[0], color="blue")
334            ax1.plot(
335                optimised_periods[:-1],
336                generate_spacing_series(optimised_periods)[0],
337                "*",
338                ls="solid",
339                color="orange",
340                label="optimised",
341            )
342            ax1.plot(
343                1 / freqs[:-1], generate_spacing_series(1 / freqs)[0], ".", ls="solid", label="initial", color="green"
344            )
345
346            fig2, ax2 = plt.subplots()
347
348            ax2.errorbar(obs, obs, fmt="o", xerr=combined_obs_err, label="obs", color="blue", alpha=0.8)
349            ax2.plot(optimised_periods, optimised_periods, "*", color="orange", label="optimised")
350            ax2.plot(1 / freqs, 1 / freqs, ".", label="initial", color="green")
351
352            ax1.legend(prop={"size": 14})
353            ax2.legend(prop={"size": 14})
354            ax1.set_title(
355                f"initial omega = {estimated_rotation}, optimised omega = {result_minimizer.params['rotation'].value}"
356            )
357            ax2.set_title(
358                f"initial omega = {estimated_rotation}, optimised omega = {result_minimizer.params['rotation'].value}"
359            )
360            plt.show()
361
362        # Create list with rotation, its error, all the input parameters, and the optimised pulsations
363        list_out = [result_minimizer.params["rotation"].value, result_minimizer.params["rotation"].stderr]
364        for parameter in grid_parameters:
365            list_out.append(summary_grid_row[1][parameter])
366        list_out.extend(optimised_pulsations)
367
368    return list_out
369
370
371# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
372def rescale_rotation_and_select_theoretical_pattern(
373    params,
374    asymptotic_object,
375    estimated_rotation,
376    freqs_input,
377    orders,
378    obs,
379    obs_pattern_parts,
380    obs_err_pattern_parts,
381    which_observable,
382    method_build_series,
383    pattern_starting_pulsation,
384):
385    """
386    If rotation is not optimised in the modelling, asymptotic_object should be set to 'None' and
387    this function will just select a theoretical pulsation pattern based on the specified method.
388    If asymptotic_object is supplied, the rotation will be optimised and this will be the
389    objective function for the lmfit Minimizer class to minimise the output of.
390    Rescales the theoretical pulsation pattern using the asymptotic object from
391    'gmode_rotation_scaling', and selects a theoretical pulsation pattern based
392    on the specified method.
393
394    Parameters
395    ----------
396    params: Parameter (object of lmfit)
397        The parameter used to optimise the fit. In this case rotation in units of 1/day.
398    asymptotic_object: asymptotic (see 'gmode_rotation_scaling')
399        Object to calculate g-mode period spacing patterns in the asymptotic regime using the TAR.
400    estimated_rotation: float
401        Estimated rotation rate in units of 1/day. Used as initial value in the optimisation problem.
402    freqs_input: numpy array, dtype=float
403        Array of the frequencies as computed by GYRE, to be scaled to the optimal rotation rate.
404    orders: numpy array, dtype=int
405        Array with the radial orders of the theoretical input frequencies.
406    obs: numpy array, dtype=float
407        Array of observed frequencies or periods. (Ordered increasing in frequency.)
408    obs_pattern_parts: list of numpy array, dtype=float
409        Holds a numpy array with observed frequencies or periods
410        per split off part of the observed pattern if it is interrupted.
411        The list contains only one array if the observed pattern is uninterrupted.
412    obs_err_pattern_parts: list of numpy array, dtype=float
413        The errors on obs_pattern_parts
414    which_observable: string
415        Which observables are used in the pattern building, options are 'frequency' or 'period'.
416    method_build_series: string
417        way to generate the theoretical frequency pattern from each model
418        to match the observed pattern. Options are:
419            provided-pulsation: build pattern from the provided pulsation    (function 'puls_series_from_given_puls')
420            highest-frequency: build pattern from the observed highest frequency    (function 'puls_series_from_given_puls')
421            chisq-longest-sequence: build pattern based on longest, best matching sequence of pulsations    (function 'chisq_longest_sequence')
422    pattern_starting_pulsation: list of floats
423        Only needed if you set method_build_series=provided-pulsation
424        Value of the pulsation to start building the pattern from, one for each separated part of the pattern.
425        The unit of this value needs to be the same as the observable set through which_observable.
426
427    Returns
428    ----------
429    (output_pulsations - obs): numpy array, dtype=float
430        Differences between the scaled pulsations and the observations.
431        The array to be minimised by the lmfit Minimizer if rotation is optimised.
432    """
433    if asymptotic_object is not None:
434        v = params.valuesdict()
435        # To avoid division by zero in scale_pattern
436        if estimated_rotation == 0:
437            estimated_rotation = 1e-99
438        freqs = asymptotic_object.scale_pattern(freqs_input / u.d, estimated_rotation / u.d, v["rotation"] / u.d) * u.d
439        # convert astropy quantities back to floats
440        freqs = np.asarray(freqs, dtype=np.float64)
441    else:
442        freqs = freqs_input
443
444    periods = 1 / freqs
445
446    output_pulsations = []
447    for obs_part, obs_err_part, pattern_starting_pulsation_part in zip(
448        obs_pattern_parts, obs_err_pattern_parts, pattern_starting_pulsation
449    ):
450        # To indicate interruptions in the pattern
451        if len(output_pulsations) > 0:
452            output_pulsations.append(0)
453
454        if which_observable == "frequency":
455            # remove frequencies that were already chosen in a different, split-off part of the pattern
456            if len(output_pulsations) > 0:
457                # If input is in increasing radial order (decreasing n_pg, since n_pg is negative for g-modes)
458                if orders[1] == orders[0] - 1:
459                    # index -2 to get lowest, non-zero freq
460                    np.delete(freqs, np.where(freqs >= output_pulsations[-2]))
461                # If input is in decreasing radial order
462                else:
463                    np.delete(freqs, np.where(freqs <= max(output_pulsations)))
464            theory_value = freqs
465            obs_period = 1 / obs_part
466            obs_period_err = obs_err_part / obs_part**2
467            highest_obs_freq = max(obs_part)
468
469        elif which_observable == "period":
470            # remove periods that were already chosen in a different, split-off part of the pattern
471            if len(output_pulsations) > 0:
472                # If input is in increasing radial order (decreasing n_pg, since n_pg is negative for g-modes)
473                if orders[1] == orders[0] - 1:
474                    np.delete(periods, np.where(periods <= max(output_pulsations)))
475                # If input is in decreasing radial order
476                else:
477                    # index -2 to get lowest, non-zero period
478                    np.delete(periods, np.where(periods >= output_pulsations[-2]))
479            theory_value = periods
480            obs_period = obs_part
481            obs_period_err = obs_err_part
482            # highest frequency is lowest period
483            highest_obs_freq = min(obs_part)
484        else:
485            sys.exit(logger.error("Unknown observable to fit"))
486
487        if method_build_series == "provided-pulsation":
488            selected_theoretical_pulsations = puls_series_from_given_puls(
489                theory_value, obs_part, pattern_starting_pulsation_part
490            )
491        elif method_build_series == "highest-frequency":
492            selected_theoretical_pulsations = puls_series_from_given_puls(theory_value, obs_part, highest_obs_freq)
493        elif method_build_series == "chisq-longest-sequence":
494            series_chi2, final_theoretical_periods, corresponding_orders = chisq_longest_sequence(
495                periods, orders, obs_period, obs_period_err
496            )
497            if which_observable == "frequency":
498                selected_theoretical_pulsations = 1 / np.asarray(final_theoretical_periods)
499            elif which_observable == "period":
500                selected_theoretical_pulsations = final_theoretical_periods
501        else:
502            sys.exit(logger.error(f"Unrecognised method to build pulsation series: {method_build_series}"))
503
504        output_pulsations.extend(selected_theoretical_pulsations)
505
506    return output_pulsations - obs
507
508
509################################################################################
510def puls_series_from_given_puls(theory_in, obs, obs_to_build_from, plot=False):
511    """
512    Generate a theoretical pulsation pattern (can be in frequency or period) from the given observations.
513    Build consecutively in radial order, starting from the theoretical value closest to the provided observational value.
514
515    Parameters
516    ----------
517    theory_in: numpy array, dtype=float
518        Array of theoretical frequencies or periods.
519    obs: numpy array, dtype=float
520        Array of observed frequencies or periods.
521    obs_to_build_from: float
522        observed frequency or period value to start building the pattern from.
523    plot: boolean
524        Make a period spacing diagram for the constructed series.
525
526    Returns
527    ----------
528    theory_sequence: list of float
529        The constructed theoretical frequency pattern
530    """
531    # get index of observation to build the series from
532    nth_obs = np.where(obs == obs_to_build_from)[0][0]
533    # search theoretical freq closest to the given observed one
534    diff = abs(theory_in - obs_to_build_from)
535    # get index of this theoretical frequency
536    index = np.where(diff == min(diff))[0][0]
537
538    # Insert a value of -1 if observations miss a theoretical counterpart in the beginning
539    theory_sequence = []
540    if (index - nth_obs) < 0:
541        for i in range(abs((index - nth_obs))):
542            theory_sequence.append(-1)
543        theory_sequence.extend(theory_in[0 : index + (len(obs) - nth_obs)])
544    else:
545        theory_sequence.extend(theory_in[index - nth_obs : index + (len(obs) - nth_obs)])
546
547    # Insert a value of -1 if observations miss a theoretical counterpart at the end
548    if index + (len(obs) - nth_obs) > len(theory_in):
549        for i in range((index + (len(obs) - nth_obs)) - len(theory_in)):
550            theory_sequence.append(-1)
551
552    if plot is True:
553        fig = plt.figure()
554        ax = fig.add_subplot(111)
555        theory = np.asarray(theory_sequence)
556        ax.plot((1 / obs)[::-1][:-1], np.diff((1 / obs)[::-1]) * 86400, "ko", lw=1.5, linestyle="-")
557        ax.plot(
558            (1.0 / theory)[::-1][:-1],
559            -np.diff(1.0 / theory)[::-1] * 86400,
560            "ko",
561            color="blue",
562            lw=1.5,
563            linestyle="--",
564            markersize=6,
565            markeredgewidth=0.0,
566        )
567        plt.show()
568
569    return theory_sequence
570
571
572################################################################################
573def chisq_longest_sequence(theory_periods, orders, obs_periods, obs_periods_errors):
574    """
575    Method to extract the theoretical pattern that best matches the observed one.
576    Match each observed mode period to its best matching theoretical counterpart,
577    and adopt the longest sequence of consecutive modes found this way.
578    In case of multiple mode series with the same length, a final pattern selection
579    is made based on the best (chi-square) match between theory and observations.
580
581    Parameters
582    ----------
583    theory_periods : numpy array, dtype=float
584        Theoretical periods and their radial orders.
585    orders : numpy array, dtype=int
586        Radial orders corresponding to the periods in theory_periods.
587    obs_periods: numpy array, dtype=float
588        Observational periods.
589    obs_periods_errors: numpy array, dtype=float
590        The errors on obs_periods.
591
592    Returns
593    ----------
594    series_chi2: float
595        chi2 value of the selected theoretical frequencies.
596    final_theoretical_periods: numpy array, dtype=float
597        The selected theoretical periods that best match the observed pattern.
598    corresponding_orders: list of integers
599        The radial orders of the returned theoretical periods.
600    """
601    if len(theory_periods) < len(obs_periods):
602        return 1e16, [-1.0 for i in range(len(obs_periods))], [-1 for i in range(len(obs_periods))]
603    else:
604        # Find the best matches per observed period
605        pairs_orders = []
606        for ii, period in enumerate(obs_periods):
607            ## Chi_squared array definition
608            chi_sqrs = np.array(
609                [((period - theory_period) / obs_periods_errors[ii]) ** 2 for theory_period in theory_periods]
610            )
611
612            ## Locate the theoretical frequency (and accompanying order) with the best chi2
613            min_ind = np.where(chi_sqrs == min(chi_sqrs))[0]
614            best_match = theory_periods[min_ind][0]
615            best_order = orders[min_ind][0]
616
617            ## Toss everything together for bookkeeping
618            pairs_orders.append([period, best_match, int(best_order), chi_sqrs[min_ind][0]])
619
620        pairs_orders = np.array(pairs_orders)
621
622        # If input is in increasing radial order (decreasing n_pg, since n_pg is negative for g-modes)
623        if orders[1] == orders[0] - 1:
624            increase_or_decrease = -1
625        # If input is in decreasing radial order
626        else:
627            increase_or_decrease = 1
628
629        sequences = []
630        ## Go through all pairs of obs and theoretical frequencies and
631        ## check if the next observed frequency has a corresponding theoretical frequency
632        ## with the consecutive radial order.
633        current = []
634        lp = len(pairs_orders[:-1])
635        for ii, sett in enumerate(pairs_orders[:-1]):
636            if abs(sett[2]) == abs(pairs_orders[ii + 1][2]) + increase_or_decrease:
637                current.append(sett)
638            # If not consecutive radial order, save the current sequence and start a new one.
639            else:
640                current.append(sett)
641                sequences.append(np.array(current).reshape(len(current), 4))
642                current = []
643            if ii == lp - 1:
644                current.append(sett)
645                sequences.append(np.array(current).reshape(len(current), 4))
646                current = []
647        len_list = np.array([len(x) for x in sequences])
648        longest = np.where(len_list == max(len_list))[0]
649
650        ## Test if there really is one longest sequence
651        if len(longest) == 1:
652            lseq = sequences[longest[0]]
653
654        ## if not, pick, of all the sequences with the same length, the best based on chi2
655        else:
656            scores = [np.sum(sequences[ii][:, -1]) / len(sequences[ii]) for ii in longest]
657            min_score = np.where(scores == min(scores))[0][0]
658            lseq = sequences[longest[min_score]]
659
660        obs_ordering_ind = np.where(obs_periods == lseq[:, 0][0])[0][0]
661        thr_ordering_ind = np.where(theory_periods == lseq[:, 1][0])[0][0]
662
663        ordered_theoretical_periods = []
664        corresponding_orders = []
665
666        thr_ind_start = thr_ordering_ind - obs_ordering_ind
667        thr_ind_current = thr_ind_start
668
669        for i, oper in enumerate(obs_periods):
670            thr_ind_current = thr_ind_start + i
671            if thr_ind_current < 0:
672                tper = -1
673                ordr = -1
674            elif thr_ind_current >= len(theory_periods):
675                tper = -1
676                ordr = -1
677            else:
678                tper = theory_periods[thr_ind_current]
679                ordr = orders[thr_ind_current]
680            ordered_theoretical_periods.append(tper)
681            corresponding_orders.append(ordr)
682
683        final_theoretical_periods = np.array(ordered_theoretical_periods)
684
685        obs_series, obs_series_errors = generate_spacing_series(obs_periods, obs_periods_errors)
686        thr_series, _ = generate_spacing_series(final_theoretical_periods)
687
688        obs_series = np.array(obs_series)
689        obs_series_errors = np.array(obs_series_errors)
690        thr_series = np.array(thr_series)
691
692        series_chi2 = np.sum(((obs_series - thr_series) / obs_series_errors) ** 2) / len(obs_series)
693
694        return series_chi2, final_theoretical_periods, corresponding_orders