diff --git a/straxen/plugins/events/__init__.py b/straxen/plugins/events/__init__.py index 0079976a5..44501f43f 100644 --- a/straxen/plugins/events/__init__.py +++ b/straxen/plugins/events/__init__.py @@ -55,6 +55,9 @@ from . import event_s2_position_gcn from .event_s2_position_gcn import * +from . import event_s2_position_lpf +from .event_s2_position_lpf import * + from . import local_minimum_info from .local_minimum_info import * diff --git a/straxen/plugins/events/event_s2_position_lpf.py b/straxen/plugins/events/event_s2_position_lpf.py new file mode 100644 index 000000000..db0a67104 --- /dev/null +++ b/straxen/plugins/events/event_s2_position_lpf.py @@ -0,0 +1,109 @@ +import strax +import straxen +import numpy as np +from warnings import warn +export, __all__ = strax.exporter() + +from straxen.plugins.peaks.peak_positions_lpf import lpf_execute + +@export +class EventS2PositionLPF(strax.Plugin): + """ + LinPosFit pluging for S2 position reconstruction at event level + """ + __version__ = '0.0.0' + depends_on = ('event_area_per_channel', 'event_basics') + provides = "event_s2_position_lpf" + + algorithm = 'lpf' + compressor = 'zstd' + parallel = True # can set to "process" after #82 + + min_reconstruction_area = straxen.URLConfig( + help='Skip reconstruction if area (PE) is less than this', + default=10, infer_type=False, ) + n_top_pmts = straxen.URLConfig( + default=straxen.n_top_pmts, infer_type=False, + help="Number of top PMTs") + pmt_to_lxe = straxen.URLConfig(default=7, infer_type=False, + help="Distance between the PMTs and the liquid interface") + + + def infer_dtype(self): + dtype = [] + for ptype in ['', 'alt_']: + for xy in ['x', 'y']: + dtype += [(f'event_{ptype}s2_{xy}_' + self.algorithm, np.float32, + f'Reconstructed {self.algorithm} {ptype}S2 {xy} position (cm), uncorrected'), + (f'event_{ptype}s2_err_{xy}_' + self.algorithm, np.float32, + f'Error on {self.algorithm} {ptype}S2 {xy} position (cm), 95p confidence interval'), + ] + + dtype += [(f'event_{ptype}s2_lpf_logl',np.float32, f'LinPosFit {ptype}S2 LogLikelihood value'), + (f'event_{ptype}s2_lpf_nuv',np.float32, f'LinPosFit {ptype}S2 parameter n, number of photons generated'), + (f'event_{ptype}s2_lpf_gamma', np.float32, f'LinPosFit {ptype}S2 gamma parameter, reflection'), + (f'event_{ptype}s2_lpf_n_in_fit',np.float32, f'LinPosFit {ptype}S2 number of PMTs included in the fit'), + (f'event_{ptype}s2_lpf_r0',np.float32, f'LinPosFit {ptype}S2 r0 parameter, reduced rate'), + ] + + dtype += strax.time_fields + return dtype + + def setup(self,): + inch = 2.54 # cm + pmt_pos = straxen.pmt_positions() + self.pmt_pos = list(zip(pmt_pos['x'].values,pmt_pos['y'].values,np.repeat(self.pmt_to_lxe, self.n_top_pmts))) + self.pmt_surface=(3*inch)**2*np.pi/4. + + + def compute(self, events): + + result = np.ones(len(events), dtype=self.dtype) + result['time'], result['endtime'] = events['time'], strax.endtime(events) + + for p_type in ['s2', 'alt_s2']: + result[f'event_{p_type}_x_lpf'] *= float('nan') + result[f'event_{p_type}_err_x_lpf'] *= float('nan') + result[f'event_{p_type}_y_lpf'] *= float('nan') + result[f'event_{p_type}_err_y_lpf'] *= float('nan') + result[f'event_{p_type}_lpf_logl'] *= float('nan') + result[f'event_{p_type}_lpf_nuv'] *= float('nan') + result[f'event_{p_type}_lpf_gamma'] *= float('nan') + result[f'event_{p_type}_lpf_n_in_fit'] *= float('nan') + result[f'event_{p_type}_lpf_r0'] *= float('nan') + + + for ix, ev in enumerate(events): + for p_type in ['s2', 'alt_s2']: + if ev[p_type+'_area'] > self.min_reconstruction_area: + + _top_pattern = ev[p_type + '_area_per_channel'][0:self.n_top_pmts] + + try: + + # Execute linearised position reconstruction fit + fit_result, xiter, used_hits = lpf_execute(self.pmt_pos[:self.n_top_pmts], + _top_pattern, + self.pmt_surface) + + result[ix][f'event_{p_type}_x_lpf'] = fit_result[0] + result[ix][f'event_{p_type}_y_lpf'] = fit_result[1] + result[ix][f'event_{p_type}_lpf_nuv'] = fit_result[2] + result[ix][f'event_{p_type}_lpf_logl'] = fit_result[3] + result[ix][f'event_{p_type}_lpf_n_in_fit'] = fit_result[4] + result[ix][f'event_{p_type}_lpf_gamma'] = fit_result[5] + result[ix][f'event_{p_type}_lpf_r0'] = fit_result[6] + result[ix][f'event_{p_type}_err_x_lpf'] = fit_result[7] + result[ix][f'event_{p_type}_err_y_lpf'] = fit_result[8] + + except Exception as e: + # Sometimes inverting the matrix fails.. + # Let's just leave it as a NaN + print(e) + pass + + return result + + + + diff --git a/straxen/plugins/peaks/__init__.py b/straxen/plugins/peaks/__init__.py index 38de82403..94e197ed9 100644 --- a/straxen/plugins/peaks/__init__.py +++ b/straxen/plugins/peaks/__init__.py @@ -22,6 +22,9 @@ from . import peak_positions_mlp from .peak_positions_mlp import * +from . import peak_positions_lpf +from .peak_positions_lpf import * + from . import peak_proximity from .peak_proximity import * diff --git a/straxen/plugins/peaks/peak_positions_lpf.py b/straxen/plugins/peaks/peak_positions_lpf.py new file mode 100644 index 000000000..cd78f2ccd --- /dev/null +++ b/straxen/plugins/peaks/peak_positions_lpf.py @@ -0,0 +1,646 @@ +import strax +import numpy as np +import numba +import straxen +from numba import njit + +export, __all__ = strax.exporter() + +# A.P. Colijn - Feb. 2021 +# Parameters for the fit +lpf_iter_max = 100 +lpf_min_position_cor = 0.01 # convergence criterium for position minimizer +lpf_npar = 4 # number of fit parameters +# x y r0 gamma +lpf_fix_par = np.array([False, False, False, False]) # which parameters to fix + +@export +class PeakPositionsLPF(strax.Plugin): + """ + LPF = Linearized Position Fit + Computes S2 position with linearized log-likelihood fit. + Returns xy position and the parameters of the fit. + See LinPosFit https://github.com/XAMS-nikhef/LinPosFit/ + from A.P. Colijn (Feb 2021) + """ + depends_on =('peaks','peak_basics') + provides = 'peak_positions_lpf' + rechunk_on_save = False + + n_top_pmts = straxen.URLConfig(default=straxen.n_top_pmts, infer_type=False, + help="Number of top PMTs") + + pmt_to_lxe = straxen.URLConfig(default=7, infer_type=False, + help="Distance between the PMTs and the liquid interface") + + def infer_dtype(self): + + dtype = [('x_lpf',np.float32, 'Reconstructed x (cm) position, LinPosFit'), + ('y_lpf',np.float32, 'Reconstructed y (cm) position, LinPosFit'), + ('lpf_logl',np.float32, 'LinPosFit LogLikelihood value'), + ('lpf_nuv',np.float32, 'LinPosFit parameter n, number of photons generated'), + ('lpf_gamma', np.float32, 'LinPosFit gamma parameter, reflection'), + ('lpf_n_in_fit',np.float32, 'LinPosFit number of PMTs included in the fit'), + ('lpf_r0',np.float32, 'LinPosFit r0 parameter, reduced rate'), + ('lpf_err_x',np.float32, 'LinPosFit error on x (cm), 95p conficence interval'), + ('lpf_err_y',np.float32, 'LinPosFit error on y (cm), 95p conficence interval') + ] + + dtype += strax.time_fields + return dtype + + def setup(self,): + inch = 2.54 # cm + pmt_pos = straxen.pmt_positions() + self.pmt_pos = list(zip(pmt_pos['x'].values,pmt_pos['y'].values,np.repeat(self.pmt_to_lxe, self.n_top_pmts))) + self.pmt_surface=(3*inch)**2*np.pi/4. + + def compute(self,peaks): + + result = np.ones(len(peaks), dtype=self.dtype) + + result['time'], result['endtime'] = peaks['time'], strax.endtime(peaks) + result['x_lpf'] *= float('nan') + result['y_lpf'] *= float('nan') + result['lpf_nuv'] *= float('nan') + result['lpf_logl'] *= float('nan') + result['lpf_n_in_fit'] *= float('nan') + result['lpf_gamma'] *= float('nan') + result['lpf_r0'] *= float('nan') + result['lpf_err_x'] *= float('nan') + result['lpf_err_y'] *= float('nan') + + for ix, p in enumerate(peaks): + if p['type']==2: + # Only reconstruct s2 peaks. We do need to set the time of the peaks + try: + # Execute linearised position reconstruction fit + fit_result, xiter, used_hits = lpf_execute(self.pmt_pos[:self.n_top_pmts], + p['area_per_channel'][:self.n_top_pmts], + self.pmt_surface) + + result[ix]['x_lpf'] = fit_result[0] + result[ix]['y_lpf'] = fit_result[1] + result[ix]['lpf_nuv'] = fit_result[2] + result[ix]['lpf_logl'] = fit_result[3] + result[ix]['lpf_n_in_fit'] = fit_result[4] + result[ix]['lpf_gamma'] = fit_result[5] + result[ix]['lpf_r0'] = fit_result[6] + result[ix]['lpf_err_x'] = fit_result[7] + result[ix]['lpf_err_y'] = fit_result[8] + + except Exception as e: + # Sometimes inverting the matrix fails.. + # Let's just leave it as a NaN + pass + + return result + + + +def lpf_execute(xhit, nhit, area_sensor): + #%snakeviz + """ + Execute the fitter + :param xhit: + :param nhit: + :param area_sensor: + :param kwargs: + :return: + """ + xhit = np.array(xhit) + nhit = np.array(nhit) + # minimize the likelihood + fit_result, xiter, hit_is_used = lpf_minimize(xhit, nhit, area_sensor) + + return fit_result, xiter, hit_is_used + +@njit +def lpf_factorial(x): + n = 1 + for i in range(2, x+1): + n *= i + return n + +@njit +def lpf_lnlike(xhit, nhit, xf, nuv, gamma, area_sensor): + r0 = nuv * area_sensor * xhit[0][2] / 4 / np.pi + logl = 0 + + # print('next') + for ih in range(len(nhit)): + nexp = lpf_nexp(xhit[ih], xf, r0, gamma) + # print(xhit[ih], nexp, xf, r0) + + logl = logl + nexp - nhit[ih] * np.log(nexp) + + return logl + +@njit +def lpf_lnlike_plot(xhit, nhit, hit_is_used, xf, r0, gamma): + logl = 0 + + for ih in range(len(nhit)): + nexp = lpf_nexp(xhit[ih], xf, r0, gamma) + d = lpf_dist(xhit[ih],xf) + if hit_is_used[ih] == 1: + logl = logl + nexp - nhit[ih] * np.log(nexp) + + return logl + +@njit +def lpf_lnlike_r(xhit, nhit, xf, r0, gamma): + logl = 0 + + for ih in range(len(nhit)): + nexp = lpf_nexp(xhit[ih], xf, r0, gamma) + logl = logl + nexp - nhit[ih] * np.log(nexp) + + return logl + +@njit +def lpf_initialize(xhit, nhit): + """ + Initialize the lpf_fitter. Estimate the initial position and uv photons + + :param xhit: + :param nhit: + :return: + """ + nmax = -1 + + xhit_max = np.zeros(3) + # sort hits in descending order. helps when looping over the hits + # + indices = np.argsort(-nhit) # descending order + + # find a cluster of hits + # + # 1. we start with the highest hit in the event + # 2. if it is surrrounded by other hits then we found a cluster and we are done + # 3. if not: continue at 1. with the second highest hit in the event + # 4. repeat until we checked all possibilities + # 5. if no cluster was found we use the highest hit as the initial fit position + # + + # only consider hits within distance dist_min + # + dist_min = 25. + # - maximum fraction of the cluster signal allowed to be inside the leading hit + # - below this fraction we assume that we are dealing with a hit cluster + hit_frac_cut = 0.75 + found_cluster = False + + # loop starting from the highest hit + # + xcluster = np.zeros(3) + for i in range(len(indices)-1): + index = indices[i] + xhit_max = xhit[index] + # the seed hit has signal xmax + # + nmax = nhit[index] + + # we have reached the zero-signal PMTs..... no cluster was found, so leave this loop + # + if nmax == 0: + break + + # n_around indicates the amount of signal that is found around the seed hit + # + n_around = 0 + + # check the other hits.... + # + xcluster = nhit[index]*xhit[index] + for j in range(i+1, len(indices)): + jindex = indices[j] + xtest = xhit[jindex] + dist = lpf_dist(xhit_max, xtest) + # consider the hit if it is near the seed hit + # + if dist1000: + # print('WARNING:: result = ',result) + r_new = np.sqrt((xstart[0]-(xfit[0] - result[0]))**2 + (xstart[1] - (xfit[1] - result[1]))**2) + + weight = 0.5*(1 + np.math.erf((r_max-r_new)/1.0)) + # weight = 1 + # update fit result + #if weight<1e-3: + # print('ping ping ping.... weight =',weight) + ###weight = 1.0 + xfit[0] = xfit[0] - result[0]*weight + xfit[1] = xfit[1] - result[1]*weight + + #weight = np.exp(-lpf_dist(xstart,xfit)/25.) + r0 = r0 - result[2]*weight + + gamma = gamma - result[3]*weight + + # and store the intermediate results + # + xiter[lpf_iter + 1][0] = xfit[0] + xiter[lpf_iter + 1][1] = xfit[1] + xiter[lpf_iter + 1][2] = r0 + xiter[lpf_iter + 1][3] = lpf_lnlike_r(xhit, nhit, xfit, r0, gamma) + xiter[lpf_iter + 1][4] = n_in_fit + xiter[lpf_iter + 1][5] = gamma + + + # if (lpf_iter>=5) and (abs(result[0]) < 0.01) and (abs(result[1]) < 0.01) or (r0<0): # if position no longer changes -> terminate loop + if ((lpf_iter>1) and abs(result[0]) < lpf_min_position_cor) and (abs(result[1]) < lpf_min_position_cor) or (r0<0): # if position no longer changes -> terminate loop + break + + # calculate the number of uv photons + # + nuv = 4 * np.pi * r0 / area_sensor / xhit[0][2] + + # store the fit results + # + fit_result = np.zeros(9) + fit_result[0] = xfit[0] + fit_result[1] = xfit[1] + fit_result[2] = nuv + + logl = 0 + for ih in range(len(nhit)): + if nhit[ih] > 0: + logprob = lpf_hit_prob(xhit[ih], nhit[ih], xfit, r0, gamma) + if logprob>log_pmin: + logl = logl + lpf_hit_prob(xhit[ih], nhit[ih], xfit, r0, gamma) + + fit_result[3] = logl # lpf_lnlike_r(xhit, nhit, xfit, r0) + fit_result[4] = n_in_fit + fit_result[5] = gamma + fit_result[6] = r0 + + errors = lpf_deriv_95(xhit, nhit, hit_is_used, xfit, nuv, gamma) + + fit_result[7] = errors[0] + fit_result[8] = errors[1] + + return fit_result, xiter, hit_is_used + + +@njit +def lpf_dist(x0, x1): + d2 = (x0[0] - x1[0]) ** 2 + (x0[1] - x1[1]) ** 2 + (x0[2] - x1[2]) ** 2 + return np.sqrt(d2) + + +@njit +def lpf_nexp(xi, xf, r0, gamma): + """ + Calculate the expected number of p.e. for a lightsensor + + :param xi: sensor position + :param xf: assumed hit position + :param r0: assumed nuumber of UV photons (re-normalized) + :param gamma: assumed constant background + :return nexp: number of expected photons + + A.P. Colijn + """ + delta = lpf_dist(xi, xf) + nexp = r0 / delta ** 3 + gamma + + return nexp + +@njit +def lpf_f(i, xi, ni, xf, r0, gamma): + """ + Calculate the minimizer functions + :param i: + 0=F0 (x) + 1=F1 (y) + 2=F2 (r0) + 3=F3 (gamma) + :param xi: sensor position + :param ni: hits for the sensor + :param xf: assumed fit position + :param r0: assumed number of UV photons (re-normalized) + :param gamma: assumed constant background + :return f: function value + """ + + f = 0 + if i < 2: + f = -3 * (xf[i] - xi[i]) * (lpf_nexp(xi, xf, r0, gamma) - ni) / lpf_dist(xi, xf) ** 2 + elif i == 2: + f = (lpf_nexp(xi, xf, r0, gamma) - ni) / r0 + elif i == 3: + f = 1. - ni / lpf_nexp(xi, xf, r0, gamma) + + return f + +@njit +def lpf_deriv_f(i, j, xi, ni, xf, r0, gamma): + """ + Derivatives of the minimizer functions + + :param i: + 0=F0 + 1=F1 + 2=F2 + :param j: + 0=x + 1=y + z=r0 + :param xi: hit position + :param ni: number of hits + :param xf: fit position + :param r0: number of photons + :return: + """ + + d = lpf_dist(xi, xf) + n0 = lpf_nexp(xi, xf, r0, gamma) + + deriv = 0 + if i == 0: + dx = xf[0] - xi[0] + if j == 0: # dF0/dx + deriv = -3 * (n0 - ni) / d ** 2 + deriv = deriv - 3 * dx * (n0 - ni) * lpf_deriv_dist_min2(0, xi, xf) + deriv = deriv - 3 * dx * lpf_deriv_n(0, xi, xf, r0, gamma) / d ** 2 + elif j == 1: # dF0/dy + deriv = - 3 * dx * (n0 - ni) * lpf_deriv_dist_min2(1, xi, xf) + deriv = deriv - 3 * dx * lpf_deriv_n(1, xi, xf, r0, gamma) / d ** 2 + elif j == 2: # dF0/dr0 + deriv = -3 * dx * lpf_deriv_n(2, xi, xf, r0, gamma) / d ** 2 + elif j == 3: #dF0/dgamma + deriv = -3 * dx * lpf_deriv_n(3, xi, xf, r0, gamma) / d ** 2 + elif i == 1: + dy = xf[1] - xi[1] + if j == 0: # dF1/dx + deriv = - 3 * dy * (n0 - ni) * lpf_deriv_dist_min2(0, xi, xf) + deriv = deriv - 3 * dy * lpf_deriv_n(0, xi, xf, r0, gamma) / d ** 2 + elif j == 1: # dF1/dy + deriv = -3 * (n0 - ni) / d ** 2 + deriv = deriv - 3 * dy * (n0 - ni) * lpf_deriv_dist_min2(1, xi, xf) + deriv = deriv - 3 * dy * lpf_deriv_n(1, xi, xf, r0, gamma) / d ** 2 + elif j == 2: # dF1/dr0 + deriv = -3 * dy * lpf_deriv_n(2, xi, xf, r0, gamma) / d ** 2 + elif j == 3: #dF0/dgamma + deriv = -3 * dy * lpf_deriv_n(3, xi, xf, r0, gamma) / d ** 2 + elif i == 2: + #if j == 0: + # deriv = lpf_deriv_n(0, xi, xf, r0) + #elif j == 1: + # deriv = lpf_deriv_n(1, xi, xf, r0) + #elif j == 2: + # deriv = lpf_deriv_n(2, xi, xf, r0) + if j == 0: + deriv = lpf_deriv_n(0, xi, xf, r0, gamma) / r0 + elif j == 1: + deriv = lpf_deriv_n(1, xi, xf, r0, gamma) / r0 + elif j == 2: + deriv = lpf_deriv_n(2, xi, xf, r0, gamma) / r0 - (n0 - ni) / r0 ** 2 + elif j == 3: + deriv = lpf_deriv_n(3, xi, xf, r0, gamma) / r0 + elif i == 3: + deriv = ni * lpf_deriv_n(j, xi, xf, r0, gamma) / n0 ** 2 + else: + deriv = 0. + + + return deriv + + +@njit +def lpf_deriv_n(i, xi, xf, r0, gamma): + """ + Derivative of n wrt to fit parameters + + :param i: + 0=x + 1=y + 2=r0 + 3=gamma + :param xi: hit position + :param xf: fit position + :param r0: number of photons + :param gamma: constant background term + :return: dn/di + """ + + if i < 2: + deriv = -3 * lpf_nexp(xi, xf, r0, gamma) * (xf[i] - xi[i]) / lpf_dist(xi, xf) ** 2 + elif i == 2: + deriv = lpf_nexp(xi, xf, r0, gamma) / r0 + elif i == 3: + deriv = 1. + else: + deriv = 0. + + return deriv + + +@njit +def lpf_deriv_dist(i, xi, xf): + """ + Derivative of distance wrt fit parameters + + :param i: 0=x + 1=y + 2=r0 + :param xi: hit position + :param xf: fit position + :return: dDist/di + """ + + if i < 2: + deriv = (xf[i] - xi[i]) / lpf_dist(xi, xf) + else: + deriv = 0.0 + + return deriv + + +@njit +def lpf_deriv_dist_min2(i, xi, xf): + """ + Derivative of 1/dist**2 + + :param i: 0=x + 1=y + 2=r0 + :param xi: hit position + :param xf: fit position + :return: d(1/Dist**2)/di + """ + + deriv = 0.0 + if i < 2: + d = lpf_dist(xi, xf) + deriv = -(2 / d ** 3) * lpf_deriv_dist(i, xi, xf) + + return deriv + +@njit +def lpf_deriv_95(xhit, nhit, hit_is_used, xfit, nuv, gamma): + """ + Error estimation for LinPosFit + Bachelor thesis: Ezra de Cleen, Nikhef, June 2022 + """ + + # deriv_x = 0 + # deriv_y = 0 + + # for isensor in np.arange(len(nhit)): + + # xff = np.zeros(3) + # xff[0] = xfit[0] + # xff[1] = xfit[1] + # xff[2] = 0. + + # # m[i][j] = m[i][j] + lpf_deriv_f(i, j, xhit[isensor], nhit[isensor], xfit, r0, gamma) + + # deriv_x = deriv_x + lpf_deriv_f(float(0), float(0), xhit[isensor], nhit[isensor], xff, nuv, gamma) + # deriv_y = deriv_y + lpf_deriv_f(float(1), float(1), xhit[isensor], nhit[isensor], xff, nuv, gamma) + + + # return x_95, y_95 + + + err = np.zeros(2) + + for i in np.arange(2): + for isensor in range(len(nhit)): + err[i] = err[i] + lpf_deriv_f(i, i, xhit[isensor], nhit[isensor], xfit, nuv, gamma) + + err[i] = 1/np.sqrt(err[i]) + + return err +