Files
flaredetector/main/flaredetector/util.py
T

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")