exovetter 0.0.6__tar.gz → 0.0.8__tar.gz

{exovetter-0.0.6 → exovetter-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: exovetter
-Version: 0.0.6
+Version: 0.0.8
 Summary: Exoplanet vetting package
 Home-page: https://github.com/spacetelescope/exovetter
 Author: Susan Mullally et al.

{exovetter-0.0.6 → exovetter-0.0.8}/exovetter/centroid/centroid.py

@@ -14,6 +14,7 @@ def compute_diff_image_centroids(
         duration_days,
         remove_transits, 
         max_oot_shift_pix=1.5,
+        starloc_pix = None,
         plot=False
 ):
     """Compute difference image centroid shifts for every transit in a dataset.
@@ -42,6 +43,9 @@ def compute_diff_image_centroids(
         (float) Duration of transit. 
     remove_transits
         (list) List of 0 indexed transit integers to not calculate on.
+    starloc_pix
+        (2d array) catalog location of target star for plotting.
+        Default is None.
     max_oot_shift_pix
         (float) Passed to `fastpsffit.fastGaussianPrfFit()
 
@@ -74,29 +78,30 @@ 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,
+                starloc_pix = starloc_pix,
                 plot=plot
             )
 
             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):
@@ -140,6 +145,51 @@ def measure_centroid_shift(centroids, kept_transits, plot=False):
         fig = covar.diagnostic_plot(dcol, drow, kept_transits, flags)
     return offset_pix, signif, fig
 
+def measure_centroid_shift_cat(centroids, kept_transits, starloc_pix, plot=False):
+    """Measure the average offset of the DIC centroids from the catalog position.
+
+    Inputs
+    ----------
+    centroids
+        (2d np array) Output of :func:`compute_diff_image_centroids`
+
+    kept_transits
+        (list) List of 0 indexed transit integers to calculate on.
+    
+    starloc_pix
+        (2d np array) col,row expected location of target star
+
+    Returns
+    -----------
+    offset
+        (float) Size of offset in pixels (or whatever unit `centroids`
+        is in)
+    signif
+        (float) The statistical significance of the transit. Values
+        close to 1 mean the transit is likely on the target star.
+        Values less than ~1e-3 suggest the target is not the
+        source of the transit.
+    fig
+        A figure handle. Is **None** if plot is **False**
+    """
+
+    # DIC - catalog
+    # dcol = centroids[:, 5] - centroids[:, 0]
+    # drow = centroids[:, 4] - centroids[:, 1]
+    dcol = centroids[:, 4] - starloc_pix[0]
+    drow = centroids[:, 5] - starloc_pix[1]
+
+    flags = centroids[:, -1].astype(bool)
+
+    offset_pix, signif = covar.compute_offset_and_signif(
+        dcol[~flags], drow[~flags])
+
+    fig = None
+    if plot:
+        fig = covar.diagnostic_plot(dcol, drow, kept_transits, flags)
+        
+    return offset_pix, signif, fig
+
 
 def getIngressEgressCadences(time, period_days, epoch_btjd, duration_days):
     assert np.all(np.isfinite(time))
@@ -151,7 +201,7 @@ def getIngressEgressCadences(time, period_days, epoch_btjd, duration_days):
     return transits
 
 
-def measure_centroids(cube, cin, max_oot_shift_pix=0.5, plot=False):
+def measure_centroids(cube, cin, max_oot_shift_pix=0.5, starloc_pix = None, plot=False):
     """Private function of :func:`compute_diff_image_centroids`
 
     Computes OOT, ITR and diff images for a single transit event,
@@ -209,17 +259,31 @@ def measure_centroids(cube, cin, max_oot_shift_pix=0.5, plot=False):
         if diffSoln.success:
             clr = "green"
 
+        fig = plt.gcf()
+        axlist = fig.axes 
+        #assert len(axlist) == 3, axlist
+
         res = diffSoln.x
-        disp.plotCentroidLocation(res[0], res[1], marker="^", color=clr,
-                                  label="diff")
+        for ax in axlist:
+            if ax.get_label() == '<colorbar>':
+                continue
+
+            plt.sca(ax)
+            disp.plotCentroidLocation(res[0], res[1], marker="^", color=clr,
+                                    label="diff")
 
-        res = ootSoln.x
-        disp.plotCentroidLocation(res[0], res[1], marker="o", color=clr,
-                                  label="OOT")
+            res1 = ootSoln.x
+            disp.plotCentroidLocation(res1[0], res1[1], marker="o", color=clr,
+                                    label="OOT")
 
-        res = intransSoln.x
-        disp.plotCentroidLocation(res[0], res[1], marker="+", color=clr,
-                                  label="InT")
+            res2 = intransSoln.x
+            disp.plotCentroidLocation(res2[0], res2[1], marker="+", color=clr,
+                                    label="InT")
+            
+            if starloc_pix is not None:
+                disp.plotCentroidLocation(starloc_pix[0], starloc_pix[1], marker="*",
+                                        color='red', label="Cat", ms=10)
+            
         plt.legend(fontsize=12, framealpha=0.7, facecolor='silver')
 
     out = []
@@ -236,7 +300,7 @@ def measure_centroids(cube, cin, max_oot_shift_pix=0.5, plot=False):
     return out, ax
 
 
-def generateDiffImg(cube, transits, plot=False):
+def generateDiffImg(cube, transits, starloc_pix = None, plot=False):
     """Generate a difference image.
 
     Also generates an image for each the $n$ cadedences before
@@ -249,6 +313,8 @@ def generateDiffImg(cube, transits, plot=False):
         (np 3 array) Datacube of postage stamps
     transits
         (2-tuples) Indices of the first and last cadence
+    starloc_pix
+        (np 2 element array) col, row position of star
 
     Optional Inputs
     -----------------
@@ -286,6 +352,7 @@ def generateDiffImg(cube, transits, plot=False):
         fig = plt.figure()
         fig.set_size_inches(16, 4)
         disp.plotTransit(fig, oot, during, diff)
+
     else:
         fig = None
 

{exovetter-0.0.6 → exovetter-0.0.8}/exovetter/centroid/disp.py

@@ -8,15 +8,15 @@ import numpy as np
 
 def plotTransit(fig, oot, during, diff, **kwargs):
 
-    fig.add_subplot(131)
+    ax1 = fig.add_subplot(131)
     plotImage(oot, **kwargs)
     plt.title("OOT")
 
-    fig.add_subplot(132)
+    ax2 = fig.add_subplot(132)
     plotImage(during, **kwargs)
     plt.title("In-transit")
 
-    fig.add_subplot(133)
+    ax3 = fig.add_subplot(133)
     plotDiffImage(diff, **kwargs)
     plt.title("Difference")
 

{exovetter-0.0.6 → exovetter-0.0.8}/exovetter/utils.py

@@ -2,12 +2,19 @@
 
 import sys
 import warnings
-
+from exovetter import lightkurve_utils
+from exovetter import utils
+from exovetter import const as exo_const
+import astropy.units as u
+from exovetter import vetters as vet
+from matplotlib.backends.backend_pdf import PdfPages
+import matplotlib.pyplot as plt
 import numpy as np
+import os
 
 __all__ = ['sine', 'estimate_scatter', 'mark_transit_cadences', 'median_detrend', 
            'plateau', 'set_median_flux_to_zero', 'set_median_flux_to_one', 'sigmaClip', 
-           'get_mast_tce', 'WqedLSF', 'compute_phases', 'first_epoch']
+           'get_mast_tce', 'WqedLSF', 'compute_phases', 'first_epoch', 'run_all']
 
 def sine(x, order, period=1):
     """Sine function for SWEET vetter."""
@@ -653,3 +660,164 @@ def first_epoch(epoch, period, lc):
         first_epoch = epoch + N*period
 
     return first_epoch
+
+def run_all(tce, lc, tpf=None, vetters=None, remove_metrics=None, plot=False, plot_dir=None):
+    """Run a set of vetters on a tce and lc, returning a dictionary of all metrics collected
+    
+    Parameters
+    ----------
+    tce : tce object
+        tce object is a dictionary that contains information about the tce
+        to vet, like period, epoch, duration, depth
+
+    lc : lightkurve object
+        lightkurve object with the time and flux to use for vetting.
+
+    tpf : obj
+        ``lightkurve`` target pixel file object with pixels in column lc_name
+
+    vetters : list
+        list of vetter objects to run on, ie [vet.ModShift(), vet.OddEven(dur_frac=0.75)]
+        Defaults to all vetters
+    
+    remove_metrics : list
+        metrics to not store, defaults to removing plotting values
+
+    plot : bool
+        Option to return a pdf of the vetting diagnostic plots, defaults to False
+
+    plot_dir : str
+        Path to store diagnostic pdfs in, defaults to current working directory
+
+    Returns
+    -------
+    results_dict : dictionary
+        Dictionary of the kept vetting metrics
+    """
+    
+    # Set initial parameters
+    if not vetters:
+        vetters = [vet.VizTransits(), vet.ModShift(), vet.Lpp(), vet.OddEven(), vet.TransitPhaseCoverage(), vet.Sweet(), vet.LeoTransitEvents(), vet.Centroid()]
+
+    if not remove_metrics:
+        remove_metrics = ['plot_data']
+
+    if not plot_dir:
+        plot_dir = os.getcwd()+'/'
+
+    if not tpf:
+        if any(vetter.__class__.__name__ == 'Centroid' for vetter in vetters):
+            raise Exception("TPF file required while running centroid")
+
+    # Run all listed vetters
+    results_list = []
+
+    if not plot:
+        for vetter in vetters:
+            if vetter.__class__.__name__ != 'Centroid':
+                vetter_results = vetter.run(tce, lc, plot=False) # dictionary returned from each vetter
+                results_list.append(vetter_results)
+            else:
+                vetter_results = vetter.run(tce, tpf, plot=False) # centroid uses tpf rather than lc
+                results_list.append(vetter_results)
+
+    else:
+        plot_name = lc.LABEL
+
+        diagnostic_plot = PdfPages(plot_dir+plot_name+'.pdf') # initialize a pdf to save each figure into
+        plot_figures = []
+        
+        # Manually run viz transits with an extra mark cadences plot
+        cadences_plot = mark_cadences_plot(lc, tce)
+        diagnostic_plot.savefig(cadences_plot)
+
+        for vetter in vetters:
+            if vetter.__class__.__name__ not in ['VizTransits', 'Centroid', 'LeoTransitEvents']: # viz_transits and Centroid generate more than one figures so handle later
+                vetter_results = vetter.run(tce, lc, plot=True) # dictionary returned from each vetter
+                plot_figures.append(plt.gcf())
+                plt.close() # Make them not show up if running in a notebook
+                results_list.append(vetter_results)
+
+            if vetter.__class__.__name__ == 'Centroid': # centroid produces 2 plots, the second of which is the most useful so just collect that one
+                vetter_results = vet.Centroid(lc_name=vetter.lc_name, diff_plots=False, centroid_plots=True).run(tce, tpf)
+                plot_figures.append(plt.gcf())
+                plt.close()
+                results_list.append(vetter_results)
+                
+        # run viz_transits plots
+        transit = vet.VizTransits(transit_plot=True, folded_plot=False).run(tce, lc)
+        transit_plot = plt.gcf()
+        transit_plot.suptitle(plot_name+' Transits')
+        transit_plot.tight_layout()
+        plt.close()
+        diagnostic_plot.savefig(transit_plot)
+
+        folded = vet.VizTransits(transit_plot=False, folded_plot=True).run(tce, lc)
+        folded_plot = plt.gcf()
+        folded_plot.suptitle(plot_name+' Folded Transits')
+        folded_plot.tight_layout()
+        plt.close()
+        diagnostic_plot.savefig(folded_plot)
+        
+        # Save each diagnostic plot stored in plot_figures to diagnostic_plot pdf file
+        for plot in plot_figures:
+            diagnostic_plot.savefig(plot)
+
+        diagnostic_plot.close()
+    
+    # Convert to a single dictionary output
+    results_dict = {k: v for d in results_list for k, v in d.items()} # Combine all dictionaries returned from vetters
+
+    # delete specified metrics
+    for key in remove_metrics:
+        if results_dict.get(key):
+            del results_dict[key]
+    
+    return results_dict
+
+
+def mark_cadences_plot(lc, tce):
+    """return figure object of the lightcurve with epochs oeverplotted"""
+    
+    fig, ax1 = plt.subplots(nrows=1, ncols=1, figsize=(9,5))
+
+    # Get lightcurve data
+    time, flux, time_offset_str = lightkurve_utils.unpack_lk_version(lc, "flux")  # noqa: E50
+    time_offset_q = getattr(exo_const, time_offset_str)
+    epoch = tce.get_epoch(time_offset_q).to_value(u.day)
+    
+    # Get points of lightcurve found to be in transit
+    period = tce["period"].to_value(u.day)
+    dur = tce["duration"].to_value(u.day)
+    intransit = mark_transit_cadences(time, period, epoch, dur)
+
+    # Plot epoch
+    ax1.axvline(x=epoch, lw='0.6', color='r', label='epoch', alpha=0.5)
+
+    # Plot transit train
+    ax1.plot(time, flux, lw=0.72, alpha=0.9)
+    ax1.scatter(time, flux, color='k', s=3, label='cadences', alpha=0.5)
+    
+    # TODO This only plots forward in time from the epoch, works fine assuming tce epoch is first in light curve but not robust
+    transit_epochs = epoch-dur/2
+    while transit_epochs <= time[-1]:
+        ax1.axvline(x=transit_epochs, lw='0.6', color='r', alpha=0.3, ls='--')
+        transit_epochs = transit_epochs+dur
+        ax1.axvline(x=transit_epochs, lw='0.6', color='r', alpha=0.3, ls='--')
+        transit_epochs = transit_epochs-dur + period
+        
+    # Plot cadences in transit
+    ax1.scatter(time[intransit], flux[intransit], color='r', s=4, label='cadences in transit');
+
+    # Plotting params
+    ax1.set_ylabel('Flux')
+    ax1.set_xlabel('Time '+time_offset_str)
+    ax1.set_title(lc.label+' period='+'{0:.2f}'.format(period)+'d, dur='+'{0:.2f}'.format(dur)+'d')
+    ax1.legend();
+
+    cadences_plot = plt.gcf()
+    plt.close()
+
+    return cadences_plot
+
+    

{exovetter-0.0.6 → exovetter-0.0.8}/exovetter/version.py

@@ -5,4 +5,4 @@ try:
     from setuptools_scm import get_version
     version = get_version(root='..', relative_to=__file__)
 except Exception:
-    version = '0.0.6'
+    version = '0.0.8'