follow up: data binning for drawing 3D events

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