TreeWriterModule Class Reference

Inheritance diagram for TreeWriterModule:

Collaboration diagram for TreeWriterModule:

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

Public Attributes
	file
	Output ROOT file.
	tree
	TTree for output data.
	data
	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__()

def __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()

def 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()

def 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

Instance of EventData class.

Definition at line 66 of file pid_ttree_writer.py.

file

file

Output ROOT file.

Definition at line 62 of file pid_ttree_writer.py.

tree

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