Skip to content
Snippets Groups Projects
Commit e752d32f authored by Chao Peng's avatar Chao Peng
Browse files

add benchmark script

parent 920200e2
Branches
No related tags found
No related merge requests found
import os
import argparse
from ROOT import TCanvas, TFile, TH1F, TLegend, THStack
from ROOT import gROOT, gInterpreter, gStyle
from ROOT import kBlack, kBlue, kRed, kGreen
# interested particles
PARTICLES = {
2212: r'p',
321: r'K+',
211: r'$\pi+$',
}
parser = argparse.ArgumentParser('RICH pid analysis (simple)')
parser.add_argument('root_file', help='input root file')
parser.add_argument('-o', '--out-dir', default='.', dest='out_dir', help='directory for output files')
parser.add_argument('-n', '--number-events', type=int, default=0, dest='nev',
help='number of events to process, <= 0 means all')
parser.add_argument('--npdet-inc', default='/usr/local/include', dest='npdet_inc',
help='NPDet include directory')
parser.add_argument('--dd4hep-inc', default='/usr/local/include', dest='dd4hep_inc',
help='DD4hep include directory')
args = parser.parse_args()
# include files
gInterpreter.AddIncludePath(args.npdet_inc)
gInterpreter.AddIncludePath(args.dd4hep_inc)
gInterpreter.Declare(r'#include "npdet/PhotoMultiplierHit.h"')
gInterpreter.Declare(r'#include "DDG4/Geant4Particle.h"')
# style
gROOT.SetStyle('Modern')
# read root file and fill data
f = TFile.Open(args.root_file)
tree = f.EVENT
hists = {pid: TH1F(name, 'N.p.e', 50, 0, 201) for pid, name in PARTICLES.items()}
for iev in range(tree.GetEntries()):
tree.GetEntry(iev)
npe = tree.ForwardRICHHits.size()
if npe > 0:
pid = tree.MCParticles[2].pdgID
hists[pid].Fill(npe)
gStyle.SetOptStat(0)
# plot
c1 = TCanvas('c1', 'Dynamic Filling Example', 200, 10, 1200, 800)
c1.GetFrame().SetBorderSize(6)
c1.GetFrame().SetBorderMode(-1)
colors = [kRed, kBlue, kGreen]
leg = TLegend(.7,.7,.9,.9)
hstack = THStack("hs", "N.p.e Distributions;Number of photo-electrons;counts")
for (pid, hist), color in zip(hists.items(), colors):
hist.SetFillStyle(3001)
hist.SetFillColorAlpha(color, 0.5)
hist.SetLineColor(color)
hstack.Add(hist)
leg.AddEntry(hist, PARTICLES[pid], 'LF')
hstack.Draw("nostack")
leg.Draw()
c1.SaveAs(os.path.join(args.out_dir, "rich_npe.png"))
#!/bin/bash
# source /usr/local/bin/thisdd4hep.sh
source /usr/local/bin/thisdd4hep.sh
# python rich_benchmark.py --compact=./pid/GenericRICH_example.xml --input=./sim_output/rich_sim.root --output=./results/rich_benchmarks.root
python ./pid/rich_benchmark.py ./sim_output/rich_sim.root --out-dir=./results --npdet-inc=/home/cpeng/Apps/NPDet/include --dd4hep-inc=/home/cpeng/Apps/DD4hep/include
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment