diff --git a/main/astrodatagui/ui/FlaredetectorWidget.py b/main/astrodatagui/ui/FlaredetectorWidget.py index f66713a..70c2341 100644 --- a/main/astrodatagui/ui/FlaredetectorWidget.py +++ b/main/astrodatagui/ui/FlaredetectorWidget.py @@ -7,7 +7,9 @@ from astropy.io import fits from astropy.visualization import time_support import lightkurve as lk -from ...flaredetector.util import findMaxIndices +from ...flaredetector.util import (findMaxIndices, plotFlarePeaks, + plotFlareFits, markFlare) +from ...flaredetector.flaredetector import calculateFlareFitsForLightcurve class FlaredetectorWidget(QtWidgets.QWidget): def __init__(self, name): @@ -168,6 +170,7 @@ class FlaredetectorWidget(QtWidgets.QWidget): lc = self.currentLC label = f"{self.currentMainName}" lc.flux = lc[self.fluxType] + if(self.NormalizeState["Enabled"]): if(self.currentLCCollection is None): print("normalizing normal lc") @@ -208,7 +211,16 @@ class FlaredetectorWidget(QtWidgets.QWidget): self.periodogramInfoGroupBox.setVisible(self.PeriodogramState["Enabled"]) + if(not self.FlattenState["Enabled"] and + not self.FoldState["Enabled"] and + not self.PeriodogramState["Enabled"]): + peaks, fits = calculateFlareFitsForLightcurve(lc.flatten(window_length=self.FlattenState["WindowLength"], + polyorder=self.FlattenState["PolynomialOrder"]), num=5) + elif(self.FlattenState["Enabled"]): + peaks, fits = calculateFlareFitsForLightcurve(lc, num=5) + self.figureAxis.clear() + time_support() if(self.FoldState["Enabled"]): lc.scatter(label=label, ax=self.figureAxis) elif(self.PeriodogramState["Enabled"]): @@ -217,10 +229,10 @@ class FlaredetectorWidget(QtWidgets.QWidget): self.figureAxis.plot(lc.period[maxPowers], lc.power[maxPowers], "x") if(self.PeriodogramState["View"] == "frequency"): self.figureAxis.plot(lc.frequency[maxPowers], lc.power[maxPowers], "x") - elif(self.FlattenState["Enabled"]): - time_support() - lc.plot(label=label, ax=self.figureAxis) - self.figureAxis.plot(lc.time[maxPowers], lc.flux[maxPowers], "x") else: lc.plot(label=label, ax=self.figureAxis) + plotFlarePeaks(self.figureAxis, peaks, lc.flux) + markFlare(self.figureAxis, fits, lc.flux) + if(self.FlattenState["Enabled"]): + plotFlareFits(self.figureAxis, fits) self.figure.canvas.draw_idle() \ No newline at end of file