TreeWriterModule Class Reference

Inheritance diagram for TreeWriterModule:

Public Member Functions
	__init__ (self)
	event (self)
	terminate (self)

Public Attributes
	file = ROOT.TFile('PID_TTree.root', 'recreate')
	Output ROOT file.
	tree = ROOT.TTree('tree', '')
	TTree for output data.
	data = TreeStruct()
	Instance of EventData class.

Detailed Description

This module writes its output to a ROOT tree.
Adapted from pxd/validation/PXDValidationTTreeSimHit.py

Definition at line 49 of file pid_ttree_writer.py.

Constructor & Destructor Documentation

__init__()

__init__

(

self

)

Initialize the module

Definition at line 56 of file pid_ttree_writer.py.

    def __init__(self):
        """Initialize the module"""
 
        super().__init__()
 
        
        self.file = ROOT.TFile('PID_TTree.root', 'recreate')
        
        self.tree = ROOT.TTree('tree', '')
        
        self.data = TreeStruct()
        """ Declare tree branches """
        for key in TreeStruct.__dict__.keys():
            if '__' not in key:
                formstring = '/F'
                if isinstance(self.data.__getattribute__(key), int):
                    formstring = '/I'
                self.tree.Branch(key, AddressOf(self.data, key), key + formstring)
 

Member Function Documentation

event()

event

(

self

)

Store TOP and dE/dx info in tree

Definition at line 75 of file pid_ttree_writer.py.

    def event(self):
        """Store TOP and dE/dx info in tree"""
 
        pids = Belle2.PyStoreArray('PIDLikelihoods')
        for pid in pids:
            track = pid.getRelatedFrom('Tracks')
            if track:
                mcpart = track.getRelatedTo('MCParticles')
            try:
                pdg = abs(mcpart.getPDG())
                if pdg != 211 and pdg != 321:
                    continue
                momentumVec = mcpart.getMomentum()
                momentum = momentumVec.Mag()
                if momentum > 3.5:  # cut off
                    continue
                costheta = momentumVec.CosTheta()
                phi = momentumVec.Phi()
                fitresult = track.getTrackFitResultWithClosestMass(Belle2.Const.pion)
                if fitresult:
                    trackmomentum = fitresult.getMomentum().Mag()
                else:
                    trackmomentum = 0.0
 
                # particle to compare with pions
                selectedpart = Belle2.Const.kaon
                pid_dedx = Belle2.Const.PIDDetectorSet(Belle2.Const.CDC)
                pid_top = Belle2.Const.PIDDetectorSet(Belle2.Const.TOP)
                logl_sel = pid.getLogL(selectedpart, pid_dedx)
                logl_pi = pid.getLogL(Belle2.Const.pion, pid_dedx)
                dedx_DLL = logl_pi - logl_sel
 
                logl_sel = pid.getLogL(selectedpart, pid_top)
                logl_pi = pid.getLogL(Belle2.Const.pion, pid_top)
                top_DLL = logl_pi - logl_sel
 
                self.data.lld_dedx = dedx_DLL
                self.data.lld_top = top_DLL
                self.data.p = momentum
                self.data.phi = phi
                self.data.costheta = costheta
                self.data.trackmomentum = trackmomentum
                self.data.iskaon = pdg == 321
 
                # Fill tree
                self.file.cd()
                self.tree.Fill()
            except BaseException:
 
                # some tracks don't have an mcparticle (fixed now)
                b2.B2WARNING('problems with track <-> mcparticle relations')
                event = Belle2.PyStoreObj('EventMetaData').obj().getEvent()
                print(f'event: {int(event)}, track: {int(track.getArrayIndex())}')
 

terminate()

terminate

(

self

)

 Close the output file.

Definition at line 129 of file pid_ttree_writer.py.

    def terminate(self):
        """ Close the output file."""
 
        self.file.cd()
        self.file.Write()
        self.file.Close()
 
 
# create path

Member Data Documentation

data

data = TreeStruct()

Instance of EventData class.

Definition at line 66 of file pid_ttree_writer.py.

file

file = ROOT.TFile('PID_TTree.root', 'recreate')

Output ROOT file.

Definition at line 62 of file pid_ttree_writer.py.

tree

tree = ROOT.TTree('tree', '')

TTree for output data.

Definition at line 64 of file pid_ttree_writer.py.

---

The documentation for this class was generated from the following file:

• reconstruction/examples/pid_ttree_writer.py