Skip to content
Snippets Groups Projects
rich_benchmark.py 2.1 KiB
Newer Older
  • Learn to ignore specific revisions
  • Chao Peng's avatar
    Chao Peng committed
    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"))