small updates to plot-energy

- apply muon follower and muon cuts to atmospheric sample
- print warnings in red
- fix how events with fmin = nan are counted

1 files changed, 64 insertions, 10 deletions

diff --git a/utils/plot-energy b/utils/plot-energy
index 705e53e..a0274b7 100755
--- a/utils/plot-energy
+++ b/utils/plot-energy
@@ -24,6 +24,17 @@ import numpy as np
 from scipy.stats import iqr
 from matplotlib.lines import Line2D
 
+# from https://stackoverflow.com/questions/287871/how-to-print-colored-text-in-terminal-in-python
+class bcolors:
+    HEADER = '\033[95m'
+    OKBLUE = '\033[94m'
+    OKGREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
 # on retina screens, the default plots are way too small
 # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
 # Qt5 will scale everything using the dpi in ~/.Xresources
@@ -87,6 +98,9 @@ def chunks(l, n):
     for i in range(0, len(l), n):
         yield l[i:i + n]
 
+def print_warning(msg):
+    print(bcolors.FAIL + msg + bcolors.ENDC,file=sys.stderr)
+
 if __name__ == '__main__':
     import argparse
     import matplotlib.pyplot as plt
@@ -104,7 +118,7 @@ if __name__ == '__main__':
     for filename in args.filenames:
         print(filename)
         with open(filename) as f:
-            data = yaml.load(f.read(),Loader=Loader)
+            data = yaml.load(f,Loader=Loader)
 
         for i, event in enumerate(data['data']):
             for ev in event['ev']:
@@ -270,7 +284,7 @@ if __name__ == '__main__':
                                (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 = pd.DataFrame().reindex_like(follower)
+        michel = pd.DataFrame(columns=follower.columns)
 
     if prompt.size and follower.size:
         # neutron followers have to obey stricter set of data cleaning cuts
@@ -287,15 +301,15 @@ if __name__ == '__main__':
         df_atm = prompt[neutron_mask]
         prompt = prompt[~neutron_mask]
     else:
-        prompt = prompt
-        df_atm = pd.DataFrame().reindex_like(prompt)
+        df_atm = pd.DataFrame(columns=prompt.columns)
 
-    print("number of events after neutron follower cut = %i" % len(df_ev))
+    print("number of events after neutron follower cut = %i" % len(prompt))
 
     # Get rid of muon events in our main event sample
     prompt = prompt[(prompt.dc & DC_MUON) == 0]
+    df_atm = df_atm[(df_atm.dc & DC_MUON) == 0]
 
-    print("number of events after muon cut = %i" % len(df_ev))
+    print("number of events after muon cut = %i" % len(prompt))
 
     if muons.size:
         # FIXME: need to deal with 50 MHz clock rollover
@@ -303,7 +317,33 @@ if __name__ == '__main__':
         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 = df_atm[~np.any((df_atm.run.values == muons.run.values[:,np.newaxis]) & \
+                                (df_atm.gtr.values > muons.gtr.values[:,np.newaxis]) & \
+                                (df_atm.gtr.values <= (muons.gtr.values[:,np.newaxis] + 200e3)),axis=0)]
+
+    # Check to see if there are any events with missing fit information
+    df_atm_ra = df_atm[['run','gtid']].to_records(index=False)
+    muons_ra = muons[['run','gtid']].to_records(index=False)
+    prompt_ra = prompt[['run','gtid']].to_records(index=False)
+    michel_ra = michel[['run','gtid']].to_records(index=False)
+    df_ra = df[['run','gtid']].to_records(index=False)
+
+    if np.count_nonzero(~np.isin(df_atm_ra,df_ra)):
+        print_warning("skipping %i atmospheric events because they are missing fit information!" % np.count_nonzero(~np.isin(df_atm_ra,df_ra)))
+
+    if np.count_nonzero(~np.isin(muons_ra,df_ra)):
+        print_warning("skipping %i muon events because they are missing fit information!" % np.count_nonzero(~np.isin(muons_ra,df_ra)))
 
+    if np.count_nonzero(~np.isin(prompt_ra,df_ra)):
+        print_warning("skipping %i signal events because they are missing fit information!" % np.count_nonzero(~np.isin(prompt_ra,df_ra)))
+
+    if np.count_nonzero(~np.isin(michel_ra,df_ra)):
+        print_warning("skipping %i Michel events because they are missing fit information!" % np.count_nonzero(~np.isin(michel_ra,df_ra)))
+
+    # Now, we merge the event info with the fitter info.
+    #
+    # Note: This means that the dataframe now contains multiple rows for each
+    # event, one for each fit hypothesis.
     df_atm = pd.merge(df,df_atm,how='inner',on=['run','gtid'])
     muons = pd.merge(df,muons,how='inner',on=['run','gtid'])
     michel = pd.merge(df,michel,how='inner',on=['run','gtid'])
@@ -311,19 +351,32 @@ if __name__ == '__main__':
 
     # get rid of events which don't have a fit
     nan = np.isnan(df.fmin.values)
+
+    if np.count_nonzero(nan):
+        print_warning("skipping %i signal events because the negative log likelihood is nan!" % len(df[nan].groupby(['run','gtid'])))
+
     df = df[~nan]
 
     nan_atm = np.isnan(df_atm.fmin.values)
+
+    if np.count_nonzero(nan_atm):
+        print_warning("skipping %i atmospheric events because the negative log likelihood is nan!" % len(df_atm[nan_atm].groupby(['run','gtid'])))
+
     df_atm = df_atm[~nan_atm]
 
     nan_muon = np.isnan(muons.fmin.values)
+
+    if np.count_nonzero(nan_muon):
+        print_warning("skipping %i muons because the negative log likelihood is nan!" % len(muons[nan_muon].groupby(['run','gtid'])))
+
     muons = muons[~nan_muon]
 
     nan_michel = np.isnan(michel.fmin.values)
-    michel = michel[~nan_michel]
 
-    if np.count_nonzero(nan):
-        print("skipping %i events because they are missing fit information!" % np.count_nonzero(nan),file=sys.stderr)
+    if np.count_nonzero(nan_michel):
+        print_warning("skipping %i michel electron events because the negative log likelihood is nan!" % len(michel[nan_michel].groupby(['run','gtid'])))
+
+    michel = michel[~nan_michel]
 
     # get the best fit
     df = df.sort_values('fmin').groupby(['run','gtid']).first()
@@ -366,7 +419,8 @@ if __name__ == '__main__':
 
     plt.figure()
     plt.hist(michel.ke.values, bins=np.linspace(20,100,100), histtype='step', label="My fitter")
-    plt.hist(michel[~np.isnan(michel.rsp_energy.values)].rsp_energy.values, bins=np.linspace(20,100,100), histtype='step',label="RSP")
+    if michel.size:
+        plt.hist(michel[~np.isnan(michel.rsp_energy.values)].rsp_energy.values, bins=np.linspace(20,100,100), histtype='step',label="RSP")
     plt.xlabel("Energy (MeV)")
     plt.title("Michel Electrons")
     plt.legend()