1 files changed, 268 insertions, 0 deletions
diff --git a/utils/plot-muons b/utils/plot-muons
new file mode 100755
index 0000000..2a6e0a2
--- /dev/null
+++ b/utils/plot-muons
@@ -0,0 +1,268 @@
+#!/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 matplotlib.lines import Line2D
+from scipy.stats import iqr, norm, beta
+from scipy.special import spence
+from itertools import izip_longest
+
+particle_id = {20: 'e', 22: r'\mu'}
+
+def plot_hist2(df, muons=False, norm=1.0, label=None, color=None):
+    for id, df_id in sorted(df.groupby('id')):
+        if id == 20:
+            plt.subplot(2,3,1)
+        elif id == 22:
+            plt.subplot(2,3,2)
+        elif id == 2020:
+            plt.subplot(2,3,4)
+        elif id == 2022:
+            plt.subplot(2,3,5)
+        elif id == 2222:
+            plt.subplot(2,3,6)
+
+        if muons:
+            plt.hist(np.log10(df_id.ke.values/1000), bins=np.linspace(0,4.5,100), histtype='step', weights=np.tile(norm,len(df_id.ke.values)), label=label, color=color)
+            plt.xlabel("log10(Energy (GeV))")
+        else:
+            bins = np.logspace(np.log10(20),np.log10(10e3),21)
+            plt.hist(df_id.ke.values, bins=bins, histtype='step', weights=np.tile(norm,len(df_id.ke.values)), label=label, color=color)
+            plt.gca().set_xscale("log")
+            plt.xlabel("Energy (MeV)")
+        plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')
+
+    if len(df):
+        plt.tight_layout()
+
+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,'muon'] = True
+
+    # FIXME: TESTING
+    ev_mc.loc[ev_mc.prompt & (ev_mc.id == 22) & (ev_mc.r_psup > 0.9),'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)
+
+    muons = muons[muons.psi < 6]
+    muons_mc = muons_mc[muons_mc.psi < 6]
+
+    handles = [Line2D([0], [0], color='C0'),
+               Line2D([0], [0], color='C1')]
+    labels = ('Data','Monte Carlo')
+
+    fig = plt.figure()
+    plot_hist2(michel,color='C0')
+    plot_hist2(michel_mc,norm=len(michel)/len(michel_mc),color='C1')
+    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()
+    plot_hist2(muons,muons=True,color='C0')
+    plot_hist2(muons_mc,norm=len(muons)/len(muons_mc),muons=True,color='C1')
+    despine(fig,trim=True)
+    if len(muons):
+        plt.tight_layout()
+    fig.legend(handles,labels,loc='upper right')
+    if args.save:
+        plt.savefig("external_muons.pdf")
+        plt.savefig("external_muons.eps")
+    else:
+        plt.suptitle("External Muons")
+
+    # Plot the energy and angular distribution for external muons
+    fig = plt.figure()
+    plt.subplot(2,1,1)
+    plt.hist(muons.ke.values, bins=np.logspace(3,7,100), histtype='step', color='C0', label="Data")
+    plt.hist(muons_mc.ke.values, bins=np.logspace(3,7,100), histtype='step', color='C1', label="Monte Carlo")
+    plt.legend()
+    plt.xlabel("Energy (MeV)")
+    plt.gca().set_xscale("log")
+    plt.subplot(2,1,2)
+    plt.hist(np.cos(muons.theta.values), bins=np.linspace(-1,1,100), histtype='step', color='C0', label="Data")
+    plt.hist(np.cos(muons_mc.theta.values), bins=np.linspace(-1,1,100), histtype='step', color='C1', label="Monte Carlo")
+    plt.legend()
+    despine(fig,trim=True)
+    plt.xlabel(r"$\cos(\theta)$")
+    plt.tight_layout()
+    if args.save:
+        plt.savefig("muon_energy_cos_theta.pdf")
+        plt.savefig("muon_energy_cos_theta.eps")
+    else:
+        plt.suptitle("Muons")
+
+    stopping_muons = pd.merge(ev[ev.muon & ev.stopping_muon],michel,left_on=['run','gtid'],right_on=['run','muon_gtid'],suffixes=('','_michel'))
+    stopping_muons_mc = pd.merge(ev_mc[ev_mc.muon & ev_mc.stopping_muon],michel_mc,left_on=['run','gtid'],right_on=['run','muon_gtid'],suffixes=('','_michel'))
+
+    if len(stopping_muons):
+        # project muon to PSUP
+        stopping_muons['dx'] = stopping_muons.apply(get_dx,axis=1)
+        stopping_muons_mc['dx'] = stopping_muons_mc.apply(get_dx,axis=1)
+        # energy based on distance travelled
+        stopping_muons['T_dx'] = dx_to_energy(stopping_muons.dx)
+        stopping_muons_mc['T_dx'] = dx_to_energy(stopping_muons_mc.dx)
+        stopping_muons['dT'] = stopping_muons['energy1'] - stopping_muons['T_dx']
+        stopping_muons_mc['dT'] = stopping_muons_mc['energy1'] - stopping_muons_mc['T_dx']
+
+        print(stopping_muons[['run','gtid','energy1','T_dx','dT','gtid_michel','r_michel','ftp_r_michel','id','r']])
+
+        fig = plt.figure()
+        plt.hist((stopping_muons['energy1']-stopping_muons['T_dx'])*100/stopping_muons['T_dx'], bins=np.linspace(-100,100,200), histtype='step', color='C0', label="Data")
+        plt.hist((stopping_muons_mc['energy1']-stopping_muons_mc['T_dx'])*100/stopping_muons_mc['T_dx'], bins=np.linspace(-100,100,200), histtype='step', color='C1', label="Monte Carlo")
+        plt.legend()
+        despine(fig,trim=True)
+        plt.xlabel("Fractional energy difference (\%)")
+        plt.title("Fractional energy difference for Stopping Muons")
+        plt.tight_layout()
+        if args.save:
+            plt.savefig("stopping_muon_fractional_energy_difference.pdf")
+            plt.savefig("stopping_muon_fractional_energy_difference.eps")
+        else:
+            plt.title("Stopping Muon Fractional Energy Difference")
+
+        # 100 bins between 50 MeV and 10 GeV
+        bins = np.linspace(50,2000,10)
+
+        pd_bins = pd.cut(stopping_muons['T_dx'],bins)
+        pd_bins_mc = pd.cut(stopping_muons_mc['T_dx'],bins)
+
+        T = (bins[1:] + bins[:-1])/2
+
+        dT = stopping_muons.groupby(pd_bins)['dT'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
+        dT_mc = stopping_muons_mc.groupby(pd_bins_mc)['dT'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
+
+        fig = plt.figure()
+        plt.errorbar(T, dT['median']*100/T, yerr=dT['median_err']*100/T, color='C0', label="Data")
+        plt.errorbar(T, dT_mc['median']*100/T, yerr=dT_mc['median_err']*100/T, color='C1', label="Monte Carlo")
+        plt.gca().set_ylim(-500,500)
+        plt.legend()
+        despine(fig,trim=True)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel(r"Energy bias (\%)")
+        plt.tight_layout()
+        if args.save:
+            plt.savefig("stopping_muon_energy_bias.pdf")
+            plt.savefig("stopping_muon_energy_bias.eps")
+        else:
+            plt.title("Stopping Muon Energy Bias")
+
+        fig = plt.figure()
+        plt.errorbar(T, dT['iqr_std']*100/T, yerr=dT['iqr_std_err']*100/T, color='C0', label="Data")
+        plt.errorbar(T, dT_mc['iqr_std']*100/T, yerr=dT_mc['iqr_std_err']*100/T, color='C1', label="Monte Carlo")
+        plt.gca().set_ylim(-500,500)
+        despine(fig,trim=True)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel(r"Energy resolution (\%)")
+        plt.tight_layout()
+        if args.save:
+            plt.savefig("stopping_muon_energy_resolution.pdf")
+            plt.savefig("stopping_muon_energy_resolution.eps")
+        else:
+            plt.title("Stopping Muon Energy Resolution")
+
+    # For the Michel energy plot, we only look at the single particle electron
+    # fit
+    michel = michel[michel.id == 20]
+
+    fig = plt.figure()
+    bins = np.linspace(0,100,41)
+    plt.hist(michel.ke.values, bins=bins, histtype='step', color='C0', label="Data")
+    plt.hist(michel_mc.ke.values, weights=np.tile(len(michel)/len(michel_mc),len(michel_mc.ke.values)), bins=bins, histtype='step', color='C1', label="Monte Carlo")
+    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()