1 files changed, 31 insertions, 0 deletions
diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py
index 1c1b7e3..b4882d9 100755
--- a/utils/sddm/plot_energy.py
+++ b/utils/sddm/plot_energy.py
@@ -141,6 +141,20 @@ def michel_cut(ev):
         michel['muon_gtid'] = -1
         return michel
 
+def tag_michels(ev):
+    """
+    Tags michel events with the michel column and also adds a column linking to
+    the muon GTID. The muon also gets labelled with the stopping_muon column.
+    """
+    michels = michel_cut(ev)
+    ev['michel'] = np.zeros(len(ev),dtype=np.bool)
+    ev.loc[michels.index,'michel'] = True
+    ev['muon_gtid'] = -np.ones(len(ev),dtype=np.int)
+    ev.loc[michels.index,'muon_gtid'] = michels['muon_gtid']
+    ev['stopping_muon'] = np.zeros(len(ev),dtype=np.bool)
+    ev.loc[ev.gtid.isin(ev.muon_gtid[ev.muon_gtid > 0].values),'stopping_muon'] = 1
+    return ev
+
 def atmospheric_events(ev):
     """
     Tags atmospheric events which have a neutron follower.
@@ -452,6 +466,12 @@ def get_events(filenames, merge_fits=False):
     fits = pd.concat([pd.read_hdf(filename, "fits") for filename in filenames],ignore_index=True)
     rhdr = pd.concat([pd.read_hdf(filename, "rhdr") for filename in filenames],ignore_index=True)
 
+    if len(ev) == 0:
+        if merge_fits:
+            return ev
+        else:
+            return ev, fits
+
     first_gtid = rhdr.set_index('run').to_dict()['first_gtid']
 
     # First, remove junk events since orphans won't have a 50 MHz clock and so
@@ -566,6 +586,17 @@ def get_events(filenames, merge_fits=False):
     # retrigger cut
     ev = ev.groupby('run',group_keys=False).apply(retrigger_cut)
 
+    # Label instrumentals
+    ev['noise'] = ev.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_ITC | DC_ESUM) != 0
+    ev['neck'] = ((ev.dc & DC_NECK) != 0) & ~ev.noise
+    ev['flasher'] = ((ev.dc & DC_FLASHER) != 0) & ~(ev.noise | ev.neck) & (ev.nhit < 1000)
+    ev['breakdown'] = ((ev.dc & (DC_FLASHER | DC_BREAKDOWN)) != 0) & ~(ev.noise | ev.neck) & (ev.nhit >= 1000)
+    ev['muon'] = ((ev.dc & DC_MUON) != 0) & ~(ev.noise | ev.neck | ev.flasher | ev.breakdown)
+    ev['signal'] = ~(ev.noise | ev.neck | ev.flasher | ev.breakdown | ev.muon)
+    ev['instrumental'] = ~ev.signal
+
+    ev = ev.groupby('run',group_keys=False).apply(tag_michels)
+
     if merge_fits:
         ev = pd.merge(fits,ev,how='inner',on=['run','gtid'])
         # Set the index to (run, gtid) so we can set columns from the single particle results