diff --git a/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd4e3ea6a4c4f58260316ae3ecee31ee63abace --- /dev/null +++ b/benchmarks/imaging_shower_ML/scripts/draw_imaging_event3d.py @@ -0,0 +1,68 @@ +""" + A script to draw imaging events in 3D + Assumed the input reconstruction file has only single-particle events + + Chao Peng (ANL) + 2022/12/06 +""" +import ROOT +import os +import gc +import sys +import json +import numpy as np +import pandas as pd +import argparse +from utils import flatten_collection, cartesian_to_polar, imcal_info + + +pd.set_option('display.max_rows', 500) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Visualize the cluster (layer-wise) from analysis') + parser.add_argument( + 'file', type=str, + help='path to the input file.' + ) + parser.add_argument( + '-o', type=str, + default='.', + dest='outdir', + help='output directory.' + ) + parser.add_argument( + '-e', type=int, + default=0, + dest='event', + help='event number to draw.' + ) + parser.add_argument( + '--nhits', type=int, + default=50, + dest='nhits', + help='number of hits per layer.' + ) + parser.add_argument( + '-b', '--branch', type=str, + dest='branch', + default='EcalBarrelImagingHitsReco', + help='name of data branch (edm4eic::CalorimeterHitCollection).' + ) + args = parser.parse_args() + + # read data and MCParticles + rdf = ROOT.RDataFrame("events", args.file) + # event range + ntotal = rdf.Count().GetValue() + iev = int(np.clip(args.event, 0, ntotal)) + + # data + data = flatten_collection(rdf, args.branch, (iev, iev + 1), cols=[ + 'layer', + 'energy', + 'position.x', 'position.y', 'position.z', + ]) + print(data) + +