Add more info in E/p cut scan script, improve statistical uncertainty calculation

Add more info in E/p cut scan script, improve statistical uncertainty calculation

benchmarks/imaging_shower_ML/scripts/epcut_scan.py

@@ -170,6 +170,10 @@ if __name__ == '__main__':
         # print(dfr)
 
 
+    # NOTE assumed e-pi only
+    nel = np.sum(el_hist, axis=1)[0]
+    npi = ntotal - nel
+
     # prepare output container
     best = OrderedDict(
         layer=int(nlayers + 1),
@@ -179,9 +183,13 @@ if __name__ == '__main__':
         )
     ep_dict = OrderedDict(
         info= OrderedDict(
-            nsamples=int(ntotal),
+            nsamples=OrderedDict(
+                total=int(ntotal),
+                electron=nel,
+                pion=npi,
+                ),
             targeted_efficiency=args.eff,
-            tracking_resolution=args.res
+            tracking_resolution=args.res,
             ),
         best=best,
         )
@@ -200,7 +208,16 @@ if __name__ == '__main__':
         eff = np.sum(elvals[idx:])/np.sum(elvals)
         # rejection power
         rej = np.sum(pivals[:idx])/np.sum(pivals)
-        ep_dict['layer_{:d}'.format(i + 1)] = {'ep_cut': ep_cut, 'el_eff': eff, 'pi_rej': rej}
+        ep_dict['layer_{:d}'.format(i + 1)] = OrderedDict(
+                ep_cut=ep_cut,
+                el_eff=eff,
+                pi_rej=rej,
+                nsamples_after_cut=OrderedDict(
+                    total=np.sum(elvals[idx:] + pivals[idx:]) ,
+                    electron=np.sum(elvals[idx:]),
+                    pion=np.sum(pivals[idx:]),
+                    )
+                )
         # greater or [equal] here because always favor the cut on more layers
         if rej >= best['pi_rej']:
             best.update({
@@ -230,15 +247,16 @@ if __name__ == '__main__':
     cuts_rej = np.array([c.get('pi_rej') for c in cuts])
 
     # estimated uncertainty (binomial)
-    nerr = np.sqrt(cuts_rej*(1. - cuts_rej)*ntotal) # npq
+    nerr = np.sqrt(cuts_rej*(1. - cuts_rej)*npi) # npq
     # leftover pions
-    nres = ntotal * (1. - cuts_rej)
-    nres_lo = np.clip(nres - nerr, 1, ntotal)
-    nres_hi = np.clip(nres + nerr, 1, ntotal)
+    nres = npi * (1. - cuts_rej)
+    nres_lo = np.clip(nres - nerr, 1, npi)
+    nres_hi = np.clip(nres + nerr, 1, npi)
 
     # rejection power
-    rej_pow = ntotal/nres
-    rej_err = (rej_pow - ntotal/nres_hi, ntotal/nres_lo - rej_pow)
+    rej_pow = npi/nres
+    # NOTE: assumed half of them are pions, so need a factor of sqrt(2)
+    rej_err = ((rej_pow - npi/nres_hi)*np.sqrt(2), (npi/nres_lo - rej_pow)*np.sqrt(2))
 
     fig, ax1 = plt.subplots(figsize=(8, 8))
     ax2 = ax1.twinx()