update plot-energy to apply neutron follower cut

1 files changed, 61 insertions, 31 deletions

diff --git a/utils/plot-energy b/utils/plot-energy
index 292ae7b..6867983 100755
--- a/utils/plot-energy
+++ b/utils/plot-energy
@@ -75,6 +75,7 @@ if __name__ == '__main__':
     parser.add_argument("filenames", nargs='+', help="input files")
     args = parser.parse_args()
 
+    events = []
     fit_results = []
 
     for filename in args.filenames:
@@ -84,22 +85,21 @@ if __name__ == '__main__':
 
         for i, event in enumerate(data['data']):
             for ev in event['ev']:
+                events.append((
+                    ev['run'],
+                    ev['gtr'],
+                    ev['nhit'],
+                    ev['gtid'],
+                    ev['dc'],
+                    ev['ftp']['x'] if 'ftp' in ev else np.nan,
+                    ev['ftp']['y'] if 'ftp' in ev else np.nan,
+                    ev['ftp']['z'] if 'ftp' in ev else np.nan,
+                    ev['rsp']['energy'] if 'rsp' in ev and ev['rsp']['energy'] > 0 else np.nan,
+                    ))
+
                 if 'fit' not in ev:
-                    fit_results.append((
-                        ev['run'],
-                        ev['gtr'],
-                        ev['nhit'],
-                        ev['gtid'],
-                        ev['dc'],
-                        0,
-                        np.nan,
-                        np.nan,
-                        np.nan,
-                        np.nan,
-                        np.nan,
-                        np.nan,
-                        np.nan))
                     continue
+
                 for id, fit_result in [x for x in ev['fit'].iteritems() if isinstance(x[0],int)]:
                     # FIXME: Should I just store the particle ids in the YAML
                     # output as a list of particle ids instead of a single
@@ -141,10 +141,7 @@ if __name__ == '__main__':
 
                     fit_results.append((
                         ev['run'],
-                        ev['gtr'],
-                        ev['nhit'],
                         ev['gtid'],
-                        ev['dc'],
                         id,
                         fit_result['posx'],
                         fit_result['posy'],
@@ -160,10 +157,7 @@ if __name__ == '__main__':
     # https://github.com/pandas-dev/pandas/issues/4464
     array = np.array(fit_results,
                      dtype=[('run',np.int),      # run number
-                            ('gtr',np.double),   # 50 MHz clock in ns
-                            ('nhit',np.int),     # number of PMTs hit
                             ('gtid',np.int),     # gtid
-                            ('dc',np.int),       # data cleaning word
                             ('id',np.int),       # particle id
                             ('x', np.double),    # x
                             ('y',np.double),     # y
@@ -175,29 +169,63 @@ if __name__ == '__main__':
                    )
     df = pd.DataFrame.from_records(array)
 
+    array = np.array(events,
+                     dtype=[('run',np.int),      # run number
+                            ('gtr',np.double),   # 50 MHz clock in ns
+                            ('nhit',np.int),     # number of PMTs hit
+                            ('gtid',np.int),     # gtid
+                            ('dc',np.int),       # data cleaning word
+                            ('ftpx',np.double),       # data cleaning word
+                            ('ftpy',np.double),       # data cleaning word
+                            ('ftpz',np.double),       # data cleaning word
+                            ('rsp_energy',np.double)]       # data cleaning word
+                   )
+    df_ev = pd.DataFrame.from_records(array)
+
     # remove events 200 microseconds after a muon
-    muons = df[(df.dc & DC_MUON) != 0]
+    muons = df_ev[(df_ev.dc & DC_MUON) != 0]
+
+    print("number of events = %i" % len(df_ev))
+
+    print("number of muons = %i" % len(muons))
 
-    print(len(df))
-    print("nmuons = %i" % len(muons))
+    df_ev = df_ev[(df_ev.dc & DC_MUON) == 0]
 
-    df = df[(df.dc & DC_MUON) == 0]
+    print("number of events after muon cut = %i" % len(df_ev))
 
     if muons.size:
         # FIXME: need to deal with 50 MHz clock rollover
-        df = df[~np.any((df.gtr.values > muons.gtr.values[:,np.newaxis]) & (df.gtr.values <= (muons.gtr.values[:,np.newaxis] + 200e3)),axis=0)]
+        df_ev = df_ev[~np.any((df_ev.gtr.values > muons.gtr.values[:,np.newaxis]) & (df_ev.gtr.values <= (muons.gtr.values[:,np.newaxis] + 200e3)),axis=0)]
 
-    print(len(df))
+    print("number of events after muon follower cut = %i" % len(df_ev))
 
     # perform prompt event data cleaning
-    df = df[df.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK) == 0]
+    df_ev = df_ev[df_ev.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK) == 0]
 
-    print(len(df))
+    print("number of events after data cleaning = %i" % len(df_ev))
 
-    # apply prompt event selection
-    df = df[df.nhit >= 100]
+    # select prompt events
+    # FIXME: how to deal with two prompt events one after another
+    prompt = df_ev[df_ev.nhit >= 100]
 
-    print(len(df))
+    print("number of events after prompt nhit cut = %i" % len(prompt))
+
+    if prompt.size:
+        # FIXME: need to deal with 50 MHz clock rollover
+        # neutron followers have to obey stricter set of data cleaning cuts
+        neutron = df_ev[df_ev.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)]
+        r = np.sqrt(neutron.ftpx**2 + neutron.ftpy**2 + neutron.ftpz**2)
+        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_ev = prompt[~np.any((neutron.gtr.values > prompt.gtr.values[:,np.newaxis] + 20e3) & (neutron.gtr.values < prompt.gtr.values[:,np.newaxis] + 250e6),axis=1)]
+    else:
+        df_ev = prompt
+
+    print("number of events after neutron follower cut = %i" % len(df_ev))
+
+    df = pd.merge(df,df_ev,how='inner',on=['run','gtid'])
 
     # get rid of events which don't have a fit
     nan = np.isnan(df.fmin.values)
@@ -212,6 +240,8 @@ if __name__ == '__main__':
     # require r < 6 meters
     df = df[np.sqrt(df.x.values**2 + df.y.values**2 + df.z.values**2) < AV_RADIUS]
 
+    print("number of events after radius cut = %i" % len(df))
+
     # Note: Need to design and apply a psi based cut here, and apply the muon
     # and neutron follower cuts.