diff options
Diffstat (limited to 'utils/sddm/plot_energy.py')
| -rwxr-xr-x | utils/sddm/plot_energy.py | 352 |
1 files changed, 352 insertions, 0 deletions
diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py new file mode 100755 index 0000000..a24ea00 --- /dev/null +++ b/utils/sddm/plot_energy.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python +# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago> +# +# This program is free software: you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the Free +# Software Foundation, either version 3 of the License, or (at your option) +# any later version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for +# more details. +# +# You should have received a copy of the GNU General Public License along with +# this program. If not, see <https://www.gnu.org/licenses/>. + +from __future__ import print_function, division +import numpy as np +from scipy.stats import iqr, poisson +from matplotlib.lines import Line2D +from scipy.stats import iqr, norm, beta +from scipy.special import spence +from itertools import izip_longest +from .dc import DC_MUON, DC_JUNK, DC_CRATE_ISOTROPY, DC_QVNHIT, DC_NECK, DC_FLASHER, DC_ESUM, DC_OWL, DC_OWL_TRIGGER, DC_FTS, DC_ITC, DC_BREAKDOWN +from . import grouper, print_warning, AV_RADIUS, PSUP_RADIUS +import pandas as pd + +# Fermi constant +GF = 1.16637887e-5 # 1/MeV^2 +ELECTRON_MASS = 0.5109989461 # MeV +MUON_MASS = 105.6583745 # MeV +PROTON_MASS = 938.272081 # MeV +FINE_STRUCTURE_CONSTANT = 7.297352566417e-3 + +def unwrap(p, delta, axis=-1): + """ + A modified version of np.unwrap() useful for unwrapping the 50 MHz clock. + It unwraps discontinuities bigger than delta/2 by delta. + + Example: + + >>> a = np.arange(10) % 5 + >>> a + array([0, 1, 2, 3, 4, 0, 1, 2, 3, 4]) + >>> unwrap(a,5) + array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + + In the case of the 50 MHz clock delta should be 0x7ffffffffff*20.0. + """ + p = np.asarray(p) + nd = p.ndim + dd = np.diff(p, axis=axis) + slice1 = [slice(None, None)]*nd # full slices + slice1[axis] = slice(1, None) + slice1 = tuple(slice1) + ddmod = np.mod(dd + delta/2, delta) - delta/2 + np.copyto(ddmod, delta/2, where=(ddmod == -delta/2) & (dd > 0)) + ph_correct = ddmod - dd + np.copyto(ph_correct, 0, where=abs(dd) < delta/2) + up = np.array(p, copy=True, dtype='d') + up[slice1] = p[slice1] + ph_correct.cumsum(axis) + return up + +def unwrap_50_mhz_clock(gtr): + """ + Unwrap an array with 50 MHz clock times. These times should all be in + nanoseconds and come from the KEV_GTR variable in the EV bank. + + Note: We assume here that the events are already ordered contiguously by + GTID, so you shouldn't pass an array with multiple runs! + """ + return unwrap(gtr,0x7ffffffffff*20.0) + +def retrigger_cut(ev): + """ + Cuts all retrigger events. + """ + return ev[ev.dt > 500] + +def breakdown_follower_cut(ev): + """ + Cuts all events within 1 second of breakdown events. + """ + breakdowns = ev[ev.dc & DC_BREAKDOWN != 0] + return ev[~np.any((ev.gtr.values > breakdowns.gtr.values[:,np.newaxis]) & \ + (ev.gtr.values < breakdowns.gtr.values[:,np.newaxis] + 1e9),axis=0)] + +def flasher_follower_cut(ev): + """ + Cuts all events within 200 microseconds of flasher events. + """ + flashers = ev[ev.dc & DC_FLASHER != 0] + return ev[~np.any((ev.gtr.values > flashers.gtr.values[:,np.newaxis]) & \ + (ev.gtr.values < flashers.gtr.values[:,np.newaxis] + 200e3),axis=0)] + +def muon_follower_cut(ev): + """ + Cuts all events 200 microseconds after a muon. + """ + muons = ev[ev.dc & DC_MUON != 0] + return ev[~np.any((ev.gtr.values > muons.gtr.values[:,np.newaxis]) & \ + (ev.gtr.values < muons.gtr.values[:,np.newaxis] + 200e3),axis=0)] + +def michel_cut(ev): + """ + Looks for Michel electrons after muons. + """ + prompt_plus_muons = ev[ev.prompt | ((ev.dc & DC_MUON) != 0)] + + # Michel electrons and neutrons can be any event which is not a prompt + # event + follower = ev[~ev.prompt] + + # require Michel events to pass more of the SNO data cleaning cuts + michel = follower[follower.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK | DC_ESUM | DC_OWL | DC_OWL_TRIGGER | DC_FTS) == 0] + + michel = michel[michel.nhit >= 100] + + # Accept events which had a muon more than 800 nanoseconds but less than 20 + # microseconds before them. The 800 nanoseconds cut comes from Richie's + # thesis. He also mentions that the In Time Channel Spread Cut is very + # effective at cutting electron events caused by muons, so I should + # implement that. + # + # Note: We currently don't look across run boundaries. This should be a + # *very* small effect, and the logic to do so would be very complicated + # since I would have to deal with 50 MHz clock rollovers, etc. + if prompt_plus_muons.size and michel.size: + mask = (michel.gtr.values > prompt_plus_muons.gtr.values[:,np.newaxis] + 800) & \ + (michel.gtr.values < prompt_plus_muons.gtr.values[:,np.newaxis] + 20e3) + michel = michel.iloc[np.any(mask,axis=0)] + michel['muon_gtid'] = pd.Series(prompt_plus_muons['gtid'].iloc[np.argmax(mask[:,np.any(mask,axis=0)],axis=0)].values, + index=michel.index.values, + dtype=np.int32) + return michel + else: + # Return an empty slice since we need it to have the same datatype as + # the other dataframes + michel = ev[:0] + michel['muon_gtid'] = -1 + return michel + +def atmospheric_events(ev): + """ + Tags atmospheric events which have a neutron follower. + """ + prompt = ev[ev.prompt] + + # Michel electrons and neutrons can be any event which is not a prompt + # event + follower = ev[~ev.prompt] + + ev['atm'] = np.zeros(len(ev),dtype=np.bool) + + if prompt.size and follower.size: + # neutron followers have to obey stricter set of data cleaning cuts + neutron = follower[follower.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK | DC_ESUM | DC_OWL | DC_OWL_TRIGGER | DC_FTS) == 0] + neutron = neutron[~np.isnan(neutron.ftp_x) & ~np.isnan(neutron.rsp_energy)] + # FIXME: What should the radius cut be here? AV? (r/r_psup)^3 < 0.9? + neutron = neutron[neutron.ftp_r < AV_RADIUS] + neutron = neutron[neutron.rsp_energy > 4.0] + + # neutron events accepted after 20 microseconds and before 250 ms (50 ms during salt) + ev.loc[ev.prompt,'atm'] = np.any((neutron.gtr.values > prompt.gtr.values[:,np.newaxis] + 20e3) & \ + (neutron.gtr.values < prompt.gtr.values[:,np.newaxis] + 250e6),axis=1) + + return ev + +def gtid_sort(ev, first_gtid): + """ + Adds 0x1000000 to the gtid_sort column for all gtids before the first gtid + in a run, which should be passed as a dictionary. This column can then be + used to sort the events sequentially. + + This function should be passed to ev.groupby('run').apply(). We use this + idiom instead of just looping over the groupby results since groupby() + makes a copy of the dataframe, i.e. + + for run, ev_run in ev.groupby('run'): + ev_run.loc[ev_run.gtid < first_gtid[run],'gtid_sort'] += 0x1000000 + + would produce a SettingWithCopyWarning, so instead we use: + + ev = ev.groupby('run',as_index=False).apply(gtid_sort,first_gtid=first_gtid) + + which doesn't have this problem. + """ + # see https://stackoverflow.com/questions/32460593/including-the-group-name-in-the-apply-function-pandas-python + run = ev.name + + if run not in first_gtid: + print_warning("No RHDR bank for run %i! Assuming first event is the first GTID." % run) + first_gtid[run] = ev.gtid.iloc[0] + + ev.loc[ev.gtid < first_gtid[run],'gtid_sort'] += 0x1000000 + + return ev + +def prompt_event(ev): + ev['prompt'] = (ev.nhit >= 100) + ev.loc[ev.prompt,'prompt'] &= np.concatenate(([True],np.diff(ev[ev.prompt].gtr.values) > 250e6)) + return ev + +def plot_corner_plot(ev, title, save=None): + import matplotlib.pyplot as plt + variables = ['r_psup','psi','z','udotr'] + labels = [r'$(r/r_\mathrm{PSUP})^3$',r'$\psi$','z',r'$\vec{u}\cdot\vec{r}$'] + limits = [(0,1),(0,10),(-840,840),(-1,1)] + cuts = [0.9,6,0,-0.5] + + ev = ev.dropna(subset=variables) + + fig = plt.figure(figsize=(FIGSIZE[0],FIGSIZE[0])) + despine(fig,trim=True) + for i in range(len(variables)): + for j in range(len(variables)): + if j > i: + continue + ax = plt.subplot(len(variables),len(variables),i*len(variables)+j+1) + if i == j: + plt.hist(ev[variables[i]],bins=np.linspace(limits[i][0],limits[i][1],100),histtype='step') + plt.gca().set_xlim(limits[i]) + else: + plt.scatter(ev[variables[j]],ev[variables[i]],s=0.5) + plt.gca().set_xlim(limits[j]) + plt.gca().set_ylim(limits[i]) + n = len(ev) + if n: + p_i_lo = np.count_nonzero(ev[variables[i]] < cuts[i])/n + p_j_lo = np.count_nonzero(ev[variables[j]] < cuts[j])/n + p_lolo = p_i_lo*p_j_lo + p_lohi = p_i_lo*(1-p_j_lo) + p_hilo = (1-p_i_lo)*p_j_lo + p_hihi = (1-p_i_lo)*(1-p_j_lo) + n_lolo = np.count_nonzero((ev[variables[i]] < cuts[i]) & (ev[variables[j]] < cuts[j])) + n_lohi = np.count_nonzero((ev[variables[i]] < cuts[i]) & (ev[variables[j]] >= cuts[j])) + n_hilo = np.count_nonzero((ev[variables[i]] >= cuts[i]) & (ev[variables[j]] < cuts[j])) + n_hihi = np.count_nonzero((ev[variables[i]] >= cuts[i]) & (ev[variables[j]] >= cuts[j])) + observed = np.array([n_lolo,n_lohi,n_hilo,n_hihi]) + expected = n*np.array([p_lolo,p_lohi,p_hilo,p_hihi]) + psi = -poisson.logpmf(observed,expected).sum() + poisson.logpmf(observed,observed).sum() + psi /= np.std(-poisson.logpmf(np.random.poisson(observed,size=(10000,4)),observed).sum(axis=1) + poisson.logpmf(observed,observed).sum()) + plt.title(r"$\psi = %.1f$" % psi) + if i == len(variables) - 1: + plt.xlabel(labels[j]) + else: + plt.setp(ax.get_xticklabels(),visible=False) + if j == 0: + plt.ylabel(labels[i]) + else: + plt.setp(ax.get_yticklabels(),visible=False) + plt.axvline(cuts[j],color='k',ls='--',alpha=0.5) + if i != j: + plt.axhline(cuts[i],color='k',ls='--',alpha=0.5) + + plt.tight_layout() + + if save: + plt.savefig(save + ".pdf") + plt.savefig(save + ".eps") + + plt.suptitle(title) + +def intersect_sphere(pos, dir, R): + """ + Compute the first intersection of a ray starting at `pos` with direction + `dir` and a sphere centered at the origin with radius `R`. The distance to + the intersection is returned. + + Example: + + pos = np.array([0,0,0]) + dir = np.array([1,0,0]) + + l = intersect_sphere(pos,dir,PSUP_RADIUS): + if l is not None: + hit = pos + l*dir + print("ray intersects sphere at %.2f %.2f %.2f", hit[0], hit[1], hit[2]) + else: + print("ray didn't intersect sphere") + """ + + b = 2*np.dot(dir,pos) + c = np.dot(pos,pos) - R*R + + if b*b - 4*c <= 0: + # Ray doesn't intersect the sphere. + return None + + # First, check the shorter solution. + l = (-b - np.sqrt(b*b - 4*c))/2 + + # If the shorter solution is less than 0, check the second solution. + if l < 0: + l = (-b + np.sqrt(b*b - 4*c))/2 + + # If the distance is still negative, we didn't intersect the sphere. + if l < 0: + return None + + return l + +def get_dx(row): + pos = np.array([row.x,row.y,row.z]) + dir = np.array([np.sin(row.theta1)*np.cos(row.phi1), + np.sin(row.theta1)*np.sin(row.phi1), + np.cos(row.theta1)]) + l = intersect_sphere(pos,-dir,PSUP_RADIUS) + if l is not None: + pos -= dir*l + michel_pos = np.array([row.x_michel,row.y_michel,row.z_michel]) + return np.linalg.norm(michel_pos-pos) + else: + return 0 + +def dx_to_energy(dx): + lines = [] + with open("../src/muE_water_liquid.txt") as f: + for i, line in enumerate(f): + if i < 10: + continue + if 'Minimum ionization' in line: + continue + if 'Muon critical energy' in line: + continue + lines.append(line) + data = np.genfromtxt(lines) + return np.interp(dx,data[:,8],data[:,0]) + +def f(x): + y = (5/(3*x**2) + 16*x/3 + 4/x + (12-8*x)*np.log(1/x-1) - 8)*np.log(MUON_MASS/ELECTRON_MASS) + y += (6-4*x)*(2*spence(x) - 2*np.log(x)**2 + np.log(x) + np.log(1-x)*(3*np.log(x)-1/x-1) - np.pi**2/3-2) + y += (1-x)*(34*x**2+(5-34*x**2+17*x)*np.log(x) - 22*x)/(3*x**2) + y += 6*(1-x)*np.log(x) + return y + +def michel_spectrum(T): + """ + Michel electron energy spectrum for a free muon. `T` should be the kinetic + energy of the electron or positron in MeV. + + Note: The result is not normalized. + + From https://arxiv.org/abs/1406.3575. + """ + E = T + ELECTRON_MASS + x = 2*E/MUON_MASS + mask = (x > 0) & (x < 1) + y = np.zeros_like(x,dtype=np.double) + y[mask] = GF**2*MUON_MASS**5*x[mask]**2*(6-4*x[mask]+FINE_STRUCTURE_CONSTANT*f(x[mask])/np.pi)/(192*np.pi**3) + y *= 2*MUON_MASS + return y |