Skip to content
Snippets Groups Projects

Add plot script for e/pi separation performance

Open Chao Peng requested to merge update_imaging_ml_benchmarks into master
1 file
+ 68
0
Compare changes
  • Side-by-side
  • Inline
 
"""
 
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)
 
 
Loading