update plot-energy script

Tidy up the code in the script by creating a function to plot the energy
distributions for each particle combo. Also fixed a bug which was causing the
neutron follower cut not to be applied.

1 files changed, 59 insertions, 140 deletions

diff --git a/utils/plot-energy b/utils/plot-energy
index f62ac0e..705e53e 100755
--- a/utils/plot-energy
+++ b/utils/plot-energy
@@ -51,13 +51,36 @@ DC_OWL            = 0x80
 DC_OWL_TRIGGER    = 0x100
 DC_FTS            = 0x200
 
-def plot_hist(x, label=None):
-    # determine the bin width using the Freedman Diaconis rule
-    # see https://en.wikipedia.org/wiki/Freedman%E2%80%93Diaconis_rule
-    h = 2*iqr(x)/len(x)**(1/3)
-    n = max(int((np.max(x)-np.min(x))/h),10)
-    bins = np.linspace(np.min(x),np.max(x),n)
-    plt.hist(x, bins=bins, histtype='step', label=label)
+def plot_hist(df, title=None):
+    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)
+
+        plt.hist(df_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
+        plt.xlabel("Energy (MeV)")
+        plt.title(str(id))
+
+    if title:
+        plt.suptitle(title)
+
+    if len(df):
+        plt.tight_layout()
 
 def chunks(l, n):
     """Yield successive n-sized chunks from l."""
@@ -195,27 +218,26 @@ if __name__ == '__main__':
 
     print("number of events = %i" % len(df_ev))
 
-    # First, do basic data cleaning which is done for all events
+    # First, do basic data cleaning which is done for all events.
     df_ev = df_ev[df_ev.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK) == 0]
 
     print("number of events after data cleaning = %i" % len(df_ev))
 
-    # 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
+    # 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 = df_ev[(df_ev.dc & DC_MUON) != 0]
 
     print("number of muons = %i" % len(muons))
 
-    # select prompt events
-    # FIXME: how to deal with two prompt events one after another
-    prompt = df_ev[df_ev.nhit >= 100]
-
+    # Now, select prompt events.
+    #
     # We define a prompt event here as any event with an NHIT > 100 and whose
     # previous > 100 nhit event was more than 250 ms ago
     #
     # Should I use 10 MHz clock here since the time isn't critical and then I
     # don't have to worry about 50 MHz clock rollover?
+    prompt = df_ev[df_ev.nhit >= 100]
     prompt_mask = np.concatenate(([True],np.diff(prompt.gtr.values) > 250e6))
 
     # Michel electrons and neutrons can be any event which is not a prompt
@@ -244,12 +266,13 @@ if __name__ == '__main__':
         # The 800 nanoseconds cut comes from Richie's thesis. He also mentions
         # that the In Time Channel Spread Cut is very effective at cutting
         # electrons events caused by muons, so I should implement that
-        michel = michel[np.any((michel.gtr.values > prompt_plus_muons.gtr.values[:,np.newaxis] + 800) & (michel.gtr.values < prompt_plus_muons.gtr.values[:,np.newaxis] + 20e3),axis=0)]
+        michel = michel[np.any((michel.run.values == prompt_plus_muons.run.values[:,np.newaxis]) & \
+                               (michel.gtr.values > prompt_plus_muons.gtr.values[:,np.newaxis] + 800) & \
+                               (michel.gtr.values < prompt_plus_muons.gtr.values[:,np.newaxis] + 20e3),axis=0)]
     else:
-        michel = prompt_plus_muons.copy()
+        michel = pd.DataFrame().reindex_like(follower)
 
     if prompt.size and follower.size:
-        # FIXME: need to deal with 50 MHz clock rollover
         # 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) == DC_FTS]
         neutron = neutron[~np.isnan(neutron.ftpx) & ~np.isnan(neutron.rsp_energy)]
@@ -257,11 +280,15 @@ if __name__ == '__main__':
         neutron = neutron[r < AV_RADIUS]
         neutron = neutron[neutron.rsp_energy > 4.0]
         # neutron events accepted after 20 microseconds and before 250 ms (50 ms during salt)
-        df_atm = prompt[np.any((neutron.gtr.values > prompt.gtr.values[:,np.newaxis] + 20e3) & (neutron.gtr.values < prompt.gtr.values[:,np.newaxis] + 250e6),axis=1)]
-        df_ev = prompt[~np.any((neutron.gtr.values > prompt.gtr.values[:,np.newaxis] + 20e3) & (neutron.gtr.values < prompt.gtr.values[:,np.newaxis] + 250e6),axis=1)]
+        # FIXME: need to deal with 50 MHz clock rollover
+        neutron_mask = np.any((neutron.run.values == prompt.run.values[:,np.newaxis]) & \
+                              (neutron.gtr.values > prompt.gtr.values[:,np.newaxis] + 20e3) & \
+                              (neutron.gtr.values < prompt.gtr.values[:,np.newaxis] + 250e6),axis=1)
+        df_atm = prompt[neutron_mask]
+        prompt = prompt[~neutron_mask]
     else:
-        df_ev = prompt
-        df_atm = prompt
+        prompt = prompt
+        df_atm = pd.DataFrame().reindex_like(prompt)
 
     print("number of events after neutron follower cut = %i" % len(df_ev))
 
@@ -273,9 +300,9 @@ if __name__ == '__main__':
     if muons.size:
         # FIXME: need to deal with 50 MHz clock rollover
         # remove events 200 microseconds after a muon
-        prompt = prompt[~np.any((prompt.gtr.values > muons.gtr.values[:,np.newaxis]) & (prompt.gtr.values <= (muons.gtr.values[:,np.newaxis] + 200e3)),axis=0)]
-
-    print("number of events after muon follower cut = %i" % len(df_ev))
+        prompt = prompt[~np.any((prompt.run.values == muons.run.values[:,np.newaxis]) & \
+                                (prompt.gtr.values > muons.gtr.values[:,np.newaxis]) & \
+                                (prompt.gtr.values <= (muons.gtr.values[:,np.newaxis] + 200e3)),axis=0)]
 
     df_atm = pd.merge(df,df_atm,how='inner',on=['run','gtid'])
     muons = pd.merge(df,muons,how='inner',on=['run','gtid'])
@@ -314,125 +341,15 @@ if __name__ == '__main__':
     # Note: Need to design and apply a psi based cut here
 
     plt.figure()
-    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)
-
-        plt.hist(df_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
-        plt.xlabel("Energy (MeV)")
-        plt.title(str(id))
-
-    plt.suptitle("Without Neutron Follower")
-
-    if len(df):
-        plt.tight_layout()
-
+    plot_hist(df,"Without Neutron Follower")
     plt.figure()
-    for id, df_atm_id in sorted(df_atm.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)
-
-        plt.hist(df_atm_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
-        plt.xlabel("Energy (MeV)")
-        plt.title(str(id))
-
-    plt.suptitle("With Neutron Follower")
-
-    if len(df_atm):
-        plt.tight_layout()
-
+    plot_hist(df_atm,"With Neutron Follower")
     plt.figure()
-    for id, michel_id in sorted(michel_best_fit.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)
-
-        plt.hist(michel_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
-        plt.xlabel("Energy (MeV)")
-        plt.title(str(id))
-
-    plt.suptitle("Michel Electrons")
-
-    if len(michel):
-        plt.tight_layout()
-
+    plot_hist(michel_best_fit,"Michel Electrons")
     plt.figure()
-    for id, muon_id in sorted(muon_best_fit.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)
-
-        plt.hist(muon_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
-        plt.xlabel("Energy (MeV)")
-        plt.title(str(id))
-
-    plt.suptitle("External Muons")
-
-    if len(muon_best_fit):
-        plt.tight_layout()
+    plot_hist(muon_best_fit,"External Muons")
 
+    # Plot the energy and angular distribution for external muons
     plt.figure()
     plt.subplot(2,1,1)
     plt.hist(muons.ke.values, bins=np.logspace(3,7,100), histtype='step')
@@ -443,6 +360,8 @@ if __name__ == '__main__':
     plt.xlabel(r"$\cos(\theta)$")
     plt.suptitle("Muons")
 
+    # For the Michel energy plot, we only look at the single particle electron
+    # fit
     michel = michel[michel.id == 20]
 
     plt.figure()