219 lines
7.2 KiB
Python
219 lines
7.2 KiB
Python
from collections import deque
|
|
import numpy as np
|
|
from scipy.optimize import curve_fit
|
|
from scipy.special import erf
|
|
from scipy.stats import binned_statistic
|
|
from scipy.signal import hilbert, chirp, find_peaks, savgol_filter
|
|
from scipy import ndimage
|
|
from collections import OrderedDict
|
|
from cam_server.pipeline.data_processing import functions
|
|
from logging import getLogger
|
|
import numpy as np
|
|
|
|
|
|
initialized = False
|
|
p0 = None
|
|
|
|
def initialize(parameters, sig):
|
|
global initialized, buffer
|
|
buffer = deque(maxlen=parameters["dark_buffer_length"])
|
|
buffer.append(np.ones(np.shape(sig)))
|
|
buffer.append(np.ones(np.shape(sig)))
|
|
initialized = True
|
|
|
|
def do_fft(t, tr, lm=False, lx=False, padding=False, bins=None):
|
|
if bins is not None:
|
|
xbin = np.linspace(0,t[-1]+(t[1]-t[0])/2,bins+1)
|
|
tr, _, _ = binned_statistic(t, tr, statistic='mean', bins=xbin, range=None)
|
|
t = np.linspace(0,t[-1],bins)
|
|
if (lm and lx):
|
|
lm = np.argmin(abs(t-lm));
|
|
lx = np.argmin(abs(t-lx))
|
|
tr=tr[lm:lx]
|
|
t=t[lm:lx]
|
|
|
|
if padding:
|
|
N = t.size
|
|
te = t[1]-t[0]
|
|
t = np.arange(N+pad)*te+t[0]
|
|
trp = np.zeros(N+pad)
|
|
trp[:N]=tr
|
|
tr=trp
|
|
|
|
N = t.size
|
|
T = t[N-1]-t[0]
|
|
te = t[1]-t[0]
|
|
fe = 1.0/te
|
|
tfd=np.fft.fft(tr)/N
|
|
freq=np.arange(N)*1.0/T
|
|
tfd = tfd[:int(len(tfd)/2)]
|
|
freq = freq[:int(len(freq)/2)]
|
|
return freq, np.absolute(tfd), np.real(tfd), np.imag(tfd)
|
|
|
|
def gauss(x, x0, fwhm, a):
|
|
return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))
|
|
|
|
def gauss_cos( x, x0, fwhm, a, f, phase):
|
|
return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))*np.cos(2*np.pi*f*(x-x0)-phase)
|
|
|
|
def gauss_sin(x, x0, fwhm, a, f, phase):
|
|
return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))*np.sin(2*np.pi*f*(x-x0)+phase)
|
|
|
|
def fit_data(x, y, bins=None, diag=False, sigma=None, distance=None):
|
|
bounds=[[x[abs(y).argmax()]-0.3, 0.03, 0.005, 15., -np.inf], [x[abs(y).argmax()]+0.3, 2., 10., 35., np.inf]]
|
|
global p0
|
|
if p0 is None or np.any([~(lb<p<hb) for p, lb, hb in zip(p0, bounds[0], bounds[1])]):
|
|
p0 = find_signal(x, y, distance=distance)
|
|
#p0 = [x[abs(y).argmax()], abs(x[-1]-x[0])/4, abs(y).max(), 24, np.pi]
|
|
|
|
# assumes that x is in ps
|
|
if bins is not None:
|
|
xbin = np.linspace(0,x[-1]+(x[1]-x[0])/2,bins+1)
|
|
y, _, _ = binned_statistic(x, y, statistic='mean', bins=xbin, range=None)
|
|
x = np.linspace(0,x[-1],bins)
|
|
popt, pcov = curve_fit(gauss_cos, x, y, p0, sigma=sigma, bounds=bounds)
|
|
#x0, fwhm, a, f, phase = p0
|
|
p0 = popt
|
|
x0, fwhm, a, f, phase = popt
|
|
|
|
if diag:
|
|
return x0, fwhm, a, f, np.remainder(phase, 2*np.pi), p0, popt
|
|
else:
|
|
return x0, fwhm, a, f, np.remainder(phase, 2*np.pi)
|
|
|
|
def fit_extrema(time, y, distance, diag=None):
|
|
xp, _ = find_peaks(abs(y), height=abs(y).max()/5, distance=distance)
|
|
yp = y[xp]
|
|
return fit_data(time[xp], yp, bins=None, diag=diag)
|
|
|
|
def get_max(c, dpx_poly=None, offset=0):
|
|
"""getting maximum from a correlation curve (optionally using polynomial fit)"""
|
|
im = c.argmax()
|
|
mx = c[im]
|
|
|
|
if dpx_poly:
|
|
try:
|
|
poly = np.polyfit(np.arange(-dpx_poly//2,dpx_poly//2),c[im-dpx_poly//2:im+dpx_poly//2], 2)
|
|
root = np.roots(np.polyder(poly,1))
|
|
im+= np.real(root[np.argmin(np.abs(root))])
|
|
except:
|
|
im=im
|
|
return im + offset, mx
|
|
|
|
def find_signal(time, y, distance):
|
|
dt = (time[1]-time[0])
|
|
xp, _ = find_peaks(abs(y), height=abs(y).max()/5, distance=distance)
|
|
yp = abs(y)[xp]
|
|
p0 = [xp[yp.argmax()], (xp[-1]-xp[0])/3, yp.max()]
|
|
popt, pcov = curve_fit(gauss, xp, yp, p0)
|
|
x0, fwhm, a = popt
|
|
xp_sign = xp * np.sign(y[xp])
|
|
# idx_closest = np.argmin(np.abs(xp_sign-x0))
|
|
# f = 1/(dt*np.mean(np.diff(xp_sign[idx_closest-1:idx_closest+1]))) Metric that doesnt make sense but is sensitive to the phase
|
|
xp_pos = xp_sign[xp_sign>0]
|
|
idx_closest = np.argmin(np.abs(xp_pos-x0))
|
|
weights=abs(y)[xp[:-1]]+abs(y)[xp[1:]]
|
|
f = 1/(dt*np.average(np.diff(xp), weights=weights)*2)
|
|
#f = 1/(dt*np.mean(np.diff(xp_pos[idx_closest-2:idx_closest+1]))) works but uses only two points
|
|
#f = 1/(dt*np.mean(np.diff(xp_pos))) <-- rather noisy
|
|
phase = f*(np.min(xp_pos[xp_pos>x0]-x0)*dt)*2*np.pi
|
|
x0 = x0*dt
|
|
fwhm = fwhm*dt
|
|
return x0, fwhm, a, f, phase
|
|
|
|
def normstep(step):
|
|
"""normalizing a test signal for np.correlate"""
|
|
step = step - np.mean(step)
|
|
step = step / np.sum(step**2)
|
|
return step
|
|
|
|
def get_reference_function(sigma_px=30, reflen=150):
|
|
rng = reflen / np.sqrt(2) / sigma_px / 2
|
|
ref = -erf(np.linspace(-rng, rng, reflen)) / 2
|
|
ref = normstep(ref)
|
|
return ref
|
|
|
|
def process(bsdata, pid, timestanp, parameters):
|
|
global initialized
|
|
global p0
|
|
## initialize all values sent to BS stream to None
|
|
x=None
|
|
sig=None
|
|
bg=None
|
|
time=None
|
|
ratio=None
|
|
freq=None
|
|
fft_frequency = None
|
|
fft_amplitude = None
|
|
fft_amplitude_real = None
|
|
fft_amplitude_imag = None
|
|
x0=None
|
|
fwhm=None
|
|
a=None
|
|
f=None
|
|
phase=None
|
|
c=None
|
|
|
|
|
|
## params used for evaluation
|
|
calibration = parameters["calibration"]
|
|
padding = parameters["padding"]
|
|
lm = parameters["time_min"]
|
|
lx = parameters["time_max"]
|
|
|
|
|
|
## get data and events for sorting
|
|
events = bsdata[parameters["events"]]
|
|
sig = bsdata[parameters["sig"]]
|
|
x = np.arange(0,len(sig))
|
|
|
|
if not initialized:
|
|
initialize(parameters, sig)
|
|
|
|
is_midir_dark = parameters["is_midir_dark"]
|
|
is_midir_dark_evt = False
|
|
|
|
# add to dark reference
|
|
if is_midir_dark:
|
|
is_midir_dark_evt = events[parameters["event_code_midir_dark"]]
|
|
if is_midir_dark_evt:
|
|
buffer.append(sig)
|
|
# or analyze
|
|
bg = np.mean(buffer, axis=0)
|
|
if not is_midir_dark_evt:
|
|
ratio = sig/bg
|
|
time = np.polyval(calibration, x)
|
|
#fft_frequency, fft_amplitude, fft_amplitude_real, fft_amplitude_imag = do_fft(time, ratio-1, lm=False, lx=False, padding=False, bins=parameters["fft_bins"])
|
|
ref=get_reference_function(parameters["sigma_px"], parameters["reflen"])
|
|
c = np.correlate(ratio-1, ref, "same")
|
|
if parameters["fit_data"]:
|
|
x0, fwhm, a, f, phase = fit_data(time, c, bins=parameters["bins"], distance=parameters["distance"])
|
|
else:
|
|
x0, fwhm, a, f, phase = find_signal(time, c, distance=parameters["distance"])
|
|
|
|
data = {
|
|
"SARES20-CEP02:spectrometer_x": x,
|
|
"SARES20-CEP02:spectrometer_signal": sig,
|
|
"SARES20-CEP02:spectrometer_background": bg,
|
|
"SARES20-CEP02:spectrometer_time": time,
|
|
"SARES20-CEP02:spectrometer_ratio": ratio,
|
|
"SARES20-CEP02:spectrometer_correlation": c,
|
|
"SARES20-CEP02:fft_frequency": fft_frequency,
|
|
"SARES20-CEP02:fft_amplitude_abs": fft_amplitude,
|
|
"SARES20-CEP02:fft_amplitude_real": fft_amplitude_real,
|
|
"SARES20-CEP02:fft_amplitude_imag": fft_amplitude_imag,
|
|
"SARES20-CEP02:fit_x0": x0,
|
|
"SARES20-CEP02:fit_fwhm": fwhm,
|
|
"SARES20-CEP02:fit_amplitude": a,
|
|
"SARES20-CEP01:fit_frequency": f,
|
|
"SARES20-CEP01:fit_phase": phase,
|
|
}
|
|
|
|
#data.update({"p_est": pinit, "popt": popt})
|
|
return data
|
|
|
|
|
|
|
|
|
|
|