From b54cfc4ea19f34a995bdbc4788a0219c672afb00 Mon Sep 17 00:00:00 2001
From: Chao Peng <cpeng@anl.gov>
Date: Fri, 9 Dec 2022 16:39:49 -0600
Subject: [PATCH] follow up: data binning for drawing 3D events

---
 .../scripts/draw_imaging_event3d.py              | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py
index c5741ab7..8e685ef4 100644
--- a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py
+++ b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py
@@ -52,7 +52,7 @@ if __name__ == '__main__':
     parser.add_argument(
             '--eta-ngrid', type=int,
             dest='eta_ngrid',
-            default=100,
+            default=40,
             help='number of eta grids for drawing.')
     parser.add_argument(
             '--eta-grid', type=float,
@@ -62,7 +62,7 @@ if __name__ == '__main__':
     parser.add_argument(
             '--phi-ngrid', type=int,
             dest='phi_ngrid',
-            default=100,
+            default=40,
             help='number of phi grids for drawing.')
     parser.add_argument(
             '--phi-grid', type=float,
@@ -92,5 +92,17 @@ if __name__ == '__main__':
     rc, thetac, phic, r0c, etac = cartesian_to_polar(xc, yc, zc)
     r, theta, phi, rc, eta = cartesian_to_polar(*data[['position.x', 'position.y', 'position.z']].values.T)
 
+    dfg = pd.DataFrame(data=np.vstack((*data[['layer', 'energy']].values.T, eta, phi)).T,
+                       columns=['layer', 'energy', 'eta', 'phi'])
+    eta_bins = np.arange(args.eta_ngrid + 1)*args.eta_grid
+    eta_bins = eta_bins - eta_bins[-1]/2. + etac
+    phi_bins = np.arange(args.phi_ngrid + 1)*args.phi_grid
+    phi_bins = phi_bins - phi_bins[-1]/2. + phic
+    dfg.loc[:, 'eta_bin'] = np.digitize(dfg['eta'], eta_bins)
+    dfg.loc[:, 'phi_bin'] = np.digitize(dfg['phi'], phi_bins)
+    print(dfg)
+    # eta = eta - etac
+    # phi = phi - phic
+
     # determine eta and phi windows according to event (cluster) center
 
-- 
GitLab