pid_ttree_writer.py

#!/usr/bin/env python3
 
 
 
 
 
import basf2 as b2
 
# Some ROOT tools
import ROOT
from ROOT import Belle2
from ROOT import gROOT, AddressOf
 
# Define a ROOT struct to hold output data in the TTree.
gROOT.ProcessLine('struct TreeStruct {\
    float lld_dedx;\
    float lld_top;\
    float p;\
    float phi;\
    float costheta;\
    float trackmomentum;\
    int iskaon;\
};'
                  )
 
 
from ROOT import TreeStruct  # noqa
# define the python module to save the PID information
 
 
class TreeWriterModule(b2.Module):
 
    """
    This module writes its output to a ROOT tree.
    Adapted from pxd/validation/PXDValidationTTreeSimHit.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)
 
    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())}')
 
    def terminate(self):
        """ Close the output file."""
 
        self.file.cd()
        self.file.Write()
        self.file.Close()
 
 
# create path
main = b2.create_path()
 
# use the input file defined via command line
main.add_module(b2.register_module('RootInput'))
 
# add the python module defined above
main.add_module(TreeWriterModule())
 
# process events and print call statistics
b2.process(main)
print(b2.statistics)