diff --git a/pyke/__init__.py b/pyke/__init__.py index e847ee9d..830bb535 100644 --- a/pyke/__init__.py +++ b/pyke/__init__.py @@ -50,4 +50,5 @@ from .prf import * from .lightcurve import * from .targetpixelfile import * +from .periodogram import * from .utils import * diff --git a/pyke/lightcurve.py b/pyke/lightcurve.py index b8645b54..8b26baf0 100644 --- a/pyke/lightcurve.py +++ b/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. diff --git a/pyke/periodogram.py b/pyke/periodogram.py new file mode 100644 index 00000000..b332500f --- /dev/null +++ b/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 diff --git a/pyke/tests/test_periodogram.py b/pyke/tests/test_periodogram.py new file mode 100644 index 00000000..9ab5ebf0 --- /dev/null +++ b/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