From 47cf99a43e6d2c5d4461289720b7154c434659d3 Mon Sep 17 00:00:00 2001
From: Chao Peng <cpeng@anl.gov>
Date: Fri, 9 Dec 2022 15:47:33 -0600
Subject: [PATCH] follow up: update drawing script
---
.../scripts/draw_imaging_event3d.py | 32 +++++++++++++++++--
.../scripts/ml_data_preprocess.py | 1 -
benchmarks/imaging_shower_ML/scripts/utils.py | 5 +++
3 files changed, 35 insertions(+), 3 deletions(-)
diff --git a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py
index 7fd4e3ea..c5741ab7 100644
--- a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py
+++ b/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
diff --git a/benchmarks/imaging_shower_ML/scripts/ml_data_preprocess.py b/benchmarks/imaging_shower_ML/scripts/ml_data_preprocess.py
index dd605f73..076eb4f3 100644
--- a/benchmarks/imaging_shower_ML/scripts/ml_data_preprocess.py
+++ b/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
diff --git a/benchmarks/imaging_shower_ML/scripts/utils.py b/benchmarks/imaging_shower_ML/scripts/utils.py
index bd3e9536..78b40c29 100644
--- a/benchmarks/imaging_shower_ML/scripts/utils.py
+++ b/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)
+
--
GitLab