Added function to determine rotation period using hardcorr

dot/io.py

@@ -4,11 +4,13 @@
 import math
 import pandas as pd
 from lightkurve import LightCurve, search_lightcurvefile
+from hardcorr import interpolated_acf
+from hardcorr import dominant_period
 import h5py
 
 
 __all__ = ['save_results', 'load_results', 'ab_dor_example_lc',
-           'load_light_curve', 'load_rotation_period']
+           'load_light_curve', 'load_rotation_period', 'period_finder_hardcorr']
 
 hdf5_archive_disk = '/110k_pdcsap/'
 hdf5_index_path = '110k_rotation_mcquillan_pdcsap_smooth_index_0724.csv'
@@ -162,7 +164,7 @@ def download_from_lightkurve(kic):
 
 def load_rotation_period(kic, data_path=None, index_file=None):
     """
-    Extract McQuillan rotation period from KIC number on google cloud platform.
+    Extract McQuillan rotation period from csv file on google cloud platform.
 
     Parameters
     ----------
@@ -197,3 +199,27 @@ def load_rotation_period(kic, data_path=None, index_file=None):
         print(f'Using McQuillan period for KIC {kic}.')
         star_prot = star_prot_list[0]
         return star_prot
+
+
+def period_finder_hardcorr(lc, n_peaks=1):
+    """
+    Run auto-correlation function using hardcorr to determine rotation period.
+
+    Parameters
+    ----------
+    lc : array
+        light curve object
+    Returns
+    -------
+    period_list : list
+        list of n_peaks dominant rotation period as determined through ACF
+    """
+
+    time = lc.time
+    flux = np.array([f/np.nanmedian(lc.flux) - 1 for f in lc.flux]) # normalize around 0
+
+    lag, acf = interpolated_acf(time[~np.isnan(flux)], flux[~np.isnan(flux)])
+    period_list = [dominant_period(lag, acf, n_peaks=n_peaks)]
+    return period_list
+
+

run_dot.py

@@ -12,7 +12,7 @@
 import matplotlib.pyplot as plt
 from tqdm import tqdm
 
-from dot import Model, save_results, load_results, ab_dor_example_lc, load_light_curve, load_rotation_period
+from dot import Model, save_results, load_results, ab_dor_example_lc, load_light_curve, load_rotation_period, period_finder_hardcorr
 from dot.plots import corner, posterior_shear,  posterior_predictive, movie
 
 import warnings
@@ -49,12 +49,6 @@ def run_dot(args):
     if mission=="Kepler":
         print(f'Loading light curve for KIC {target}...')
         lc = load_light_curve(target)
-        # Extract rotation periods
-        try:
-            rotation_period = load_rotation_period(target)
-        except:
-            print(f'Target KIC {target} does not have a McQuillan period.')
-            rotation_period = args.rotation_period
 
     else:
         # Get data from lightkurve
@@ -66,6 +60,25 @@ def run_dot(args):
             sector=sector
         ).download_all().PDCSAP_FLUX.stitch()
 
+    # --------Cleaning the light curve--------
+
+    lc = lc.remove_nans()
+    lc = lc.remove_outliers(sigma=5.0) # double check range for sigma clipping
+
+    # --------Extract rotation periods--------
+
+    if args.rotation_period is None:
+        try: # Load rotation period locally
+            rotation_period = load_rotation_period(target)
+        except:
+            print('Running ACF using hardcorr...')
+            rotation_period_list = period_finder_hardcorr(lc, n_peaks=1)
+            rotation_period = rotation_period_list[0]
+    else:
+        rotation_period = args.rotation_period
+
+    print(f'Rotation period: {rotation_period} days')
+
     # --------Setting up model--------
 
     min_time = lc.time.min()
@@ -160,7 +173,7 @@ def run_dot(args):
         "--sector", type=int, default=None, help="Quarter or Sector",
     )
     parser.add_argument(
-        "--rotation-period", type=float, default=0.5, help="Stellar rotation period",
+        "--rotation-period", type=float, default=None, help="Stellar rotation period",
     )
     # =================================
     # FIT PARAMS
@@ -175,7 +188,7 @@ def run_dot(args):
         "--limit-duration", type=int, default=2, help="Length of interval to fit in days",
     )
     parser.add_argument(
-        "--n-intervals", type=int, default=2, help="Number of time intervals to fit",
+        "--n-intervals", type=int, default=1, help="Number of time intervals to fit",
     )
     # =================================
     # SMC PARAMS