Source code for gatspy.periodic.lomb_scargle

from __future__ import division, print_function, absolute_import

__all__ = ['LombScargle', 'LombScargleAstroML']

import warnings

import numpy as np

from .modeler import PeriodicModeler


class LeastSquaresMixin(object):
    """Mixin for matrix-based Least Squares periodic analysis"""
    def _construct_X(self, omega, weighted=True, **kwargs):
        raise NotImplementedError()

    def _construct_y(self, weighted=True, **kwargs):
        raise NotImplementedError()

    def _construct_X_M(self, omega, **kwargs):
        """Construct the weighted normal matrix of the problem"""
        X = self._construct_X(omega, weighted=True, **kwargs)
        M = np.dot(X.T, X)

        if hasattr(self, 'regularization') and self.regularization is not None:
            diag = M.ravel(order='K')[::M.shape[0] + 1]
            if self.regularize_by_trace:
                diag += diag.sum() * np.asarray(self.regularization)
            else:
                diag += np.asarray(self.regularization)

        return X, M

    def _compute_ymean(self, **kwargs):
        """Compute the (weighted) mean of the y data"""
        y = np.asarray(kwargs.get('y', self.y))
        dy = np.asarray(kwargs.get('dy', self.dy))

        if dy.size == 1:
            return np.mean(y)
        else:
            return np.average(y, weights=1 / dy ** 2)

    def _construct_y(self, weighted=True, **kwargs):
        y = kwargs.get('y', self.y)
        dy = kwargs.get('dy', self.dy)
        center_data = kwargs.get('center_data', self.center_data)

        y = np.asarray(y)
        dy = np.asarray(dy)

        if center_data:
            y = y - self._compute_ymean(y=y, dy=dy)

        if weighted:
            return y / dy
        else:
            return y

    def _best_params(self, omega):
        Xw, XTX = self._construct_X_M(omega)
        XTy = np.dot(Xw.T, self.yw_)
        return np.linalg.solve(XTX, XTy)

    def _score(self, periods):
        omegas = 2 * np.pi / periods

        # Set up the reference chi2. Note that this entire function would
        # be much simpler if we did not allow center_data=False.
        # We keep it just to make sure our math is correct
        chi2_0 = np.dot(self.yw_.T, self.yw_)
        if self.center_data:
            chi2_ref = chi2_0
        else:
            yref = self._construct_y(weighted=True, center_data=True)
            chi2_ref = np.dot(yref.T, yref)
        chi2_0_minus_chi2 = np.zeros(omegas.size, dtype=float)

        # Iterate through the omegas and compute the power for each
        for i, omega in enumerate(omegas.flat):
            Xw, XTX = self._construct_X_M(omega)
            XTy = np.dot(Xw.T, self.yw_)
            chi2_0_minus_chi2[i] = np.dot(XTy.T, np.linalg.solve(XTX, XTy))

        # construct and return the power from the chi2 difference
        if self.center_data:
            P = chi2_0_minus_chi2 / chi2_ref
        else:
            P = 1 + (chi2_0_minus_chi2 - chi2_0) / chi2_ref

        return P


[docs]class LombScargle(LeastSquaresMixin, PeriodicModeler): """Lomb-Scargle Periodogram Implementation This is a generalized periodogram implementation using the matrix formalism outlined in VanderPlas & Ivezic 2015. It allows computation of periodograms and best-fit models for both the classic normalized periodogram and truncated Fourier series generalizations. Parameters ---------- optimizer : PeriodicOptimizer instance Optimizer to use to find the best period. If not specified, the LinearScanOptimizer will be used. center_data : boolean (default = True) If True, then compute the weighted mean of the input data and subtract before fitting the model. fit_offset : boolean (default = True) If True, then fit a floating-mean sinusoid model. Nterms : int (default = 1) Number of Fourier frequencies to fit in the model regularization : float, vector or None (default = None) If specified, then add this regularization penalty to the least squares fit. regularize_by_trace : boolean (default = True) If True, multiply regularization by the trace of the matrix 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 Examples -------- >>> rng = np.random.RandomState(0) >>> t = 100 * rng.rand(100) >>> dy = 0.1 >>> omega = 10 >>> y = np.sin(omega * t) + dy * rng.randn(100) >>> ls = LombScargle().fit(t, y, dy) >>> ls.optimizer.period_range = (0.2, 1.2) >>> ls.best_period Finding optimal frequency: - Estimated peak width = 0.0639 - Using 5 steps per peak; omega_step = 0.0128 - User-specified period range: 0.2 to 1.2 - Computing periods at 2051 steps Zooming-in on 5 candidate peaks: - Computing periods at 1000 steps 0.62827068275990694 >>> ls.predict([0, 0.5]) array([-0.01445853, -0.92762251]) See Also -------- LombScargleAstroML LombScargleMultiband LombScargleMultibandFast """ def __init__(self, optimizer=None, center_data=True, fit_offset=True, Nterms=1, regularization=None, regularize_by_trace=True, fit_period=False, optimizer_kwds=None): self.center_data = center_data self.fit_offset = fit_offset self.Nterms = int(Nterms) self.regularization = regularization self.regularize_by_trace = regularize_by_trace PeriodicModeler.__init__(self, optimizer, fit_period=fit_period, optimizer_kwds=optimizer_kwds) if not self.center_data and not self.fit_offset: warnings.warn("Not centering data or fitting offset can lead " "to poor results") if self.Nterms < 0: raise ValueError("Nterms must be non-negative") if self.Nterms == 0 and not fit_offset: raise ValueError("Empty model: try larger Nterms.") def _construct_X(self, omega, weighted=True, **kwargs): """Construct the design matrix for the problem""" t = kwargs.get('t', self.t) dy = kwargs.get('dy', self.dy) fit_offset = kwargs.get('fit_offset', self.fit_offset) if fit_offset: offsets = [np.ones(len(t))] else: offsets = [] cols = sum(([np.sin((i + 1) * omega * t), np.cos((i + 1) * omega * t)] for i in range(self.Nterms)), offsets) if weighted: return np.transpose(np.vstack(cols) / dy) else: return np.transpose(np.vstack(cols)) def _fit(self, t, y, dy): self.yw_ = self._construct_y(weighted=True) self.ymean_ = self._compute_ymean() def _predict(self, t, period): omega = 2 * np.pi / period t = np.asarray(t) outshape = t.shape theta = self._best_params(omega) X = self._construct_X(omega, weighted=False, t=t.ravel()) return np.reshape(self.ymean_ + np.dot(X, theta), outshape) def _score(self, periods): return LeastSquaresMixin._score(self, periods)
[docs]class LombScargleAstroML(LombScargle): """Lomb-Scargle Periodogram Implementation using AstroML This is a generalized periodogram implementation which uses the periodogram functions from astroML. Compared to LombScargle, this implementation is both faster and more memory-efficient. Parameters ---------- optimizer : PeriodicOptimizer instance Optimizer to use to find the best period. If not specified, the LinearScanOptimizer will be used. Nterms : int (default = 1) Number of terms for the fit. Only Nterms=1 is currently supported. center_data : boolean (default = True) If True, then compute the weighted mean of the input data and subtract before fitting the model. fit_offset : boolean (default = True) If True, then fit a floating-mean sinusoid model. slow_version : boolean (default = False) If True, use the slower pure-python version from astroML. Otherwise, use the faster version of the code from astroML_addons 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 Examples -------- >>> rng = np.random.RandomState(0) >>> t = 100 * rng.rand(100) >>> dy = 0.1 >>> omega = 10 >>> y = np.sin(omega * t) + dy * rng.randn(100) >>> ls = LombScargleAstroML().fit(t, y, dy) >>> ls.optimizer.period_range = (0.2, 1.2) >>> ls.best_period Finding optimal frequency: - Estimated peak width = 0.0639 - Using 5 steps per peak; omega_step = 0.0128 - User-specified period range: 0.2 to 1.2 - Computing periods at 2051 steps Zooming-in on 5 candidate peaks: - Computing periods at 1000 steps 0.62827068275990694 See Also -------- LombScargle LombScargleMultiband LombScargleMultibandFast """ def __init__(self, optimizer=None, Nterms=1, fit_offset=True, center_data=True, slow_version=False, fit_period=False, optimizer_kwds=None): if Nterms != 1: raise ValueError("Only Nterms=1 is supported") LombScargle.__init__(self, optimizer=optimizer, Nterms=1, center_data=center_data, fit_offset=fit_offset, fit_period=fit_period, optimizer_kwds=optimizer_kwds) if slow_version: from astroML.time_series._periodogram import lomb_scargle else: from astroML.time_series import lomb_scargle self._LS_func = lomb_scargle def _score(self, periods): return self._LS_func(self.t, self.y, self.dy, 2 * np.pi / periods, generalized=self.fit_offset, subtract_mean=self.center_data)