flaredetector: util: add argument if peaks should be sorted or not by highest

main/flaredetector/util.py

@@ -4,7 +4,7 @@ from scipy.signal import find_peaks
 from lightkurve.periodogram import Periodogram
 from lightkurve.lightcurve import LightCurve
 
-def findMaxIndices(lc, num=3, distance=100, height=(None, None)):
+def findMaxIndices(lc, num=3, distance=100, height=(None, None), sortByHighest=False):
     if(isinstance(lc, Periodogram)):
         data = lc.power
     elif(isinstance(lc, LightCurve)):
@@ -14,5 +14,7 @@ def findMaxIndices(lc, num=3, distance=100, height=(None, None)):
     peak_indices, peak_dict = find_peaks(data, height=height,
                                          distance=distance)
     peak_heights = peak_dict["peak_heights"]
-    maxIndices = peak_indices[np.argsort(peak_heights)[-num:][::-1]]
-    return maxIndices
+    if(sortByHighest):
+        return peak_indices[np.argsort(peak_heights)[-num:][::-1]]
+    else:
+        return peak_indices[0:num]