create new plot-dc script to replace plot-energy --dc

1 files changed, 152 insertions, 0 deletions

diff --git a/utils/plot-dc b/utils/plot-dc
new file mode 100755
index 0000000..8048192
--- /dev/null
+++ b/utils/plot-dc
@@ -0,0 +1,152 @@
+#!/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 reconstructed quantities for instrumentals. To run it just run:
+
+    $ ./plot-dc [list of fit results]
+
+This will produce corner plots showing the distribution of the high level
+variables used in the contamination analysis for all the different instrumental
+backgrounds and external muons.
+"""
+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):
+    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')
+            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')
+            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()
+
+def plot_hist(df, muons=False):
+    for id, df_id in sorted(df.groupby('id')):
+        if id == 20:
+            plt.subplot(3,4,1)
+        elif id == 22:
+            plt.subplot(3,4,2)
+        elif id == 2020:
+            plt.subplot(3,4,5)
+        elif id == 2022:
+            plt.subplot(3,4,6)
+        elif id == 2222:
+            plt.subplot(3,4,7)
+        elif id == 202020:
+            plt.subplot(3,4,9)
+        elif id == 202022:
+            plt.subplot(3,4,10)
+        elif id == 202222:
+            plt.subplot(3,4,11)
+        elif id == 222222:
+            plt.subplot(3,4,12)
+
+        if muons:
+            plt.hist(np.log10(df_id.ke.values/1000), bins=np.linspace(0,4.5,100), histtype='step')
+            plt.xlabel("log10(Energy (GeV))")
+        else:
+            plt.hist(df_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
+            plt.xlabel("Energy (MeV)")
+        plt.title(str(id))
+
+    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")
+    args = parser.parse_args()
+
+    setup_matplotlib(args.save)
+
+    import matplotlib.pyplot as plt
+
+    ev = get_events(args.filenames, merge_fits=True)
+
+    ev = ev[ev.prompt & ~np.isnan(ev.fmin)]
+
+    with pd.option_context('display.max_rows', None, 'display.max_columns', None):
+        print("Noise events")
+        print(ev[ev.noise][['psi','x','y','z','id1','id2']])
+        print("Muons")
+        print(ev[ev.muon][['psi','r','id1','id2','id3','energy1','energy2','energy3']])
+        print("Neck")
+        print(ev[ev.neck & ev.psi < 6][['psi','r','id1','cos_theta']])
+        print("Flashers")
+        print(ev[ev.flasher & ev.udotr > 0])
+        print("Signal")
+        print(ev[ev.signal])
+
+    # save as PDF b/c EPS doesn't support alpha values
+    if args.save:
+        plot_corner_plot(ev[ev.breakdown],"Breakdowns",save="breakdown_corner_plot")
+        plot_corner_plot(ev[ev.muon],"Muons",save="muon_corner_plot")
+        plot_corner_plot(ev[ev.flasher],"Flashers",save="flashers_corner_plot")
+        plot_corner_plot(ev[ev.neck],"Neck",save="neck_corner_plot")
+        plot_corner_plot(ev[ev.noise],"Noise",save="noise_corner_plot")
+        plot_corner_plot(ev[ev.signal],"Signal",save="signal_corner_plot")
+    else:
+        plot_corner_plot(ev[ev.breakdown],"Breakdowns")
+        plot_corner_plot(ev[ev.muon],"Muons")
+        plot_corner_plot(ev[ev.flasher],"Flashers")
+        plot_corner_plot(ev[ev.neck],"Neck")
+        plot_corner_plot(ev[ev.noise],"Noise")
+        plot_corner_plot(ev[ev.signal],"Signal")
+
+    fig = plt.figure()
+    plot_hist2(ev[ev.flasher])
+    despine(fig,trim=True)
+    plt.suptitle("Flashers")
+    fig = plt.figure()
+    plot_hist2(ev[ev.muon],muons=True)
+    despine(fig,trim=True)
+    plt.suptitle("Muons")
+    plt.show()