Merge pull request #95 from spacetelescope/md_centroid_figs

Md centroid figs

exovetter/centroid/centroid.py

@@ -78,13 +78,13 @@ def compute_diff_image_centroids(
     transits = getIngressEgressCadences(
         time, period_days, epoch, duration_days)
 
-    figs = []
+    axs = []
     centroids = []
 
     for i in range(len(transits)):
         if i not in remove_transits:
             cin = transits[i]
-            cents, fig = measure_centroids(
+            cents, ax = measure_centroids(
                 cube, 
                 cin, 
                 max_oot_shift_pix=max_oot_shift_pix,
@@ -93,15 +93,15 @@ def compute_diff_image_centroids(
             )
 
             if plot == True:
-                fig.suptitle('Transit '+str(i))
+                plt.gcf().suptitle('Transit '+str(i))
 
             centroids.append(cents)
-            figs.append(fig)
+            axs.append(ax)
 
     centroids = np.array(centroids)
     all_transits = list(np.arange(len(transits)))
     kept_transits = [x for x in all_transits if x not in remove_transits]
-    return centroids, figs, kept_transits
+    return centroids, axs, kept_transits
 
 
 def measure_centroid_shift(centroids, kept_transits, plot=False):
@@ -297,7 +297,7 @@ def measure_centroids(cube, cin, max_oot_shift_pix=0.5, starloc_pix = None, plot
         flag = 2
     out.append(flag)
 
-    return out, fig
+    return out, ax
 
 
 def generateDiffImg(cube, transits, starloc_pix = None, plot=False):

exovetter/vetters.py

@@ -585,7 +585,7 @@ def plot(self):  # pragma: no cover
 class Centroid(BaseVetter):
     """Class to handle centroid vetting"""
 
-    def __init__(self, lc_name="flux", diff_plots=False, centroid_plots=False):
+    def __init__(self, lc_name="flux", diff_plots=False, centroid_plots=False, starloc_pix=None):
         """
         Parameters
         ----------
@@ -616,8 +616,9 @@ def __init__(self, lc_name="flux", diff_plots=False, centroid_plots=False):
         self.metrics = None
         self.diff_plots = diff_plots
         self.centroid_plots = centroid_plots
+        self.starloc_pix = starloc_pix
 
-    def run(self, tce, lk_tpf, starloc_pix = None, plot=False, remove_transits=None):
+    def run(self, tce, lk_tpf, plot=False, remove_transits=None):
         """Runs cent.compute_diff_image_centroids and cent.measure_centroid_shift
         to populate the vetter object.
 
@@ -668,10 +669,10 @@ def run(self, tce, lk_tpf, starloc_pix = None, plot=False, remove_transits=None)
 
         centroids, figs, kept_transits = cent.compute_diff_image_centroids(
             time, cube, period_days, epoch, duration_days, 
-            remove_transits, starloc_pix=starloc_pix, plot=self.diff_plots)
+            remove_transits, starloc_pix=self.starloc_pix, plot=self.diff_plots)
 
-        if (starloc_pix is not None) and (len(starloc_pix) == 2):
-            offset, signif, fig = cent.measure_centroid_shift_cat(centroids, kept_transits, starloc_pix, self.centroid_plots)
+        if (self.starloc_pix is not None) and (len(self.starloc_pix) == 2):
+            offset, signif, fig = cent.measure_centroid_shift_cat(centroids, kept_transits, self.starloc_pix, self.centroid_plots)
         else:
             offset, signif, fig = cent.measure_centroid_shift(centroids, kept_transits, self.centroid_plots)