CalcAllFlaresThread: generate plots when calculating data

main/astrodatagui/CalcAllFlaresThread.py

@@ -1,7 +1,9 @@
 from PyQt5.QtCore import pyqtSignal, QThread
 import multiprocessing
 import concurrent.futures
-from ..util.MinimalLightCurve import read
+#from ..util.MinimalLightCurve import read
+import lightkurve as lk
+import matplotlib.pyplot as plt
 import pandas as pd
 from ..flaredetector.flaredetector import calculateFlareFitsForLightcurve
 from ..flaredetector.util import *
@@ -9,12 +11,39 @@ from ..flaredetector.util import *
 import warnings
 warnings.filterwarnings("ignore")
 
+from errno import EEXIST
+from os import makedirs, path
+from datetime import datetime
+
+def mkdir_p(mypath):
+    '''Creates a directory. equivalent to using mkdir -p on the command line'''
+
+    try:
+        makedirs(mypath)
+    except OSError as exc: # Python >2.5
+        if exc.errno == EEXIST and path.isdir(mypath):
+            pass
+        else: raise
+
+current = datetime.now()
+date = f"{current.year}-{current.month}-{current.day}"
+time = f"{current.hour}-{current.minute}-{current.second}"
+folderPath = f"../{date}/"
+mkdir_p(folderPath)
+
 def getFlareCount(filesDict):
     try:
         print(f"Starting {filesDict['StarName']}, {filesDict['Sequence']}")
-        lc = read(filesDict["FilePath"])
+        starName = filesDict['StarName']
+        starNameR = filesDict['StarName'].replace("*", "_star_")
+        sequence = filesDict['Sequence']
+        starFolder = f"{folderPath}/stars/{starNameR}/"
+        mkdir_p(starFolder)
+
+        lc = lk.read(filesDict["FilePath"])
         lc.flux = lc["sap_flux"]
         lc.flux_err = lc["sap_flux_err"]
+        lc = lc.normalize()
         sapPeaks, sapFits = calculateFlareFitsForLightcurve(lc.flatten())
         sapValSec, sapTds = getTotalValidDataInSeconds(lc, "sap_flux")
         sapPeriodogram = lc.to_periodogram()
@@ -33,10 +62,51 @@ def getFlareCount(filesDict):
 
         lc.flux = lc["pdcsap_flux"]
         lc.flux_err = lc["pdcsap_flux_err"]
-        pdcsapPeaks, pdcsapFits = calculateFlareFitsForLightcurve(lc.flatten())
+        lc = lc.normalize()
+
+        lc.plot()
+        plt.title(f"{starName} - normalized lightcurve")
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-lc.png")
+        plt.close()
+
+        flattenedLc = lc.flatten()
+        flattenedLc.plot()
+        plt.title(f"{starName} - flattened lightcurve")
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-flattened_lc.png")
+        plt.close()
+
+        pdcsapPeaks, pdcsapFits = calculateFlareFitsForLightcurve(flattenedLc, normalizedLC=lc)
+        lc.plot()
+        plt.title(f"{starName} - normalized lightcurve")
+        for p in pdcsapPeaks:
+            plt.plot(p["FlarePeakTime"].value, lc.flux[p["StandardIndex"]], "x", color="red")
+        plt.plot([], [], "x", color="red", label="Flare peaks")
+        plt.legend()
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-lc-marked_flares.png")
+        plt.close()
+
+        flattenedLc.plot()
+        plt.title(f"{starName} - flattened lightcurve")
+        for f in pdcsapFits:
+            plt.plot(f["FlareFitTime"].value, f["FlareFit"], "g--")
+        plt.plot([], [], "g--", label="Flare fits")
+        for p in pdcsapPeaks:
+            plt.plot(p["FlarePeakTime"].value, p["FlarePeak"], "x", color="red")
+        plt.plot([], [], "x", color="red", label="Flare peaks")
+        plt.legend()
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-flattened_lc-marked_flares.png")
+        plt.close()
+
         pdcsapValSec, pdcsapTds = getTotalValidDataInSeconds(lc, "pdcsap_flux")
         pdcsapPeriodogram = lc.to_periodogram()
-        pdcsapPeakPeriod = pdcsapPeriodogram.period[findMaxIndices(pdcsapPeriodogram, num=4, distance=100, sortByHighest=True)[0]]
+        maxPeriodIndex = findMaxIndices(pdcsapPeriodogram, num=4, distance=100, sortByHighest=True)[0]
+        pdcsapPeakPeriod = pdcsapPeriodogram.period[maxPeriodIndex]
+
+        pdcsapPeriodogram.plot(view="period")
+        plt.plot(pdcsapPeakPeriod, pdcsapPeriodogram.power[maxPeriodIndex], "x", color="red")
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-periodogram-marked_max.png")
+        plt.close()
+
         pdcsapEpochTime = getEpochTime(lc)
         pdcsapFoldedLC = lc.fold(period=pdcsapPeakPeriod, epoch_time=pdcsapEpochTime)
         pdcsapPhase, pdcsapSineFit, pdcsapFitType = getFoldedBestFit(pdcsapFoldedLC, fitType=filesDict["FitType"])
@@ -44,10 +114,21 @@ def getFlareCount(filesDict):
         pdcsapminPhasesBounds, pdcsapmaxPhasesBounds = getPhaseRangesNearPeak((pdcsapMinima, pdcsapMaxima), pdcsapPhase, returnPhaseValue=True)
         pdcsapFoldedPeaks = []
         pdcsapFoldedPeaksPhasePair = []
+
+        pdcsapFoldedLC.scatter()
+        plt.title(f"{starName} - folded lightcurve")
+        plt.plot(pdcsapPhase, pdcsapSineFit, color="blue", label=f"{pdcsapFitType}-fit")
         for peak in pdcsapPeaks:
             cycle, foldedIndex = convertStarndardIndexToFoldedIndex(pdcsapFoldedLC, peak["StandardIndex"])
             pdcsapFoldedPeaks.append({"Cycle: ": cycle, "Index": foldedIndex})
             pdcsapFoldedPeaksPhasePair.append({"Phase": pdcsapFoldedLC.phase[pdcsapFoldedLC.cycle == cycle][foldedIndex], "Peak": peak})
+            plt.plot(pdcsapFoldedLC.phase[pdcsapFoldedLC.cycle == cycle][foldedIndex].value,
+                     pdcsapFoldedLC.flux[pdcsapFoldedLC.cycle == cycle][foldedIndex], "x", color="red")
+        plt.plot([], [], "x", color="red", label="Flare peaks")
+        plt.legend()
+        plt.savefig(f"{starFolder}/{starNameR}_{sequence}-foldedLC-marked_fit_flares.png")
+        plt.close()
+        
 
         filesDict["sapPeaks"] = sapPeaks
         filesDict["sapPeaksCount"] = len(sapPeaks)
@@ -79,6 +160,11 @@ def getFlareCount(filesDict):
         filesDict["pdcsapPeriodMinimaBoundaries"] = pdcsapminPhasesBounds
         filesDict["pdcsapPeriodMaxima"] = pdcsapMaxima
         filesDict["pdcsapPeriodMaximaBoundaries"] = pdcsapmaxPhasesBounds
+
+        csvFile = open(f"{starFolder}/{starNameR}_{sequence}.csv", "a")
+        csvFile.write("StarName,Spectral Type,Rotational Velocity,Rotenional Velocity Unit,Distance,Distance Unit,Source,Sequence,File Path,Initial folded Fit Type,Used folded Fit Type")
+        csvFile.write(f"{filesDict['StarName']},{filesDict['SpType']},{filesDict['RotVel']},{filesDict['RotVelUnit']},{filesDict['Distance']},{filesDict['DistanceUnit']},{filesDict['Source']},{filesDict['Sequence']},{filesDict['FilePath']},{filesDict['FitType']},{pdcsapFitType}")
+        csvFile.close()
         del lc
         print(f"Finished {filesDict['StarName']}, {filesDict['Sequence']}")
         return filesDict