follow up: update drawing script

benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py

@@ -13,7 +13,7 @@ import json
 import numpy as np
 import pandas as pd
 import argparse
-from utils import flatten_collection, cartesian_to_polar, imcal_info
+from utils import flatten_collection, cartesian_to_polar, imcal_info, center_of_gravity
 
 
 pd.set_option('display.max_rows', 500)
@@ -49,6 +49,28 @@ if __name__ == '__main__':
             default='EcalBarrelImagingHitsReco',
             help='name of data branch (edm4eic::CalorimeterHitCollection).'
             )
+    parser.add_argument(
+            '--eta-ngrid', type=int,
+            dest='eta_ngrid',
+            default=100,
+            help='number of eta grids for drawing.')
+    parser.add_argument(
+            '--eta-grid', type=float,
+            dest='eta_grid',
+            default=0.01,
+            help='eta grid size for drawing.')
+    parser.add_argument(
+            '--phi-ngrid', type=int,
+            dest='phi_ngrid',
+            default=100,
+            help='number of phi grids for drawing.')
+    parser.add_argument(
+            '--phi-grid', type=float,
+            dest='phi_grid',
+            default=0.01,
+            help='phi grid size for drawing (rad).')
+    parser.add_argument('--topo-range', type=float, default=50.0, dest='topo_range',
+                        help='half range for projection plot (mrad)')
     args = parser.parse_args()
 
     # read data and MCParticles
@@ -63,6 +85,12 @@ if __name__ == '__main__':
             'energy',
             'position.x', 'position.y', 'position.z',
             ])
-    print(data)
+    etot = data['energy'].sum()
+    xc, yc, zc = center_of_gravity(*data[['position.x', 'position.y', 'position.z']].values.T,
+                                   np.log(data['energy'].values/etot), 4.6)
+
+    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)
 
+    # determine eta and phi windows according to event (cluster) center
 

benchmarks/imaging_shower_ML/scripts/ml_data_preprocess.py

@@ -74,7 +74,6 @@ def format_ml_data(dfh, dft, nlayers=20, nhits=50,
 
     # calculate weighted center of the event, in the future this can be replaced by tracking input
     # NOTE: assumed single particle events
-    # log weighting (4.6
     dfh.loc[:, 'logw'] = log_weights(dfh['eratio'].values)
     dfh.loc[:, 'wtotal'] = event_group['logw'].transform('sum')
     # sanity check

benchmarks/imaging_shower_ML/scripts/utils.py

@@ -86,3 +86,8 @@ def cartesian_to_polar(x, y, z):
     return r, theta, phi, rc, eta
 
 
+def center_of_gravity(x, y, z, w, wbase=0.):
+    wclip = np.clip(w + wbase, 0., None)
+    wn = wclip/np.sum(wclip)
+    return np.sum(x*wn), np.sum(y*wn), np.sum(z*wn)
+