235 lines
8.4 KiB
Python
235 lines
8.4 KiB
Python
import numpy as np
|
|
from numpy.polynomial.polynomial import Polynomial
|
|
from scipy.signal import find_peaks, argrelextrema
|
|
from scipy.optimize import curve_fit
|
|
|
|
def findMaxIndices(lc, num=3, distance=100, height=(None, None), sortByHighest=False):
|
|
if(hasattr(lc, "power")):
|
|
data = lc.power
|
|
elif(hasattr(lc, "flux")):
|
|
data = lc.flux
|
|
else:
|
|
return np.zeros(num)-1
|
|
peak_indices, peak_dict = find_peaks(data, height=height,
|
|
distance=distance)
|
|
peak_heights = peak_dict["peak_heights"]
|
|
if(sortByHighest):
|
|
return peak_indices[np.argsort(peak_heights)[-num:][::-1]]
|
|
else:
|
|
return peak_indices[0:num]
|
|
|
|
def findNonNanTimeDeltas(lc, column):
|
|
timedeltas = []
|
|
nonNanIndices = np.where(~np.isnan(lc[column]))[0]
|
|
|
|
if len(nonNanIndices) == 0:
|
|
return timedeltas
|
|
|
|
start = nonNanIndices[0]
|
|
end = start
|
|
|
|
for i in range(1, len(nonNanIndices)):
|
|
if nonNanIndices[i] == nonNanIndices[i-1] + 1:
|
|
end = nonNanIndices[i]
|
|
else:
|
|
timedeltas.append((lc.time[start], lc.time[end]))
|
|
start = nonNanIndices[i]
|
|
end = start
|
|
|
|
timedeltas.append((lc.time[start], lc.time[end]))
|
|
|
|
return [timedelta[1] - timedelta[0] for timedelta in timedeltas]
|
|
|
|
def getTimeDeltaInSeconds(timedelta):
|
|
return int(timedelta.to_value("sec"))
|
|
|
|
def getTotalValidDataInSeconds(lc, column):
|
|
tds = findNonNanTimeDeltas(lc, column)
|
|
validDataTime = 0
|
|
for td in tds:
|
|
validDataTime += getTimeDeltaInSeconds(td)
|
|
return validDataTime, tds
|
|
|
|
def plotFlarePeaks(ax, peaks, flux = None):
|
|
if(flux is None):
|
|
for peak in peaks:
|
|
ax.plot(peak["FlarePeakTime"], peak["FlarePeak"], "x", color="red")
|
|
else:
|
|
for peak in peaks:
|
|
ax.plot(peak["FlarePeakTime"], flux[peak["StandardIndex"]], "x", color="red")
|
|
|
|
def plotFlareFits(ax, fits):
|
|
for fit in fits:
|
|
ax.plot(fit["FlareFitTime"], fit["FlareFit"], "g--")
|
|
|
|
def markFlare(ax, fits, flux):
|
|
for fit in fits:
|
|
ax.plot(fit["FlareFitTime"], flux[fit["StandardIndex"]], color="red")
|
|
|
|
def convertStarndardIndexToFoldedIndex(foldedLc, standardIndex):
|
|
foldedIndex = standardIndex; cycle = 0
|
|
for i in range(max(foldedLc.cycle)+1):
|
|
if(foldedIndex - len(foldedLc.phase[foldedLc.cycle == i]) >= 0):
|
|
foldedIndex = foldedIndex - len(foldedLc.phase[foldedLc.cycle == i])
|
|
else:
|
|
cycle = i
|
|
break
|
|
|
|
return cycle, foldedIndex
|
|
|
|
def singleSine(t, A, w, p, c):
|
|
return A * np.sin(w*t + p) + c
|
|
|
|
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=300)
|
|
return singleSine(phase, *popt)
|
|
|
|
def polynomial(phase, flux, degree):
|
|
return Polynomial.fit(phase, flux, degree).convert().coef
|
|
|
|
def fitPolynomial(phase, flux, degree):
|
|
return sum(p * phase**i for i, p in enumerate(polynomial(phase, flux, degree)))
|
|
|
|
def compureRSS(fit, flux):
|
|
residuals = flux - fit
|
|
return np.sum(residuals**2)
|
|
|
|
def computeTotalSoS(flux):
|
|
return np.sum((flux - np.mean(flux))**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):
|
|
flux = foldedLc.flux
|
|
filt = ~np.isnan(flux)
|
|
multi = 1
|
|
if(np.mean(flux).value > 10):
|
|
multi = float(np.mean(flux).value)
|
|
flux = foldedLc.normalize().flux
|
|
phase = foldedLc.phase[filt].value
|
|
flux = flux[filt]
|
|
polyDegree = 7
|
|
fitThreshold = 0.0
|
|
smoothed_flux = np.convolve(flux, np.ones(len(flux)//100)/(len(flux)//100), mode="valid")
|
|
peaks, _ = find_peaks(smoothed_flux, height=np.mean(smoothed_flux))
|
|
#print("Num peaks: ", peaks)
|
|
|
|
try:
|
|
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:
|
|
retFit = fitPolynomial(phase, flux, polyDegree)
|
|
fitType = "poly"
|
|
#print("Polynomial preferred")
|
|
|
|
retFit *= multi
|
|
return phase, retFit, fitType
|
|
|
|
def getFoldedFitPeakValley(sineFit):
|
|
minima = argrelextrema(sineFit, np.less)[0]
|
|
maxima = argrelextrema(sineFit, np.greater)[0]
|
|
#print(minima)
|
|
#print(maxima)
|
|
if(len(minima) > 0 and len(maxima) > 0):
|
|
if(minima[0] < maxima[0] and minima[-1] > maxima[-1]):
|
|
if(sineFit[minima[0]] > sineFit[minima[-1]]):
|
|
minima = np.delete(minima, 0)
|
|
else:
|
|
minima = np.delete(minima, -1)
|
|
elif(maxima[0] < minima[0] and maxima[-1] > minima[-1]):
|
|
if(sineFit[maxima[0]] > sineFit[maxima[-1]]):
|
|
maxima = np.delete(maxima, -1)
|
|
else:
|
|
maxima = np.delete(maxima, 0)
|
|
|
|
return minima, maxima
|
|
|
|
def findNearestIndexOfValue(array, value):
|
|
array = np.asarray(array)
|
|
idx = (np.abs(array - value)).argmin()
|
|
return idx
|
|
|
|
def getPhaseRangesNearPeak(maxArgs, phase):
|
|
minPhases = phase[maxArgs[0]]
|
|
maxPhases = phase[maxArgs[1]]
|
|
totalPhase = abs(phase[0]) + abs(phase[-1])
|
|
|
|
phasePart = totalPhase * 0.3 / (len(minPhases) + len(maxPhases))
|
|
minPhasesBounds = []
|
|
for minPhase in minPhases:
|
|
minPhaseBounds = []
|
|
if(minPhase - phasePart < phase[0]):
|
|
minPhaseBounds.append((phase[0], minPhase + phasePart))
|
|
minPhaseBounds.append((phase[-1] + (minPhase - phasePart - phase[0]), phase[-1]))
|
|
elif(minPhase + phasePart > phase[-1]):
|
|
minPhaseBounds.append((minPhase - phasePart, phase[-1]))
|
|
minPhaseBounds.append((phase[0], phase[0] + (minPhase + phasePart - phase[-1])))
|
|
else:
|
|
minPhaseBounds.append((minPhase - phasePart, minPhase + phasePart))
|
|
minPhasesBounds.append(minPhaseBounds)
|
|
|
|
maxPhasesBounds = []
|
|
for maxPhase in maxPhases:
|
|
maxPhaseBounds = []
|
|
if(maxPhase - phasePart < phase[0]):
|
|
maxPhaseBounds.append((phase[0], maxPhase + phasePart))
|
|
maxPhaseBounds.append((phase[-1] + (maxPhase - phasePart - phase[0]), phase[-1]))
|
|
elif(maxPhase + phasePart > phase[-1]):
|
|
maxPhaseBounds.append((maxPhase - phasePart, phase[-1]))
|
|
maxPhaseBounds.append((phase[0], phase[0] + (maxPhase + phasePart - phase[-1])))
|
|
else:
|
|
maxPhaseBounds.append((maxPhase - phasePart, maxPhase + phasePart))
|
|
maxPhasesBounds.append(maxPhaseBounds)
|
|
|
|
minPhasesBoundsIndices = []
|
|
for minPhaseBounds in minPhasesBounds:
|
|
minPhaseBoundsIndices = []
|
|
for l in minPhaseBounds:
|
|
minPhaseBoundsIndices.append([findNearestIndexOfValue(phase, lv) for lv in l])
|
|
minPhasesBoundsIndices.append(minPhaseBoundsIndices)
|
|
|
|
maxPhasesBoundsIndices = []
|
|
for maxPhaseBounds in maxPhasesBounds:
|
|
maxPhaseBoundsIndices = []
|
|
for l in maxPhaseBounds:
|
|
maxPhaseBoundsIndices.append([findNearestIndexOfValue(phase, lv) for lv in l])
|
|
maxPhasesBoundsIndices.append(maxPhaseBoundsIndices)
|
|
|
|
return minPhasesBoundsIndices, maxPhasesBoundsIndices
|
|
|
|
def plotPhaseRangesNearPeak(indices, phase, ax=None):
|
|
minPhasesBoundsIndices = indices[0]
|
|
maxPhasesBoundsIndices = indices[1]
|
|
|
|
if(ax is None):
|
|
import matplotlib.pyplot as plt
|
|
for minPhaseBoundsIndices in minPhasesBoundsIndices:
|
|
for pair in minPhaseBoundsIndices:
|
|
for l in pair:
|
|
plt.axes.axvline(phase[l], color="violet")
|
|
for maxPhaseBoundsIndices in maxPhasesBoundsIndices:
|
|
for pair in maxPhaseBoundsIndices:
|
|
for l in pair:
|
|
plt.axes.axvline(phase[l], color="orange")
|
|
else:
|
|
for minPhaseBoundsIndices in minPhasesBoundsIndices:
|
|
for pair in minPhaseBoundsIndices:
|
|
for l in pair:
|
|
ax.axvline(phase[l], color="violet")
|
|
for maxPhaseBoundsIndices in maxPhasesBoundsIndices:
|
|
for pair in maxPhaseBoundsIndices:
|
|
for l in pair:
|
|
ax.axvline(phase[l], color="orange")
|
|
|
|
def getEpochTime(lc):
|
|
return lc.time.value[argrelextrema(np.asarray(lc.flux.value), np.less, order=500)[0]][0] |