flaredetector: util: add linear folded fit and add fit type override

main/flaredetector/util.py

@@ -91,6 +91,13 @@ def polynomial(phase, flux, degree):
 def fitPolynomial(phase, flux, degree):
     return sum(p * phase**i for i, p in enumerate(polynomial(phase, flux, degree)))
 
+def linear(t, k, d):
+    return k*t + d
+
+def fitLinear(phase, flux, degree):
+    popt, _ = curve_fit(linear, phase, flux, maxfev=300)
+    return linear(phase, *popt)
+
 def compureRSS(fit, flux):
     residuals = flux - fit
     return np.sum(residuals**2)
@@ -104,7 +111,7 @@ def computeAIC(rss, numParams, numDataPoints):
 def computeBIC(rss, numParams, numDataPoints):
     return numParams * np.log(numDataPoints) + numDataPoints * np.log(rss/numDataPoints)
 
-def getFoldedBestFit(foldedLc):
+def getFoldedBestFit(foldedLc, fitType="sine"):
     flux = foldedLc.flux
     filt = ~np.isnan(flux)
     multi = 1
@@ -119,18 +126,24 @@ def getFoldedBestFit(foldedLc):
     peaks, _ = find_peaks(smoothed_flux, height=np.mean(smoothed_flux))
     #print("Num peaks: ", peaks)
 
-    try:
+    if(fitType == "sine"):
         retFit = fitSingleSine(phase, flux)
-        r2 = 1 - (compureRSS(retFit, flux)/computeTotalSoS(flux))
+    elif(fitType == "poly"):
-        fitType = "sine"
-        #print("R2: ", r2)
-        if(r2 < fitThreshold):
-            raise Exception("Sine fit is suboptimal")
-        #print("Single sine preferred")
-    except:
         retFit = fitPolynomial(phase, flux, polyDegree)
-        fitType = "poly"
+    elif(fitType == "linear"):
-        #print("Polynomial preferred")
+        retFit = fitLinear(phase, flux)
+    else:
+        print("Try to detect best fit")
+        retFit = fitLinear(phase, flux)
+        r2 = 1 - (compureRSS(retFit, flux)/computeTotalSoS(flux))
+        fitType = "linear"
+        if(r2 < fitThreshold):
+            retFit = fitSingleSine(phase, flux)
+            r2 = 1 - (compureRSS(retFit, flux)/computeTotalSoS(flux))
+            fitType = "sine"
+            if(r2 < fitThreshold):
+                retFit = fitPolynomial(phase, flux, polyDegree)
+                fitType = "poly"
 
     retFit *= multi
     return phase, retFit, fitType