diff --git a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py index 7fd4e3ea6a4c4f58260316ae3ecee31ee63abace..c5741ab7cea6b828f57662cc04e667be7782d1eb 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 dd605f7337ec33d353d08a0398e7ebb953cf85cd..076eb4f3cb94fdc0b5a4ef4434b9632e5b4811e6 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 bd3e95363e110e991728a0b47015b0ec91af0ae1..78b40c29b958721bb68ae7f8b207d299adb8574c 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) +