makeConfMatrix.py

#!/usr/bin/python                                                                                                                                                                       
from ts_digi_container import *
import ROOT as r

#import argparse
#import importlib
#import re

import sys
from optparse import OptionParser


def main(options,args) :

    inFile=str(options.inFile)
    passName=str(options.passName)
    refPassName=str(options.refPassName)
    maxEvents=int(options.maxEvents)
    nElectrons=int(options.nElectrons)
    maxDelta=float(options.maxDelta)
    modules = ['trigScintDigisTag', 'trigScintDigisUp', 'trigScintDigisDn']

    ## initialize container
    cont = ts_digi_container(inFile,'LDMX_Events')
#    for collection in modules :
#        cont.get_digi_collection(collection+'_'+passName)
    
    for pad in "Tagger", "Up", "Down" :
        cont.get_cluster_collection('TriggerPad'+pad+'Clusters_'+passName)
    cont.get_track_collection('TriggerPadTracks_'+passName)
    cont.get_track_collection('TriggerPadTracks_'+refPassName)
    
    ## configuration for pretty root plots
    r.gROOT.ProcessLine(".L tdrstyle.C")
    r.gROOT.ProcessLine("setTDRStyle()")
    
    ## initialize root histogram

    nTracksMax=7
    nClustersMax=nTracksMax+3
    refString="N_{beam el.}"
    if "truth" not in refPassName :
        print("Interpreting the reference as NOT being beam electron truth collection; just calling it 'reference'")
        refString="N_{tracks}^{ref}"
    hConfMat=r.TH2F("hConfMat", ";"+refString+";N_{tracks}", nTracksMax,-0.5,nTracksMax-0.5, nTracksMax,-0.5,nTracksMax-0.5)
    hConfMatNorm=r.TH2F("hConfMatNorm", ";"+refString+";N_{tracks}", nTracksMax,-0.5,nTracksMax-0.5, nTracksMax,-0.5,nTracksMax-0.5)
    hNtrk=r.TH1F("hNtrk", ";N_{tracks}", nTracksMax,-0.5,nTracksMax-0.5)


    
    if maxEvents == -1 :
     maxEvents = cont.tree.numentries

        
    ## loop over events
    for i in range(maxEvents):
        recoTracks=cont.get_data('TriggerPadTracks_'+passName, "centroid", i)
        refTracks=cont.get_data('TriggerPadTracks_'+refPassName, "centroid", i)
        hConfMat.Fill( nElectrons, len(recoTracks) )
        hNtrk.Fill( len(recoTracks) )
        if nElectrons > -1 and len(recoTracks) > nElectrons :
            print("Overestimate of electron count: "+str(len(recoTracks))+" in event "+str(i+1))

    nEvents = hConfMat.Integral()
    hConfMatNorm.Divide( nEvents )

    if ( nElectrons > -1) :
        eStr="-"+str(nElectrons)+"e"
    else :
        eStr=""

    #before anything else, can be neat to save the histograms to a root file.
    oFileName="confMatrices"+eStr+".root"
    outFile=r.TFile(oFileName, "RECREATE")
    outFile.cd()
    hConfMat.Write()
    hConfMatCl.Write()
    outFile.Close()
 
    #plot!
    c1 = r.TCanvas("c1", "hist canvas", 600,  500)
    c1.SetRightMargin(3.5*c1.GetRightMargin() )
    hConfMat.Draw("colz text")
    c1.SaveAs( hConfMat.GetName()+eStr+".png")


    hConfMat.Scale(1./hConfMat.Integral())
    hConfMat.Draw("colz text")
    c1.SaveAs( hConfMat.GetName()+eStr+"_fractions.png")
    
    hNtrk.Draw("h")
    c1.SaveAs( hNtrk.GetName()+eStr+".png")

    c1.SetLogz()

    hResidUp.Draw("colz")
    c1.SaveAs( hResidUp.GetName()+eStr+".png")
    hResidDn.Draw("colz")
    c1.SaveAs( hResidDn.GetName()+eStr+".png")

    hDistVsNelectrons.Draw("colz")
    c1.SaveAs( hDistVsNelectrons.GetName()+eStr+".png")
    
#    leg=r.TLegend(0.2, 0.5, 0.5, 0.9)
#    leg.AddEntry(hBeamEfrac, "Tagger clusters", "L")
#    leg.AddEntry(hBeamEfracTracks, "Tracks", "L")
#    leg.Draw()
#    c1.SetLogy();
 #   c1.SaveAs( hBeamEfrac.GetName()+".png")



if __name__ == "__main__":
#here: add any option flags needed, and then pick them up in "main" above
	parser = OptionParser()
	parser.add_option('-i', '--inFile', dest='inFile', default='test.root', help='input .root file')
	parser.add_option('-p', '--passName', dest='passName', default='sim', help='pass name to use to look up input variables')
	parser.add_option('-r', '--referencePassName', dest='refPassName', default='truth', help='pass name to use to look up reference input variables')
	parser.add_option('-N', '--maxEvents', dest='maxEvents', default=-1, help='The maximum number of events to process (default: -1 for all)')
	parser.add_option('-n', '--nElectrons', dest='nElectrons', default=-1, help='The actual number of simulated electrons (default: -1 for unset)')
	parser.add_option('-d', '--maxDelta', dest='maxDelta', default=1, help='The maximum distance from seeding cluster used in tracking (default: 1)')

        (options, args) = parser.parse_args()

        main(options,args)