diff --git a/gatspy/periodic/__init__.py b/gatspy/periodic/__init__.py index 7d20402..454c612 100644 --- a/gatspy/periodic/__init__.py +++ b/gatspy/periodic/__init__.py @@ -6,7 +6,9 @@ __all__ = ['LombScargle', 'LombScargleFast', 'LombScargleAstroML', 'LombScargleMultiband', 'LombScargleMultibandFast', - 'TrendedLombScargle', 'SuperSmoother', 'SuperSmootherMultiband', + 'TrendedLombScargle', 'TrendedLombScargleMultiband', + 'PolyTrend_LombScargleMultiband', + 'SuperSmoother', 'SuperSmootherMultiband', 'RRLyraeTemplateModeler', 'RRLyraeTemplateModelerMultiband', 'NaiveMultiband'] @@ -14,6 +16,7 @@ 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 * diff --git a/gatspy/periodic/tests/test_trended_lomb_scargle_multiband.py b/gatspy/periodic/tests/test_trended_lomb_scargle_multiband.py new file mode 100755 index 0000000..5a7242b --- /dev/null +++ b/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) diff --git a/gatspy/periodic/trended_lomb_scargle_multiband.py b/gatspy/periodic/trended_lomb_scargle_multiband.py new file mode 100644 index 0000000..f732c8d --- /dev/null +++ b/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))