flaredetector: util: add multiple fitting functions

sine, double sine (2 sine curves in 1 lightcurve one after the other),
a 10 degree polynomial

main/flaredetector/util.py

@@ -1,4 +1,5 @@
 import numpy as np
+from numpy.polynomial.polynomial import Polynomial
 from scipy.signal import find_peaks
 from scipy.optimize import curve_fit
 
@@ -82,8 +83,47 @@ def convertStarndardIndexToFoldedIndex(foldedLc, standardIndex):
 
     return cycle, foldedIndex
 
+def singleSine(t, A, w, p, c):
+    return A * np.sin(w*t + p) + c
+
+def doubleSine(t, A1, w1, p1, c1, A2, w2, p2, c2, split1, split2):
+    y = np.zeros_like(t)
+    m1 = t < split1
+    m2 = (t >= split1) & (t < split2)
+    m3 = t >= split2
+    y[m1] = A1 * np.sin(w1 * t[m1] + p1) + c1
+    y[m2] = A2 * np.sin(w2 * t[m2] + p2) + (A1 * np.sin(w1 * split1 + p1) + c1 - A2 * np.sin(w2 * split1 + p2))
+    y[m3] = A1 * np.sin(w1 * t[m3] + p1) + (A2 * np.sin(w2 * split2 + p2) + (A1 * np.sin(w1 * split1 + p1) + c1 - A2 * np.sin(w2 * split1 + p2)) - A1 * np.sin(w1 * split2 + p1))
+    return y
+
+def fitSingleSine(phase, flux):
+    initGuess = [np.ptp(flux)/2, 2 * np.pi / (np.max(phase) - np.min(phase)), 0, np.mean(flux)]
+    popt, _ = curve_fit(singleSine, phase, flux, p0=initGuess, maxfev=100000)
+    return singleSine(phase, *popt)
+
+def fitDoubleSine(phase, flux):
+    initSplit1 = abs(phase[0]) * 2 / 3 + phase[0]
+    initSplit2 = abs(phase[0]) * 4 / 3 + phase[0]
+    initGuess = [np.ptp(flux) / 2, 2 * np.pi / (np.max(phase) - np.min(phase)), 0, np.mean(flux),
+                 np.ptp(flux) / 2, 2 * np.pi / (np.max(phase) - np.min(phase)), 0, np.mean(flux),
+                 initSplit1, initSplit2]
+    popt, _ = curve_fit(doubleSine, phase, flux, p0=initGuess, maxfev=100000)
+    return doubleSine(phase, *popt)
+
+def fitPolynomial(phase, flux, degree):
+    return sum(p * phase**i for i, p in enumerate(Polynomial.fit(phase, flux, degree).convert().coef))
+
+def compureRSS(fit, flux):
+    residuals = flux - fit
+    return np.sum(residuals**2)
+
+def computeAIC(rss, numParams, numDataPoints):
+    return 2 * numParams + numDataPoints * np.log(rss/numDataPoints)
+
+def computeBIC(rss, numParams, numDataPoints):
+    return numParams * np.log(numDataPoints) + numDataPoints * np.log(rss/numDataPoints)
+
 def getFoldedBestFit(foldedLc):
-    def sine(t, A, w, p, c):  return A * np.sin(w*t + p) + c
     flux = foldedLc.flux
     filt = ~np.isnan(flux)
     multi = 1
@@ -91,8 +131,43 @@ def getFoldedBestFit(foldedLc):
         multi = float(np.mean(flux).value)
         flux = foldedLc.normalize().flux
     phase = foldedLc.phase[filt].value
-    popt, pcov = curve_fit(sine, phase, flux[filt], maxfev=100000)
+    flux = flux[filt]
-    sineFit = sine(foldedLc.phase[filt].value, *popt)
+    #popt, pcov = curve_fit(sine, phase, flux[filt], maxfev=100000)
+    import time
+    start = time.time()
+    singleFit = fitSingleSine(phase, flux)
+    singleTime = time.time() - start
+    start = time.time()
+    doubleFit = fitDoubleSine(phase, flux)
+    doubleTime = time.time() - start
+    start = time.time()
+    polyFit = fitPolynomial(phase, flux, 10)
+    polyTime = time.time() - start
+
+    singleRSS = compureRSS(singleFit, flux)
+    doubleRSS = compureRSS(doubleFit, flux)
+    polyRSS = compureRSS(polyFit, flux)
+
+    singleAIC = computeAIC(singleRSS, 4, len(flux))
+    singleBIC = computeBIC(singleRSS, 4, len(flux))
+    doubleAIC = computeAIC(doubleRSS, 10, len(flux))
+    doubleBIC = computeBIC(doubleRSS, 10, len(flux))
+    polyAIC = computeAIC(polyRSS, 10, len(flux))
+    polyBIC = computeBIC(polyRSS, 10, len(flux))
+
+    if (singleAIC < doubleAIC and singleAIC < polyAIC and singleBIC < doubleBIC and singleBIC < polyBIC):
+        print("Single sine preferred")
+        sineFit = singleFit
+    elif (doubleAIC < singleAIC and doubleAIC < polyAIC and doubleBIC < singleBIC and doubleBIC < polyBIC):
+        print("Double sine preferred")
+        sineFit = doubleFit
+    else:
+        print("Polynomial preferred")
+        sineFit = polyFit
+
+    print("Single sine time: ", singleTime)
+    print("Double sine time: ", doubleTime)
+    print("Polynomial time: ", polyTime)
     sineFit *= multi
     return phase, sineFit