add multiband lombscargle

astropy/timeseries/periodograms/__init__.py

@@ -1,3 +1,4 @@
 from astropy.timeseries.periodograms.base import *
 from astropy.timeseries.periodograms.bls import *
 from astropy.timeseries.periodograms.lombscargle import *
+from astropy.timeseries.periodograms.lombscargle_multiband import *

astropy/timeseries/periodograms/lombscargle_multiband/__init__.py

@@ -0,0 +1,3 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+from .core import LombScargleMultiband

astropy/timeseries/periodograms/lombscargle_multiband/implementations/__init__.py

@@ -0,0 +1,5 @@
+"""Various implementations of the Multiband Lomb-Scargle Periodogram"""
+
+from .main import available_methods, lombscargle_multiband
+from .mbfast_impl import lombscargle_mbfast
+from .mbflex_impl import lombscargle_mbflex

astropy/timeseries/periodograms/lombscargle_multiband/implementations/main.py

@@ -0,0 +1,73 @@
+"""
+Main Lomb-Scargle Multiband Implementation
+
+The ``lombscarglemultiband`` function here is essentially a switch
+statement for the various implementations available in this submodule
+"""
+
+__all__ = ["lombscargle_multiband", "available_methods"]
+
+from .mbfast_impl import lombscargle_mbfast
+from .mbflex_impl import lombscargle_mbflex
+
+
+def available_methods():
+    methods = ["fast", "flexible"]
+    return methods
+
+
+def lombscargle_multiband(
+    t,
+    y,
+    bands,
+    dy=None,
+    frequency=None,
+    method="flexible",
+    sb_method="auto",
+    assume_regular_frequency=False,
+    normalization="standard",
+    fit_mean=True,
+    center_data=True,
+    method_kwds=None,
+    nterms_base=1,
+    nterms_band=1,
+    reg_base=None,
+    reg_band=1e-6,
+    regularize_by_trace=True,
+    fit_period=False,
+):
+    methods = available_methods()
+
+    if method == "flexible":
+        power = lombscargle_mbflex(
+            t,
+            y,
+            bands,
+            frequency,
+            dy=dy,
+            nterms_base=nterms_base,
+            nterms_band=nterms_band,
+            reg_base=reg_base,
+            reg_band=reg_band,
+            regularize_by_trace=regularize_by_trace,
+            center_data=center_data,
+        )
+    elif method == "fast":
+        power = lombscargle_mbfast(
+            t,
+            y,
+            bands,
+            dy,
+            frequency=frequency,
+            sb_method=sb_method,
+            assume_regular_frequency=assume_regular_frequency,
+            normalization=normalization,
+            fit_mean=fit_mean,
+            center_data=center_data,
+            method_kwds=method_kwds,
+            nterms=nterms_base,
+        )  # nterms_base used for single_band nterms
+    if method not in methods:
+        raise ValueError(f"Invalid Method: {method}")
+
+    return power

astropy/timeseries/periodograms/lombscargle_multiband/implementations/mbfast_impl.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+from astropy.timeseries.periodograms.lombscargle.implementations import lombscargle
+
+
+def lombscargle_mbfast(
+    t,
+    y,
+    bands,
+    dy=None,
+    frequency=None,
+    sb_method="auto",
+    assume_regular_frequency=False,
+    normalization="standard",
+    fit_mean=True,
+    center_data=True,
+    method_kwds=None,
+    nterms=1,
+):
+    # create masks for each filter/bandpass
+    unique_bands = np.unique(bands)
+    masks = [(bands == band) for band in unique_bands]
+
+    # calculate singleband powers for each filter/bandpass
+    if dy is not None:
+        models = [
+            lombscargle(
+                t[mask],
+                y[mask],
+                dy=dy[mask],
+                frequency=frequency,
+                method=sb_method,
+                assume_regular_frequency=assume_regular_frequency,
+                normalization=normalization,
+                fit_mean=fit_mean,
+                center_data=center_data,
+                method_kwds=method_kwds,
+                nterms=nterms,
+            )
+            for mask in masks
+        ]
+    else:
+        models = [
+            lombscargle(
+                t[mask],
+                y[mask],
+                dy=None,
+                frequency=frequency,
+                method=sb_method,
+                assume_regular_frequency=assume_regular_frequency,
+                normalization=normalization,
+                fit_mean=fit_mean,
+                center_data=center_data,
+                method_kwds=method_kwds,
+                nterms=nterms,
+            )
+            for mask in masks
+        ]
+
+    # Total score is the sum of powers weighted by chi2-normalization
+    powers = np.array(models)
+    chi2_0 = np.array([np.sum(model**2) for model in models])
+    return np.dot(chi2_0 / chi2_0.sum(), powers)

astropy/timeseries/periodograms/lombscargle_multiband/implementations/mbflex_impl.py

@@ -0,0 +1,125 @@
+import numpy as np
+
+
+def lombscargle_mbflex(
+    t,
+    y,
+    bands,
+    frequency,
+    dy=None,
+    nterms_base=1,
+    nterms_band=1,
+    reg_base=None,
+    reg_band=1e-6,
+    regularize_by_trace=True,
+    center_data=True,
+):
+    if (nterms_base == 0) and (nterms_band == 0):
+        raise ValueError(
+            "At least one of nterms_base and nterms_band must be greater than 0."
+        )
+    # Inputs
+    unique_bands = np.unique(bands)
+    t = np.asarray(t)
+    y = np.asarray(y)
+    bands = np.asarray(bands)
+    frequency = np.asarray(frequency)
+
+    # Create a ones array for dy (errors) if not provided
+    if dy is not None:
+        dy = np.asarray(dy)
+    else:
+        dy = np.ones(y.shape)
+
+    # Calculate ymeans -- per band/filter
+    ymeans = np.zeros(
+        y.shape
+    )  # An array of shape y, with each index given a filter specific mean
+    for band in unique_bands:
+        mask = bands == band
+        ymeans[mask] = np.average(y[mask], weights=1 / dy[mask] ** 2)
+    # Construct weighted y matrix
+    if center_data:
+        y = y - ymeans
+
+    yw = y / dy  # weighted by dy, as above; one's array if not provided
+
+    # Construct Regularization
+    if reg_base is None and reg_band is None:
+        regularization = 0
+    else:
+        n_base = 1 + 2 * nterms_base
+        n_band = 1 + 2 * nterms_band
+        regularization = np.zeros(n_base + len(unique_bands) * n_band)
+        if reg_base is not None:
+            regularization[:n_base] = reg_base
+        if reg_band is not None:
+            regularization[n_base:] = reg_band
+
+    # Scoring
+    omegas = 2 * np.pi * frequency
+
+    # Calculate chi-squared
+    chi2_0 = np.dot(yw.T, yw)
+    chi2_ref = np.copy(chi2_0)  # reference chi2 for later comparison
+
+    # Iterate through the omegas and compute the power for each
+    chi2_0_minus_chi2 = []
+    for i, omega in enumerate(omegas.flat):
+        # Construct X - design matrix of the stationary sinusoid model
+        cols = [np.ones(len(t))]
+        cols = sum(
+            (
+                [np.sin((i + 1) * omega * t), np.cos((i + 1) * omega * t)]
+                for i in range(nterms_base)
+            ),
+            cols,
+        )
+
+        # Add band columns for the multiband model, binary flag indicating the band of a given observation
+        for band in unique_bands:
+            cols.append(np.ones(len(t)))
+            cols = sum(
+                (
+                    [np.sin((i + 1) * omega * t), np.cos((i + 1) * omega * t)]
+                    for i in range(nterms_band)
+                ),
+                cols,
+            )
+            mask = bands == band
+            for i in range(-1 - 2 * nterms_band, 0):
+                cols[i][~mask] = 0
+
+        X = np.transpose(np.vstack(cols) / dy)  # weighted
+        M = np.dot(X.T, X)
+
+        if regularization is not None:
+            diag = M.ravel(order="K")[
+                :: M.shape[0] + 1
+            ]  # M is being affected by operations on diag
+            if regularize_by_trace:
+                diag += diag.sum() * np.asarray(regularization)
+            else:
+                diag += np.asarray(regularization)
+
+        # Construct Xw, XTX, XTy
+        Xw = X
+        XTX = M
+        XTy = np.dot(Xw.T, yw)
+
+        # Matrix Algebra to calculate the Lomb-Scargle power at each omega step
+        try:
+            chi2_0_minus_chi2.append(np.dot(XTy.T, np.linalg.solve(XTX, XTy)))
+        # If X'X is not invertible, use pseudoinverse instead
+        except np.linalg.LinAlgError:
+            chi2_0_minus_chi2.append(
+                np.dot(XTy.T, np.linalg.lstsq(XTX, XTy, rcond=None)[0])
+            )
+
+    # Construct and return the power from the chi2 difference
+    if center_data:
+        P = chi2_0_minus_chi2 / chi2_ref
+    else:
+        P = 1 + (chi2_0_minus_chi2 - chi2_0) / chi2_ref
+
+    return P