Extension to Trended Lomb-Scargle for Multiple Bands

gatspy/periodic/__init__.py

@@ -6,14 +6,17 @@
 
 __all__ = ['LombScargle', 'LombScargleFast', 'LombScargleAstroML',
            'LombScargleMultiband', 'LombScargleMultibandFast',
-           'TrendedLombScargle', 'SuperSmoother', 'SuperSmootherMultiband',
+           'TrendedLombScargle', 'TrendedLombScargleMultiband',
+           'PolyTrend_LombScargleMultiband',
+           'SuperSmoother', 'SuperSmootherMultiband',
            'RRLyraeTemplateModeler', 'RRLyraeTemplateModelerMultiband',
            'NaiveMultiband']
 
 from .lomb_scargle import *
 from .lomb_scargle_fast import *
 from .lomb_scargle_multiband import *
 from .trended_lomb_scargle import *
+from .trended_lomb_scargle_multiband import *
 from .supersmoother import *
 from .template_modeler import *
 from .naive_multiband import *

gatspy/periodic/tests/test_trended_lomb_scargle_multiband.py

@@ -0,0 +1,41 @@
+from __future__ import division
+
+import numpy as np
+from numpy.testing import assert_allclose, assert_, assert_raises
+from nose import SkipTest
+
+from .. import (LombScargle, TrendedLombScargleMultiband)
+
+
+def _generate_data(N=100, period=1, theta=[10, 2, 3], dy=1, rseed=0):
+    """Generate some data for testing"""
+    rng = np.random.RandomState(rseed)
+    t = 20 * period * rng.rand(N)
+    omega = 2 * np.pi / period
+    y = theta[0] + theta[1] * np.sin(omega * t) + theta[2] * np.cos(omega * t)
+    dy = dy * (0.5 + rng.rand(N))
+    y += dy * rng.randn(N)
+    return t, y, dy
+
+
+def test_multiband_lomb_scargle_linear_trend(N=100, period=1, slope=5):
+    """Test whether the generalized lomb-scargle properly fits data with a
+    linear trend component (single filter)
+    """
+    t, y, dy = _generate_data(N, period)
+    omega = 2 * np.pi / period
+
+    # Model and optimizer for standard Lomb-Scargle
+    model = LombScargle(center_data=False, fit_offset=True)
+    model.fit(t, y, dy)
+    model.optimizer.period_range = (period - 0.5, period + 0.5)
+    y_hat = model.predict(t)
+
+    # Model and optimizer for trended multiband Lomb-Scargle
+    y_mb_t = y + slope * t
+    model_mb_t = TrendedLombScargleMultiband(center_data=False)
+    model_mb_t.fit(t, y, dy, filts=1)
+    model_mb_t.optimizer.period_range = (period - 0.5, period + 0.5)
+    y_hat_mb_t = model_mb_t.predict(t)
+
+    assert_allclose(y_hat, y_hat_mb_t - slope * t, rtol=5E-2)

gatspy/periodic/trended_lomb_scargle_multiband.py

@@ -0,0 +1,206 @@
+from __future__ import division, print_function, absolute_import
+
+__all__ = ['TrendedLombScargleMultiband', 'PolyTrend_LombScargleMultiband']
+
+import warnings
+
+import numpy as np
+
+from .modeler import PeriodicModeler
+from .lomb_scargle import LombScargle
+from .lomb_scargle_multiband import LombScargleMultiband
+
+
+class TrendedLombScargleMultiband(LombScargleMultiband):
+    """Trended Multiband Lomb-Scargle Periodogram Implementation
+
+    This is a generalized multiband periodogram implementation using the matrix
+    formalism outlined in VanderPlas & Ivezic 2015. It fits both a floating mean
+    and a trend parameter in each band.
+
+    Parameters
+    ----------
+    optimizer : PeriodicOptimizer instance
+        Optimizer to use to find the best period. If not specified, the
+        LinearScanOptimizer will be used.
+    Nterms_base : integer (default = 1)
+        number of frequency terms to use for the base model common to all bands
+    Nterms_band : integer (default = 1)
+        number of frequency terms to use for the residuals between the base
+        model and each individual band
+    reg_base : float or None (default = None)
+        amount of regularization to use on the base model parameters
+    reg_band : float or None (default = 1E-6)
+        amount of regularization to use on the band model parameters
+    regularize_by_trace : bool (default = True)
+        if True, then regularization is expressed in units of the trace of
+        the normal matrix
+    center_data : boolean (default = True)
+        if True, then center the y data prior to the fit
+    optimizer_kwds : dict (optional)
+        Dictionary of keyword arguments for constructing the optimizer. For
+        example, silence optimizer output with `optimizer_kwds={"quiet": True}`.
+
+    See Also
+    --------
+    LombScargle
+    LombScargleMultiband
+    TrendedLombScargle
+    """
+
+    def _construct_regularization(self):
+        if self.reg_base is None and self.reg_band is None:
+            reg = 0
+        else:
+            Nbase = 2 + 2 * self.Nterms_base
+            Nband = 2 + 2 * self.Nterms_band
+            reg = np.zeros(Nbase + len(self.unique_filts_) * Nband)
+            if self.reg_base is not None:
+                reg[:Nbase] = self.reg_base
+            if self.reg_band is not None:
+                reg[Nbase:] = self.reg_band
+        return reg
+
+    def _construct_X(self, omega, weighted=True, **kwargs):
+        t = kwargs.get('t', self.t)
+        dy = kwargs.get('dy', self.dy)
+        filts = kwargs.get('filts', self.filts)
+
+        # X is a huge-ass matrix that has lots of zeros depending on
+        # which filters are present...
+        #
+        # huge-ass, quantitatively speaking, is of shape
+        #  [len(t), (2 + 2 * Nterms_base + nfilts * (2 + 2 * Nterms_band))]
+
+        # TODO: the building of the matrix could be more efficient
+        cols = [np.ones(len(t))]
+        cols.append(np.copy(t))   #  coefficients for trend parameter       
+        cols = sum(([np.sin((i + 1) * omega * t),
+                     np.cos((i + 1) * omega * t)]
+                    for i in range(self.Nterms_base)), cols)
+
+        for filt in self.unique_filts_:
+            cols.append(np.ones(len(t)))
+            cols.append(np.copy(t))   #  coefficients for trend parameter (in filt)
+            cols = sum(([np.sin((i + 1) * omega * t),
+                         np.cos((i + 1) * omega * t)]
+                        for i in range(self.Nterms_band)), cols)
+            mask = (filts == filt)
+            for i in range(-2 - (2 * self.Nterms_band), 0):
+                cols[i][~mask] = 0
+
+        if weighted:
+            return np.transpose(np.vstack(cols) / dy)
+        else:
+            return np.transpose(np.vstack(cols))
+
+
+class PolyTrend_LombScargleMultiband(LombScargleMultiband):
+    """Trended Multiband Lomb-Scargle Periodogram Implementation
+
+    This is a generalized multiband periodogram implementation using the matrix
+    formalism outlined in VanderPlas & Ivezic 2015. It fits both a floating mean
+    and a trend parameter in each band.
+
+    Parameters
+    ----------
+    optimizer : PeriodicOptimizer instance
+        Optimizer to use to find the best period. If not specified, the
+        LinearScanOptimizer will be used.
+    Nterms_base : integer (default = 1)
+        number of frequency terms to use for the base model common to all bands
+    Nterms_band : integer (default = 1)
+        number of frequency terms to use for the residuals between the base
+        model and each individual band
+    reg_base : float or None (default = None)
+        amount of regularization to use on the base model parameters
+    reg_band : float or None (default = 1E-6)
+        amount of regularization to use on the band model parameters
+    regularize_by_trace : bool (default = True)
+        if True, then regularization is expressed in units of the trace of
+        the normal matrix
+    center_data : boolean (default = True)
+        if True, then center the y data prior to the fit
+    optimizer_kwds : dict (optional)
+        Dictionary of keyword arguments for constructing the optimizer. For
+        example, silence optimizer output with `optimizer_kwds={"quiet": True}`.
+
+    See Also
+    --------
+    LombScargle
+    LombScargleMultiband
+    TrendedLombScargle
+    """
+
+    def __init__(self, optimizer=None,
+                 poly_trend_order_base=2, poly_trend_order_band=1,
+                 Nterms_base=1, Nterms_band=1,
+                 reg_base=None, reg_band=1E-6, regularize_by_trace=True,
+                 center_data=True, fit_period=False, optimizer_kwds=None):
+        self.poly_trend_order_base = poly_trend_order_base
+        self.poly_trend_order_band = poly_trend_order_band
+
+        LombScargleMultiband.__init__(
+            self, optimizer=optimizer,
+            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, fit_period=fit_period,
+            optimizer_kwds=optimizer_kwds,
+        )
+
+    def _construct_regularization(self):
+        if self.reg_base is None and self.reg_band is None:
+            reg = 0
+        else:
+            Nbase = (self.poly_trend_order_base + 1) + 2 * self.Nterms_base
+            Nband = (self.poly_trend_order_band + 1) + 2 * self.Nterms_band
+            reg = np.zeros(Nbase + len(self.unique_filts_) * Nband)
+            if self.reg_base is not None:
+                reg[:Nbase] = self.reg_base
+            if self.reg_band is not None:
+                reg[Nbase:] = self.reg_band
+        return reg
+
+    def _construct_X(self, omega, weighted=True, **kwargs):
+        t = kwargs.get('t', self.t)
+        dy = kwargs.get('dy', self.dy)
+        filts = kwargs.get('filts', self.filts)
+
+        # X is a huge-ass matrix that has lots of zeros depending on
+        # which filters are present...
+        #
+        # huge-ass, quantitatively speaking, is of shape
+        #  [len(t), ((poly_trend_order_base + 1) + 2 * Nterms_base +
+        #            nfilts * ((poly_trend_order_band + 1) + 2 * Nterms_band))]
+
+        # TODO: the building of the matrix could be more efficient
+        cols = [np.ones(len(t))]    # Coefficients for constant param
+
+        for cur_poly_trend_order in range(1, self.poly_trend_order_base + 1):
+            # Coefficients for current polynomial trend order
+            cols.append(np.copy(t**cur_poly_trend_order))
+
+        cols = sum(([np.sin((i + 1) * omega * t),
+                     np.cos((i + 1) * omega * t)]
+                    for i in range(self.Nterms_base)), cols)
+
+        for filt in self.unique_filts_:
+            cols.append(np.ones(len(t)))
+
+            for cur_poly_trend_order in range(1, self.poly_trend_order_band + 1):
+                # Coefficients for current polynomial trend order
+                cols.append(np.copy(t**cur_poly_trend_order))
+
+            cols = sum(([np.sin((i + 1) * omega * t),
+                         np.cos((i + 1) * omega * t)]
+                        for i in range(self.Nterms_band)), cols)
+            mask = (filts == filt)
+            for i in range((-1 * (self.poly_trend_order_band + 1)) -
+                           (2 * self.Nterms_band), 0):
+                cols[i][~mask] = 0
+
+        if weighted:
+            return np.transpose(np.vstack(cols) / dy)
+        else:
+            return np.transpose(np.vstack(cols))