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

This commit is contained in:
2024-04-15 15:04:51 +02:00
parent 0f0001f7a3
commit f2a7cc08dc
+5 -3
View File
@@ -4,7 +4,7 @@ from scipy.signal import find_peaks
from lightkurve.periodogram import Periodogram from lightkurve.periodogram import Periodogram
from lightkurve.lightcurve import LightCurve 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)): if(isinstance(lc, Periodogram)):
data = lc.power data = lc.power
elif(isinstance(lc, LightCurve)): 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, peak_indices, peak_dict = find_peaks(data, height=height,
distance=distance) distance=distance)
peak_heights = peak_dict["peak_heights"] peak_heights = peak_dict["peak_heights"]
maxIndices = peak_indices[np.argsort(peak_heights)[-num:][::-1]] if(sortByHighest):
return maxIndices return peak_indices[np.argsort(peak_heights)[-num:][::-1]]
else:
return peak_indices[0:num]