diff --git a/main/flaredetector/util.py b/main/flaredetector/util.py index 06e6880..3c79598 100644 --- a/main/flaredetector/util.py +++ b/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)) - fitType = "sine" - #print("R2: ", r2) - if(r2 < fitThreshold): - raise Exception("Sine fit is suboptimal") - #print("Single sine preferred") - except: + elif(fitType == "poly"): retFit = fitPolynomial(phase, flux, polyDegree) - fitType = "poly" - #print("Polynomial preferred") + elif(fitType == "linear"): + 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