Source code for gatspy.periodic.lomb_scargle_multiband

"""
Multiband generalizations of Lomb-Scargle Periodograms
"""

from __future__ import division, print_function

__all__ = ['LombScargleMultiband', 'LombScargleMultibandFast']

import numpy as np

from .modeler import PeriodicModelerMultiband
from .lomb_scargle import LombScargle, LeastSquaresMixin
from .lomb_scargle_fast import LombScargleFast


[docs]class LombScargleMultiband(LeastSquaresMixin, PeriodicModelerMultiband): """Multiband Periodogram Implementation This implements the generalized multi-band periodogram described in VanderPlas & Ivezic 2015. 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 See Also -------- LombScargle LombScargleFast LombScargleMultibandFast """ def __init__(self, optimizer=None, 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.Nterms_base = Nterms_base self.Nterms_band = Nterms_band self.reg_base = reg_base self.reg_band = reg_band self.regularize_by_trace = regularize_by_trace self.center_data = center_data PeriodicModelerMultiband.__init__(self, optimizer, fit_period=fit_period, optimizer_kwds=optimizer_kwds) def _fit(self, t, y, dy, filts): self.ymean_ = self._compute_ymean() masks = [(filts == filt) for filt in self.unique_filts_] self.ymean_by_filt_ = [LeastSquaresMixin._compute_ymean(self, y=y[mask], dy=dy[mask]) for mask in masks] self.yw_ = self._construct_y(weighted=True) self.regularization = self._construct_regularization() return self def _compute_ymean(self, **kwargs): y = kwargs.get('y', self.y) dy = kwargs.get('dy', self.dy) filts = kwargs.get('filts', self.filts) ymean = np.zeros(y.shape) for filt in np.unique(filts): mask = (filts == filt) ymean[mask] = LeastSquaresMixin._compute_ymean(self, y=y[mask], dy=dy[mask]) return ymean def _construct_regularization(self): if self.reg_base is None and self.reg_band is None: reg = 0 else: Nbase = 1 + 2 * self.Nterms_base Nband = 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), (1 + 2 * Nterms_base + nfilts * (1 + 2 * Nterms_band))] # TODO: the building of the matrix could be more efficient cols = [np.ones(len(t))] 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 = 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 - 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)) def _predict(self, t, filts, period): t, filts = np.broadcast_arrays(t, filts) output_shape = t.shape t = t.ravel() filts = filts.ravel() omega = 2 * np.pi / period # TODO: broadcast this ymeans = np.zeros(len(filts)) for i, filt in enumerate(filts): j = np.where(self.unique_filts_ == filt)[0][0] ymeans[i] = self.ymean_by_filt_[j] theta = self._best_params(omega) X = self._construct_X(omega, weighted=False, t=t, filts=filts) return (ymeans + np.dot(X, theta)).reshape(output_shape)
[docs]class LombScargleMultibandFast(PeriodicModelerMultiband): """Fast Multiband Periodogram Implementation This implements the specialized multi-band periodogram described in VanderPlas & Ivezic 2015 (with Nbase=0 and Nband=1) which is essentially a weighted sum of the standard periodograms for each band. Parameters ---------- optimizer : PeriodicOptimizer instance Optimizer to use to find the best period. If not specified, the LinearScanOptimizer will be used. Nterms : integer (default = 1) Number of fourier terms to use in the model BaseModel : PeriodicModeler class (optional) The base model to use for each individual band. By default it will use :class:`LombScargleFast` if Nterms == 1, and :class:`LombScargle` otherwise. fit_period : bool (optional) If True, then fit for the best period when fit() method is called. optimizer_kwds : dict (optional Dictionary of keyword arguments for constructing the optimizer See Also -------- LombScargle LombScargleFast LombScargleMultiband """ def __init__(self, optimizer=None, Nterms=1, BaseModel=None, fit_period=False, optimizer_kwds=None): # Note: center_data must be True, or else the chi^2 weighting will fail self.Nterms = Nterms if BaseModel is None: if Nterms == 1: BaseModel = LombScargleFast else: BaseModel = LombScargle self.BaseModel = BaseModel PeriodicModelerMultiband.__init__(self, optimizer, fit_period=fit_period, optimizer_kwds=optimizer_kwds) def _fit(self, t, y, dy, filts): masks = [(filts == f) for f in self.unique_filts_] self.models_ = [self.BaseModel(Nterms=self.Nterms, center_data=True, fit_offset=True).fit(t[m], y[m], dy[m]) for m in masks] def _score(self, periods): # Total score is the sum of powers weighted by chi2-normalization powers = np.array([model.score(periods) for model in self.models_]) chi2_0 = np.array([np.sum(model.yw_ ** 2) for model in self.models_]) return np.dot(chi2_0 / chi2_0.sum(), powers) def _score_frequency_grid(self, f0, df, N): powers = np.array([model._score_frequency_grid(f0, df, N) for model in self.models_]) chi2_0 = np.array([np.sum(model.yw_ ** 2) for model in self.models_]) return np.dot(chi2_0 / chi2_0.sum(), powers) def _best_params(self, omega): return np.asarray([model._best_params(omega) for model in self.models_]) def _predict(self, t, filts, period): t, filts = np.broadcast_arrays(t, filts) result = np.zeros(t.shape, dtype=float) masks = ((filts == f) for f in self.unique_filts_) for model, mask in zip(self.models_, masks): result[mask] = model.predict(t[mask], period=period) return result