diff --git a/main/flaredetector/util.py b/main/flaredetector/util.py index e557877..5c82c5c 100644 --- a/main/flaredetector/util.py +++ b/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) - sineFit = sine(foldedLc.phase[filt].value, *popt) + flux = flux[filt] + #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