astrodatagui: FlaredetectorWidget: plot fits/mark flares in all possible plots

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