asteroid_spinprops 0.2.32__py3-none-any.whl → 1.0.1__py3-none-any.whl

asteroid_spinprops/ssolib/periodest.py

@@ -1,8 +1,5 @@
 import numpy as np
 import pandas as pd
-import rocks
-import matplotlib.pyplot as plt
-
 
 from astropy.timeseries import LombScargleMultiband, LombScargle
 from scipy.signal import find_peaks
@@ -11,21 +8,59 @@ from scipy.stats import f as ftest
 
 
 def alias_func(x, i, j, p_feat):
-    """Return aliases relation
-
-    x: float
-        Fink period
-    i: int
-        Strictly positive integer. Mode.
-    j: int
-        Alias
-    p_feat: float
-        Frequency of a feature
+    """
+    Compute the aliasing relation for a periodic signal.
+
+    Parameters
+    ----------
+    x : float
+        Input period.
+    i : int
+        Positive integer representing the mode number.
+    j : int
+        Alias index.
+    p_feat : float
+        Characteristic frequency of the feature (e.g. 1 day).
+
+    Returns
+    -------
+    float
+        Value of the aliasing relation for the given inputs.
     """
     return (i + 1) * p_feat * x / np.abs(p_feat - j * x)
 
 
 def get_period_estimate(residuals_dataframe, p_min=0.03, p_max=2):
+    """
+    Estimate significant periods in a time series of residuals using Lomb-Scargle periodograms.
+
+    Parameters
+    ----------
+    residuals_dataframe : pandas.DataFrame
+        DataFrame containing at least the following columns:
+        - 'jd' : observation times
+        - 'residuals' : residuals (observed - SHG1G2 modeled magnitudes)
+        - 'filters' : filter IDs (int)
+    p_min : float, optional
+        Minimum period to search (in days). Default is 0.03.
+    p_max : float, optional
+        Maximum period to search (in days). Default is 2.
+
+    Returns
+    -------
+    tuple
+        - single_band_results : list
+            [frequencies, power, top signal peak frequencies, corresponding powers]
+        - window_function_results : list
+            [frequencies, power, top window peak frequencies, corresponding powers]
+        - noise_level : float
+            Estimated noise level in the periodogram (mean + 3*std of power)
+
+    Notes
+    -----
+    - Uses `LombScargle` for the window function and `LombScargleMultiband` from the `nifty-ls` implementations
+    - The five strongest peaks are returned.
+    """
     period_min, period_max = p_min, p_max
     period_range = (period_min, period_max)
 
@@ -90,6 +125,50 @@ def get_period_estimate(residuals_dataframe, p_min=0.03, p_max=2):
 def get_multiband_period_estimate(
     residuals_dataframe, p_min=0.03, p_max=2, k_free=True, k_val=None
 ):
+    """
+    Estimate the period of a multiband time series using a multiband Lomb-Scargle model.
+
+    Fits a multiband Lomb-Scargle model to the SHG1G2 residuals across different filters,
+    optionally testing multiple base-term complexities (`k`) and selecting the simplest model
+    that adequately describes the data.
+
+    Parameters
+    ----------
+    residuals_dataframe : pandas.DataFrame
+        DataFrame containing:
+        - 'jd' : observation times
+        - 'residuals' : residual magnitudes (observed - model)
+        - 'filters' : filter IDs
+        - 'sigma' : observational uncertainties
+    p_min : float, optional
+        Minimum period to search (days). Default is 0.03.
+    p_max : float, optional
+        Maximum period to search (days). Default is 2.
+    k_free : bool, optional
+        If True, automatically scan multiple base-term complexities to choose optimal model. Default True.
+    k_val : int, optional
+        Fixed number of base terms to use if `k_free=False`. Default None.
+
+    Returns
+    -------
+    tuple
+        - period_in : float
+            Estimated dominant SSO period (2 / f_best) in the time series.
+        - k_val : int
+            Number of base terms used in the final multiband model.
+        - p_rms : float
+            RMS of residuals for the chosen model (NaN if k_free=False).
+        - signal_peaks : array_like
+            Frequencies of the top five peaks in the final multiband periodogram.
+        - window_peaks : array_like
+            Frequencies of the top five peaks in the single-band window function periodogram.
+
+    Notes
+    -----
+    - Uses `LombScargle` for the window function and `LombScargleMultiband` from the `nifty-ls` implementations
+    - If `k_free=True`, performs F-test comparisons to select the simplest adequate model.
+    """
+
     period_min, period_max = p_min, p_max
     period_range = (period_min, period_max)
     results = []
@@ -245,152 +324,41 @@ def get_multiband_period_estimate(
     return period_in, k_val, p_rms, signal_peaks, window_peaks
 
 
-def plot_periodograms(signal, window, name=None, axis="frequency"):
-    if name is not None:
-        r = rocks.Rock(name, datacloud="spins")
-        bib_estimates = r.spins["period"].values[pd.notna(r.spins["period"].values)]
-
-    fig, ax = plt.subplots(2, 1, figsize=(12, 6))
-    if axis == "frequency":
-        ax[0].plot(signal[0], signal[1], c="red", linewidth=2)
-        ax[1].plot(window[0], window[1], c="black", alpha=1, linewidth=2)
-
-        for i, n in enumerate(signal[2]):
-            ax[0].scatter(
-                n - 0.2,
-                signal[3][i],
-                marker="${}$".format(i + 1),
-                zorder=1000,
-                s=60,
-                c="black",
-            )
-
-        for i in [0, 1]:
-            for j in range(1, 9):
-                ax[i].axvline(x=24 / (24 / j), linestyle="-.", c="purple", linewidth=1)
-            ax[i].set_xlim(
-                1 / (2 + 0.5),
-            )
-            ax[i].set_ylim(
-                0,
-            )
-        if name is not None:
-            for p in bib_estimates:
-                ax[0].axvline(x=24 / (p / 2), linestyle="-", c="green", linewidth=1.5)
-
-        for k in [
-            int(24 * 2),
-            int(12 * 2),
-            int(8 * 2),
-            int(6 * 2),
-        ]:
-            ax[0].text(
-                x=24 / (k / 2) - 0.1,
-                y=max(signal[1]) + 0.05,
-                s="{}h".format(k),
-                rotation=90,
-                fontsize=11,
-            )
-
-        ax[0].text(
-            x=0.95,
-            y=0.95,
-            s="Signal",
-            fontsize=15,
-            transform=ax[0].transAxes,
-            ha="right",
-            va="top",
-        )
-        ax[1].text(
-            x=0.95,
-            y=0.95,
-            s="Window",
-            fontsize=15,
-            transform=ax[1].transAxes,
-            ha="right",
-            va="top",
-        )
-
-        ax[1].set_xlabel(r"Frequency /day$^{-1}$")
-
-        ax[0].set_ylabel("Power")
-        ax[1].set_ylabel("Power")
-
-        ax[0].set_xticks([])
-
-        plt.tight_layout()
-
-    if axis == "period":
-        ax[0].plot(2 / signal[0], signal[1], c="red", linewidth=2)
-        ax[1].plot(2 / window[0], window[1], c="black", alpha=1, linewidth=2)
-
-        for i, n in enumerate(2 / signal[2]):
-            ax[0].scatter(
-                2 / (n - 0.2),
-                signal[3][i],
-                marker="${}$".format(i + 1),
-                zorder=1000,
-                s=60,
-                c="black",
-            )
-
-        # for i in [0, 1]:
-        #     for j in range(1, 9):
-        #         ax[i].axvline(x=24 / (24 / j), linestyle="-.", c="purple", linewidth=1)
-        #     ax[i].set_xlim(
-        #         1 / (2 + 0.5),
-        #     )
-        #     ax[i].set_ylim(
-        #         0,
-        #     )
-        if name is not None:
-            for p in bib_estimates:
-                ax[0].axvline(x=p / 24, linestyle="-", c="green", linewidth=1.5)
-
-        for k in [
-            int(24 * 2),
-            int(12 * 2),
-            int(8 * 2),
-            int(6 * 2),
-        ]:
-            ax[0].text(
-                x=(k - 0.1) / 24,
-                y=max(signal[1]) + 0.05,
-                s="{}h".format(k),
-                rotation=90,
-                fontsize=11,
-            )
-
-        ax[0].text(
-            x=0.95,
-            y=0.95,
-            s="Signal",
-            fontsize=15,
-            transform=ax[0].transAxes,
-            ha="right",
-            va="top",
-        )
-        ax[1].text(
-            x=0.95,
-            y=0.95,
-            s="Window",
-            fontsize=15,
-            transform=ax[1].transAxes,
-            ha="right",
-            va="top",
-        )
-
-        ax[1].set_xlabel(r"Period /day")
-
-        ax[0].set_ylabel("Power")
-        ax[1].set_ylabel("Power")
-
-        ax[0].set_xticks([])
-
-        plt.tight_layout()
-
-
 def perform_residual_resampling(resid_df, p_min, p_max, k=1):
+    """
+    Estimate the robustness of a period measurement via bootstrap resampling of residuals.
+
+    This function resamples the residuals DataFrame multiple times with replacement,
+    recomputes the period for each bootstrap sample, and counts how many of the
+    resampled periods are within 1% of the original period. It supports both single-band
+    (k=1) and multiband (k>1) period estimation.
+
+    Parameters
+    ----------
+    resid_df : pandas.DataFrame
+        DataFrame of residuals containing columns required for `get_period_estimate`
+        or `get_multiband_period_estimate`.
+    p_min : float
+        Minimum period to search (days).
+    p_max : float
+        Maximum period to search (days).
+    k : int, optional
+        Number of base terms in the model; k=1 for single-band, k>1 for multiband. Default is 1.
+
+    Returns
+    -------
+    tuple
+        - BS_df : pandas.DataFrame
+            The last bootstrap-resampled residuals DataFrame.
+        - Nbs : int
+            Number of bootstrap samples whose estimated period is within 1% of the original period.
+
+    Notes
+    -----
+    - Performs 25 bootstrap resamples by default.
+    - For k=1, uses `get_period_estimate`; for k>1, uses `get_multiband_period_estimate`.
+    """
+
     if k == 1:
         sg, w, _ = get_period_estimate(resid_df)
         Pog = 48 / sg[2][0]  # in hours