[WIP]: Lightcurve Periodogram

pyke/__init__.py

@@ -50,4 +50,5 @@
 from .prf import *
 from .lightcurve import *
 from .targetpixelfile import *
+from .periodogram import *
 from .utils import *

pyke/lightcurve.py

@@ -9,6 +9,8 @@
 from bs4 import BeautifulSoup
 from .utils import running_mean, channel_to_module_output, KeplerQualityFlags
 from matplotlib import pyplot as plt
+from .periodogram import Periodogram
+
 
 __all__ = ['LightCurve', 'KeplerLightCurveFile', 'KeplerCBVCorrector',
            'SimplePixelLevelDecorrelationDetrender']
@@ -93,33 +95,13 @@ def flatten(self, window_length=101, polyorder=3, **kwargs):
             flatten_lc.flux_err = lc_clean.flux_err / trend_signal
         trend_lc = copy.copy(self)
         trend_lc.flux = trend_signal
-
         return flatten_lc, trend_lc
 
-    def fold(self, period, phase=0.):
-        """Folds the lightcurve at a specified ``period`` and ``phase``.
-
-        This method returns a new ``LightCurve`` object in which the time
-        values range between -0.5 to +0.5.  Data points which occur exactly
-        at ``phase`` or an integer multiple of `phase + n*period` have time
-        value 0.0.
-
-        Parameters
-        ----------
-        period : float
-            The period upon which to fold.
-        phase : float, optional
-            Time reference point.
+    def draw(self):
+        raise NotImplementedError("Should we implement a LightCurveDrawer class?")
 
-        Returns
-        -------
-        folded_lightcurve : LightCurve object
-            A new ``LightCurve`` in which the data are folded and sorted by
-            phase.
-        """
-        fold_time = ((self.time - phase + 0.5 * period) / period) % 1 - 0.5
-        sorted_args = np.argsort(fold_time)
-        return LightCurve(fold_time[sorted_args], self.flux[sorted_args])
+    def to_csv(self):
+        raise NotImplementedError()
 
     def remove_nans(self):
         """Removes cadences where the flux is NaN.
@@ -218,8 +200,94 @@ def cdpp(self, transit_duration=13, savgol_window=101, savgol_polyorder=2,
         cdpp_ppm = np.std(mean) * 1e6
         return cdpp_ppm
 
-    def to_csv(self):
-        raise NotImplementedError()
+
+    def periodogram(self, minper=None, maxper=None, nterms=1):
+        """
+        Creates a periodogram object
+
+        Parameters
+        ----------
+        minper : float
+            Minimum period to search
+        maxper : float
+            Maximum period to search
+        nterms : int
+            Number of terms to use for Lomb-Scargle periodogram. (Default 1)
+
+        Returns
+        -------
+        p : periodogram object
+            Periodogram object
+        """
+        p = Periodogram(self.time, self.flux, self.flux_err, minper, maxper, nterms)
+        return p
+
+    def find_period(self, minper=None, maxper=None, nterms=1):
+        """
+        Finds the best fit period in the light curve
+        Parameters
+        ----------
+        minper : float
+            Minimum period to search
+        maxper : float
+            Maximum period to search
+        nterms : int
+            Number of terms to use for Lomb-Scargle periodogram. (Default 1)
+
+        Returns
+        -------
+        period : float
+            Best fit period
+        phase : float
+            Best fit phase
+        """
+        p = self.periodogram(minper, maxper, nterms)
+        period = p.per()
+        ok = np.isfinite(self.flux)
+        s = np.argsort(self.time[ok]/period % 1)
+        smooth = signal.savgol_filter(self.flux[ok][s],21,5)
+        m = np.argmin(smooth)
+        phase = (self.time[ok][s]/period % 1)[m]
+        return period,phase
+
+    def fold(self, period=None, phase=None, plot=False, **kwargs):
+        """Folds a lightcurve on the bestfit period
+
+        Parameters
+        ----------
+        period : float
+            Period to fold at. If none, the best fit period will be found.
+        phase : float
+            Phase to fold at. If none, the best fit phase will be found.
+        plot : bool
+            Whether or not to return a plot
+        **kwargs : dict
+            Dictionary of arguments to be passed to `periodogram`.
+
+        Returns
+        -------
+        period : float
+            Best fit period
+        phase : float
+            Best fit phase
+        """
+        if period is None:
+            p,ph = self.find_period(**kwargs)
+            period = p
+            if phase is None:
+                phase = ph
+        if phase is None:
+            phase=0
+        s = np.argsort(self.time/period % 1)
+        ph = ((self.time - ((phase+0.5) * period)) / period) % 1 - 0.5
+        if plot:
+            plt.plot(ph, self.flux,marker='.',ls='',ms=1)
+            plt.xlabel('Phase')
+            plt.ylabel('Counts ($e^-s^{-1}$)')
+            plt.title('Best fit Period: {:5.3} days'.format(period))
+        folded_flux = self.flux[np.argsort(ph)]
+        folded_time = np.sort(ph)
+        return folded_time,folded_flux
 
     def plot(self, ax=None, normalize=True, xlabel='Time - 2454833 (days)',
              ylabel='Normalized Flux', title=None, color='#363636', fill=False,
@@ -274,7 +342,6 @@ def plot(self, ax=None, normalize=True, xlabel='Time - 2454833 (days)',
         ax.set_ylabel(ylabel, {'color': 'k'})
         return ax
 
-
 class KeplerLightCurve(LightCurve):
     """Defines a light curve class for NASA's Kepler and K2 missions.
 
@@ -325,6 +392,18 @@ def __init__(self, time, flux, flux_err=None, centroid_col=None,
     def to_fits(self):
         raise NotImplementedError()
 
+    def periodogram(self, **kwargs):
+        p = super(KeplerLightCurve,self).periodogram(**kwargs)
+        return p
+
+    def find_period(self,**kwargs):
+        period, phase = super(KeplerLightCurve,self).find_period(**kwargs)
+        return period, phase
+
+    def fold(self,**kwargs):
+        folded_time, folded_flux = super(KeplerLightCurve,self).fold(**kwargs)
+        return folded_time, folded_flux
+
 
 class KeplerLightCurveFile(object):
     """Defines a class for a given light curve FITS file from NASA's Kepler and
@@ -455,6 +534,21 @@ def _flux_types(self):
         types = [n for n in types if not ('ERR' in n)]
         return types
 
+    def periodogram(self, fluxtype = 'PDCSAP_FLUX', **kwargs):
+        f = self.get_lightcurve(fluxtype)
+        p = Periodogram(self.time,f.flux,f.flux_err,**kwargs)
+        return p
+
+    def find_period(self, fluxtype = 'PDCSAP_FLUX', **kwargs):
+        f = self.get_lightcurve(fluxtype)
+        period, phase = f.find_period(**kwargs)
+        return period, phase
+
+    def fold(self, fluxtype = 'PDCSAP_FLUX', **kwargs):
+        f = self.get_lightcurve(fluxtype)
+        folded_time, folded_flux = f.fold(**kwargs)
+        return folded_time, folded_flux
+
     def plot(self, plottype=None, **kwargs):
         """Plot all the flux types in a light curve.
 
@@ -477,7 +571,6 @@ def plot(self, plottype=None, **kwargs):
             kwargs['color'] = 'C{}'.format(idx)
             lc.plot(label=pl, **kwargs)
 
-
 class Detrender(object):
     """
     """
@@ -487,12 +580,10 @@ def detrend(self):
         """
         pass
 
-
 class SystematicsCorrector(object):
     def correct(self):
         pass
 
-
 class KeplerCBVCorrector(SystematicsCorrector):
     r"""Remove systematic trends from Kepler light curves by fitting
     cotrending basis vectors.

pyke/periodogram.py

@@ -0,0 +1,112 @@
+import numpy as np
+from astropy.stats import LombScargle
+import warnings
+from scipy import signal
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+
+
+__all__ = ['Periodogram']
+
+class Periodogram(object):
+    """Defines a periodogram class for Kepler/K2 data. Searches for periods using
+    astropy.LombScargle and Box Least Squares.
+
+    Attributes
+    ----------
+    time : array-like
+        Time measurements
+    flux : array-like
+        Data flux for every time point
+    flux_err : array-like
+        Uncertainty on each flux data point
+    minper : float
+        Minimum period to search
+    maxper : float
+        Maximum period to search
+    nterms : int
+        Number of terms to use for Lomb-Scargle periodogram. (Default 1)
+    """
+    def __init__(self, time, flux, flux_err=None, minper=None, maxper=None, nterms=1):
+        self.time = time
+        self.flux = flux
+        if flux_err is None:
+            flux_err = np.copy(flux)*0.
+        self.flux_err = flux_err
+        self.minper = minper
+        self.maxper = maxper
+        self.nterms = nterms
+        self.dur = time.max() - time.min()
+        self.npoints = len(self.time)
+        if self.minper is None:
+            self.minper = np.median(self.time[1:] - self.time[0:-1])*4
+        if self.maxper is None:
+            self.maxper = np.nanmax(self.time - self.time.min())/4.
+        self.clean()
+        self.LombScargle()
+
+    def clean(self):
+        """
+        Remove infinite values
+        """
+        ok = np.isfinite(self.flux)
+        self.time = self.time[ok]
+        self.flux = self.flux[ok]
+        self.flux_err = self.flux_err[ok]
+
+    def LombScargle(self,samples=40):
+        """
+        Creates a Lomb Scargle Periodogram
+
+        Parameters
+        ----------
+        samples : int
+            Number of samples to take in the
+        """
+        LS = LombScargle(self.time, self.flux, self.flux.max()-self.flux.min(), nterms=self.nterms)
+        frequency, power = LS.autopower(maximum_frequency=1./self.minper, minimum_frequency=1./self.maxper, samples_per_peak=samples)
+        self.lomb_per = 1. / frequency[np.argmax(power)]
+        self.lomb_periods = 1. / frequency
+        self.lomb_power = power
+
+        s = np.argsort(self.time / self.lomb_per % 1)
+        dp = self.npoints * (self.lomb_per * 5 / self.dur)
+        if dp % 2 == 0:
+            dp += 1
+        smooth = signal.savgol_filter(self.flux[s], dp, 3)
+        m = np.argmin(smooth)
+        self.lomb_phase = (self.time[s] / period % 1)[m]
+
+    def BLS(self):
+        '''Not implemented'''
+        pass
+
+    def plot(self, ax=None, line=True, **kwargs):
+        """
+        Plot the periodogram object
+
+        Parameters
+        ----------
+        ax : matplotlib frame
+            Frame to plot the figure into. If unspecified, creates a new figure and frame.
+        line : bool
+            Plot a line at the bestfit period
+        **kwargs : dict
+            Dictionary of keyword values to pass to 'matplotlib.pyplot.plot'
+        """
+        if ax is None:
+            fig, ax = plt.subplots()
+        with mpl.style.use('ggplot'):
+            ax.plot(self.lomb_periods, self.lomb_power,**kwargs)
+            ax.set(xlabel='Period (days)', ylabel='Lomb Scargle Power')
+            if line:
+                plt.axvline(self.lomb_per,ls='--',lw=2,color='black')
+                ax.text(self.lomb_per*1.1,self.lomb_power.max(),'Best Fit Period', ha='left',va='center')
+
+    def per(self):
+        '''Returns the best fit period.'''
+        return self.lomb_per
+
+    def phase(self):
+        '''Returns the best fit phase.'''
+        return self.lomb_phase

pyke/tests/test_periodogram.py

@@ -0,0 +1,22 @@
+import numpy as np
+from astropy.utils.data import get_pkg_data_filename
+
+from ..periodogram import Periodogram
+
+filename_lc = get_pkg_data_filename("data/golden-lc.fits")
+
+def test_init():
+    '''Not implemented'''
+    pass
+
+def test_BLS():
+    '''Not implemented'''
+    pass
+
+def test_plot():
+    '''Not implemented'''
+    pass
+
+def test_period():
+    '''Not implemented'''
+    pass