KeplerCBVCorrector can now take a KeplerLightCurve Object

pyke/lightcurve.py

@@ -307,6 +307,8 @@ class KeplerLightCurve(LightCurve):
         Array indicating the quality of each data point
     quality_bitmask : int
         Bitmask specifying quality flags of cadences that should be ignored
+    quality_mask : array-like
+        Mask applied to tpf
     channel : int
         Channel number
     campaign : int
@@ -322,14 +324,20 @@ class KeplerLightCurve(LightCurve):
     """
 
     def __init__(self, time, flux, flux_err=None, centroid_col=None,
-                 centroid_row=None, quality=None, quality_bitmask=None,
+                 centroid_row=None, quality=None,
+                 quality_bitmask=KeplerQualityFlags.DEFAULT_BITMASK,
+                 quality_mask=None,
                  channel=None, campaign=None, quarter=None, mission=None,
                  cadenceno=None, keplerid=None):
         super(KeplerLightCurve, self).__init__(time, flux, flux_err)
         self.centroid_col = centroid_col
         self.centroid_row = centroid_row
         self.quality = quality
         self.quality_bitmask = quality_bitmask
+        if quality_mask is None:
+            self.quality_mask = np.asarray(self.quality & self.quality_bitmask, dtype=bool)
+        else:
+            self.quality_mask = quality_mask
         self.channel = channel
         self.campaign = campaign
         self.quarter = quarter
@@ -520,9 +528,9 @@ class KeplerCBVCorrector(SystematicsCorrector):
 
     Attributes
     ----------
-    lc_file : KeplerLightCurveFile object or str
-        An instance from KeplerLightCurveFile or a path for the .fits
-        file of a NASA's Kepler/K2 light curve.
+    lightcurve : KeplerLightCurveFile object or KeplerLightCurve object or str
+        An instance of either KeplerLightCurveFile or KeplerLightCurve
+        or a path for the .fits file of a NASA's Kepler/K2 light curve.
     loss_function : oktopus.Likelihood subclass
         A class that describes a cost function.
         The default is :class:`oktopus.LaplacianLikelihood`, which is tantamount
@@ -544,30 +552,33 @@ class KeplerCBVCorrector(SystematicsCorrector):
     >>> plt.legend() # doctest: +SKIP
     """
 
-    def __init__(self, lc_file, loss_function=oktopus.LaplacianLikelihood):
-        self.lc_file = lc_file
+    def __init__(self, lightcurve, loss_function=oktopus.LaplacianLikelihood):
+        self.lightcurve = lightcurve
         self.loss_function = loss_function
 
-        if self.lc_file.mission == 'Kepler':
+        if self.lightcurve.mission == 'Kepler':
             self.cbv_base_url = "http://archive.stsci.edu/missions/kepler/cbv/"
-        elif self.lc_file.mission == 'K2':
+        elif self.lightcurve.mission == 'K2':
             self.cbv_base_url = "http://archive.stsci.edu/missions/k2/cbv/"
 
     @property
-    def lc_file(self):
-        return self._lc_file
+    def lightcurve(self):
+        return self._lightcurve
 
-    @lc_file.setter
-    def lc_file(self, value):
-        # this enables `lc_file` to be either a string
+    @lightcurve.setter
+    def lightcurve(self, value):
+        # this enables `lightcurve` to be either a string
         # or an object from KeplerLightCurveFile
         if isinstance(value, str):
-            self._lc_file = KeplerLightCurveFile(value)
+            self._lightcurve = KeplerLightCurveFile(value).SAP_FLUX
         elif isinstance(value, KeplerLightCurveFile):
-            self._lc_file = value
+            self._lightcurve = value.SAP_FLUX
+        elif isinstance(value, KeplerLightCurve):
+            self._lightcurve = value
         else:
-            raise ValueError("lc_file must be either a string or a"
-                             " KeplerLightCurveFile instance, got {}.".format(value))
+            raise ValueError("lightcurve must be a string or either a"
+                             " KeplerLightCurveFile or KeplerLightCurve"
+                             " instance, got {}.".format(value))
 
     @property
     def coeffs(self):
@@ -594,31 +605,30 @@ def correct(self, cbvs=[1, 2]):
             The list of cotrending basis vectors to fit to the data. For example,
             [1, 2] will fit the first two basis vectors.
         """
-        module, output = channel_to_module_output(self.lc_file.channel)
+        module, output = channel_to_module_output(self.lightcurve.channel)
         cbv_file = pyfits.open(self.get_cbv_url())
         cbv_data = cbv_file['MODOUT_{0}_{1}'.format(module, output)].data
 
         cbv_array = []
         for i in cbvs:
-            cbv_array.append(cbv_data.field('VECTOR_{}'.format(i))[self.lc_file.quality_mask])
+            cbv_array.append(cbv_data.field('VECTOR_{}'.format(i))[self.lightcurve.quality_mask])
         cbv_array = np.asarray(cbv_array)
 
-        sap_lc = self.lc_file.SAP_FLUX
-        median_sap_flux = np.nanmedian(sap_lc.flux)
-        norm_sap_flux = sap_lc.flux / median_sap_flux - 1
-        norm_err_sap_flux = sap_lc.flux_err / median_sap_flux
+        median_flux = np.nanmedian(self.lightcurve.flux)
+        norm_flux = self.lightcurve.flux / median_flux - 1
+        norm_err_flux = self.lightcurve.flux_err / median_flux
 
         def mean_model(*theta):
             coeffs = np.asarray(theta)
             return np.dot(coeffs, cbv_array)
 
-        loss = self.loss_function(data=norm_sap_flux, mean=mean_model,
-                                  var=norm_err_sap_flux)
+        loss = self.loss_function(data=norm_flux, mean=mean_model,
+                                  var=norm_err_flux)
         self._opt_result = loss.fit(x0=np.zeros(len(cbvs)), method='L-BFGS-B')
         self._coeffs = self._opt_result.x
-        flux_hat = sap_lc.flux - median_sap_flux * mean_model(self._coeffs)
+        flux_hat = self.lightcurve.flux - median_flux * mean_model(self._coeffs)
 
-        return LightCurve(time=sap_lc.time, flux=flux_hat.reshape(-1))
+        return LightCurve(time=self.lightcurve.time, flux=flux_hat.reshape(-1))
 
     def get_cbv_url(self):
         # gets the html page and finds all references to 'a' tag
@@ -627,19 +637,19 @@ def get_cbv_url(self):
         soup = BeautifulSoup(requests.get(self.cbv_base_url).text, 'html.parser')
         cbv_files = [fn['href'] for fn in soup.find_all('a') if fn['href'].endswith('fits')]
 
-        if self.lc_file.mission == 'Kepler':
-            if self.lc_file.quarter < 10:
-                quarter = 'q0' + str(self.lc_file.quarter)
+        if self.lightcurve.mission == 'Kepler':
+            if self.lightcurve.quarter < 10:
+                quarter = 'q0' + str(self.lightcurve.quarter)
             else:
-                quarter = 'q' + str(self.lc_file.quarter)
+                quarter = 'q' + str(self.lightcurve.quarter)
             for cbv_file in cbv_files:
                 if quarter + '-d25' in cbv_file:
                     break
-        elif self.lc_file.mission == 'K2':
-            if self.lc_file.campaign <= 8:
-                campaign = 'c0' + str(self.lc_file.campaign)
+        elif self.lightcurve.mission == 'K2':
+            if self.lightcurve.campaign <= 8:
+                campaign = 'c0' + str(self.lightcurve.campaign)
             else:
-                campaign = 'c' + str(self.lc_file.campaign)
+                campaign = 'c' + str(self.lightcurve.campaign)
             for cbv_file in cbv_files:
                 if campaign in cbv_file:
                     break

pyke/targetpixelfile.py

@@ -213,12 +213,13 @@ def to_lightcurve(self, aperture_mask='pipeline'):
             Array containing the summed flux within the aperture for each
             cadence.
         """
-        if aperture_mask == 'pipeline':
-            aperture_mask = self.pipeline_mask
-        elif aperture_mask == 'all':
-            mask = ~np.isnan(self.hdu[1].data['FLUX'][100])
-            aperture_mask = np.ones((self.shape[1], self.shape[2]),
-                                    dtype=bool) * mask
+        if isinstance(aperture_mask, str):
+            if aperture_mask == 'pipeline':
+                aperture_mask = self.pipeline_mask
+            elif aperture_mask == 'all':
+                mask = ~np.isnan(self.hdu[1].data['FLUX'][100])
+                aperture_mask = np.ones((self.shape[1], self.shape[2]),
+                                        dtype=bool) * mask
 
         centroid_col, centroid_row = self.centroids(aperture_mask)
 
@@ -228,11 +229,14 @@ def to_lightcurve(self, aperture_mask='pipeline'):
                                 centroid_col=centroid_col,
                                 centroid_row=centroid_row,
                                 quality=self.quality,
+                                quality_bitmask=self.quality_bitmask,
+                                quality_mask=self._quality_mask(self.quality_bitmask),
                                 channel=self.channel,
                                 campaign=self.campaign,
                                 quarter=self.quarter,
                                 mission=self.mission,
-                                cadenceno=self.cadenceno)
+                                cadenceno=self.cadenceno,
+                                keplerid=self.keplerid)
 
     def centroids(self, aperture_mask='pipeline'):
         """Returns centroids based on sample moments.
@@ -250,12 +254,13 @@ def centroids(self, aperture_mask='pipeline'):
         col_centr, row_centr : tuple
             Arrays containing centroids for column and row at each cadence
         """
-        if aperture_mask == 'pipeline':
-            aperture_mask = self.pipeline_mask
-        elif aperture_mask == 'all':
-            mask = ~np.isnan(self.hdu[1].data['FLUX'][100])
-            aperture_mask = np.ones((self.shape[1], self.shape[2]),
-                                    dtype=bool) * mask
+        if isinstance(aperture_mask, str):
+            if aperture_mask == 'pipeline':
+                aperture_mask = self.pipeline_mask
+            elif aperture_mask == 'all':
+                mask = ~np.isnan(self.hdu[1].data['FLUX'][100])
+                aperture_mask = np.ones((self.shape[1], self.shape[2]),
+                                        dtype=bool) * mask
 
         yy, xx = np.indices(self.shape[1:]) + 0.5
         yy = self.row + yy