36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import numpy as np
|
|
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), sortByHighest=False):
|
|
if(isinstance(lc, Periodogram)):
|
|
data = lc.power
|
|
elif(isinstance(lc, LightCurve)):
|
|
data = lc.flux
|
|
else:
|
|
return np.zeros(num)-1
|
|
peak_indices, peak_dict = find_peaks(data, height=height,
|
|
distance=distance)
|
|
peak_heights = peak_dict["peak_heights"]
|
|
if(sortByHighest):
|
|
return peak_indices[np.argsort(peak_heights)[-num:][::-1]]
|
|
else:
|
|
return peak_indices[0:num]
|
|
|
|
def plotFlarePeaks(ax, peaks, flux = None):
|
|
if(flux is None):
|
|
for peak in peaks:
|
|
ax.plot(peak["FlarePeakTime"], peak["FlarePeak"], "x", color="red")
|
|
else:
|
|
for peak in peaks:
|
|
ax.plot(peak["FlarePeakTime"], flux[peak["StandardIndex"]], "x", color="red")
|
|
|
|
def plotFlareFits(ax, fits):
|
|
for fit in fits:
|
|
ax.plot(fit["FlareFitTime"], fit["FlareFit"], "g--")
|
|
|
|
def markFlare(ax, fits, flux):
|
|
for fit in fits:
|
|
ax.plot(fit["FlareFitTime"], flux[fit["StandardIndex"]], color="red") |