diff options
Diffstat (limited to 'utils/plot-michels')
-rwxr-xr-x | utils/plot-michels | 163 |
1 files changed, 163 insertions, 0 deletions
diff --git a/utils/plot-michels b/utils/plot-michels new file mode 100755 index 0000000..77c9cc1 --- /dev/null +++ b/utils/plot-michels @@ -0,0 +1,163 @@ +#!/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/>. +""" +Script to plot the fit results for stopping muons and Michels. To run it just +run: + + $ ./plot-muon [list of fit results] + +Currently it will plot energy distributions for external muons, stopping muons, +and michel electrons. +""" +from __future__ import print_function, division +import numpy as np +from scipy.stats import iqr, poisson +from scipy.stats import iqr, norm, beta +from sddm.stats import * + +particle_id = {20: 'e', 22: 'u'} + +def print_particle_probs(data): + n = [len(data[data.id == id]) for id in (20,22,2020,2022,2222)] + + alpha = np.ones_like(n) + n + + mode = dirichlet_mode(alpha) + std = np.sqrt(dirichlet.var(alpha)) + + for i, id in enumerate((20,22,2020,2022,2222)): + particle_id_str = ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + print("P(%s) = %.2f +/- %.2f" % (particle_id_str,mode[i]*100,std[i]*100)) + +if __name__ == '__main__': + import argparse + import numpy as np + import pandas as pd + import sys + import h5py + from sddm.plot_energy import * + from sddm.plot import * + from sddm import setup_matplotlib + + parser = argparse.ArgumentParser("plot fit results") + parser.add_argument("filenames", nargs='+', help="input files") + parser.add_argument("--save", action='store_true', default=False, help="save corner plots for backgrounds") + parser.add_argument("--mc", nargs='+', required=True, help="atmospheric MC files") + parser.add_argument("--nhit-thresh", type=int, default=None, help="nhit threshold to apply to events before processing (should only be used for testing to speed things up)") + args = parser.parse_args() + + setup_matplotlib(args.save) + + import matplotlib.pyplot as plt + + # Loop over runs to prevent using too much memory + evs = [] + rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.filenames],ignore_index=True) + for run, df in rhdr.groupby('run'): + evs.append(get_events(df.filename.values, merge_fits=True, nhit_thresh=args.nhit_thresh)) + ev = pd.concat(evs) + ev_mc = get_events(args.mc, merge_fits=True) + + # Drop events without fits + ev = ev[~np.isnan(ev.fmin)] + ev_mc = ev_mc[~np.isnan(ev_mc.fmin)] + + ev = ev.reset_index() + ev_mc = ev_mc.reset_index() + + # Set all prompt events in the MC to be muons + ev_mc.loc[ev_mc.prompt & ev_mc.filename.str.contains("cosmic"),'muon'] = True + + # First, do basic data cleaning which is done for all events. + ev = ev[ev.signal | ev.muon] + ev_mc = ev_mc[ev_mc.signal | ev_mc.muon] + + # 00-orphan cut + ev = ev[(ev.gtid & 0xff) != 0] + ev_mc = ev_mc[(ev_mc.gtid & 0xff) != 0] + + # Now, we select events tagged by the muon tag which should tag only + # external muons. We keep the sample of muons since it's needed later to + # identify Michel electrons and to apply the muon follower cut + muons = ev[ev.muon] + muons_mc = ev_mc[ev_mc.muon] + + # Try to identify Michel electrons. Currently, the event selection is based + # on Richie's thesis. Here, we do the following: + # + # 1. Apply more data cleaning cuts to potential Michel electrons + # 2. Nhit >= 100 + # 3. It must be > 800 ns and less than 20 microseconds from a prompt event + # or a muon + michel = ev[ev.michel] + michel_mc = ev_mc[ev_mc.michel] + + # remove events 200 microseconds after a muon + ev = ev.groupby('run',group_keys=False).apply(muon_follower_cut) + + ev = ev[ev.psi < 6] + ev_mc = ev_mc[ev_mc.psi < 6] + + handles = [Line2D([0], [0], color='C0'), + Line2D([0], [0], color='C1')] + labels = ('Data','Monte Carlo') + + # For the Michel energy plot, we only look at events where the + # corresponding muon had less than 2500 nhit. The reason for only looking + # at Michel electrons from muons with less than 2500 nhit is because there + # is significant ringing and afterpulsing after a large muon which can + # cause the reconstruction to overestimate the energy. + michel_low_nhit = michel[michel.muon_gtid.isin(ev.gtid.values) & (michel.muon_nhit < 2500)] + michel_low_nhit_mc = michel_mc[michel_mc.muon_gtid.isin(ev_mc.gtid.values) & (michel_mc.muon_nhit < 2500)] + + print("Particle ID probability for Michel electrons:") + print("Data") + print_particle_probs(michel_low_nhit) + print("Monte Carlo") + print_particle_probs(michel_low_nhit_mc) + + fig = plt.figure() + plot_hist2_data_mc(michel_low_nhit,michel_low_nhit_mc) + despine(fig,trim=True) + fig.legend(handles,labels,loc='upper right') + if args.save: + plt.savefig("michel_electrons.pdf") + plt.savefig("michel_electrons.eps") + else: + plt.suptitle("Michel Electrons") + + fig = plt.figure() + bins = np.linspace(0,100,41) + plt.hist(michel_low_nhit.ke.values, bins=bins, histtype='step', color='C0', label="Data") + plt.hist(michel_low_nhit_mc.ke.values, weights=np.tile(len(michel_low_nhit)/len(michel_low_nhit_mc),len(michel_low_nhit_mc.ke.values)), bins=bins, histtype='step', color='C1', label="Monte Carlo") + hist = np.histogram(michel_low_nhit.ke.values,bins=bins)[0] + hist_mc = np.histogram(michel_low_nhit_mc.ke.values,bins=bins)[0] + if hist_mc.sum() > 0: + norm = hist.sum()/hist_mc.sum() + else: + norm = 1.0 + p = get_multinomial_prob(hist,hist_mc,norm) + plt.text(0.95,0.95,"p = %.2f" % p,horizontalalignment='right',verticalalignment='top',transform=plt.gca().transAxes) + despine(fig,trim=True) + plt.xlabel("Energy (MeV)") + plt.tight_layout() + plt.legend() + if args.save: + plt.savefig("michel_electrons_ke.pdf") + plt.savefig("michel_electrons_ke.eps") + else: + plt.title("Michel Electrons") + plt.show() |