PhokharaEvtgenAnalyze.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
"""
<header>
    <input>PhokharaEvtgenData.root</input>
    <output>PhokharaEvtgenAnalysis.root</output>
    <contact>Kirill Chilikin (K.A.Chilikin@inp.nsk.su)</contact>
    <description>Analysis of e+ e- -> J/psi eta_c events.</description>
</header>
"""
 
import basf2 as b2
import ROOT
import numpy
from ROOT import Belle2
 
 
class PhokharaEvtgenAnalysisModule(b2.Module):
    """ Analysis module for PhokharaEvtgen. """
 
    def __init__(self):
        """Initialization."""
        super(PhokharaEvtgenAnalysisModule, self).__init__()
        
        self.output_fileoutput_file = ROOT.TFile('PhokharaEvtgenAnalysis.root', 'recreate')
        
        self.treetree = ROOT.TTree('tree', '')
        
        self.ecmsecms = numpy.zeros(1, dtype=numpy.float32)
        
        self.gamma_egamma_e = numpy.zeros(1, dtype=numpy.float32)
        
        self.gamma_pxgamma_px = numpy.zeros(1, dtype=numpy.float32)
        
        self.gamma_pygamma_py = numpy.zeros(1, dtype=numpy.float32)
        
        self.gamma_pzgamma_pz = numpy.zeros(1, dtype=numpy.float32)
        
        self.jpsi_ejpsi_e = numpy.zeros(1, dtype=numpy.float32)
        
        self.jpsi_pxjpsi_px = numpy.zeros(1, dtype=numpy.float32)
        
        self.jpsi_pyjpsi_py = numpy.zeros(1, dtype=numpy.float32)
        
        self.jpsi_pzjpsi_pz = numpy.zeros(1, dtype=numpy.float32)
        
        self.lepton_elepton_e = numpy.zeros(1, dtype=numpy.float32)
        
        self.lepton_pxlepton_px = numpy.zeros(1, dtype=numpy.float32)
        
        self.lepton_pylepton_py = numpy.zeros(1, dtype=numpy.float32)
        
        self.lepton_pzlepton_pz = numpy.zeros(1, dtype=numpy.float32)
        self.treetree.Branch('ecms', self.ecmsecms, 'ecms/F')
        self.treetree.Branch('gamma_e', self.gamma_egamma_e, 'gamma_e/F')
        self.treetree.Branch('gamma_px', self.gamma_pxgamma_px, 'gamma_px/F')
        self.treetree.Branch('gamma_py', self.gamma_pygamma_py, 'gamma_py/F')
        self.treetree.Branch('gamma_pz', self.gamma_pzgamma_pz, 'gamma_pz/F')
        self.treetree.Branch('jpsi_e', self.jpsi_ejpsi_e, 'jpsi_e/F')
        self.treetree.Branch('jpsi_px', self.jpsi_pxjpsi_px, 'jpsi_px/F')
        self.treetree.Branch('jpsi_py', self.jpsi_pyjpsi_py, 'jpsi_py/F')
        self.treetree.Branch('jpsi_pz', self.jpsi_pzjpsi_pz, 'jpsi_pz/F')
        self.treetree.Branch('lepton_e', self.lepton_elepton_e, 'lepton_e/F')
        self.treetree.Branch('lepton_px', self.lepton_pxlepton_px, 'lepton_px/F')
        self.treetree.Branch('lepton_py', self.lepton_pylepton_py, 'lepton_py/F')
        self.treetree.Branch('lepton_pz', self.lepton_pzlepton_pz, 'lepton_pz/F')
 
    def event(self):
        """ Event function. """
        mc_initial_particles = Belle2.PyStoreObj('MCInitialParticles')
        self.ecmsecms[0] = mc_initial_particles.getMass()
        mc_particles = Belle2.PyStoreArray('MCParticles')
        for mc_particle in mc_particles:
            # Select virtual photons.
            if (mc_particle.getPDG() != 10022):
                continue
            p = mc_particle.getMomentum()
            self.gamma_egamma_e[0] = mc_particle.getEnergy()
            self.gamma_pxgamma_px[0] = p.X()
            self.gamma_pygamma_py[0] = p.Y()
            self.gamma_pzgamma_pz[0] = p.Z()
            jpsi = mc_particle.getDaughters()[0]
            p = jpsi.getMomentum()
            self.jpsi_ejpsi_e[0] = jpsi.getEnergy()
            self.jpsi_pxjpsi_px[0] = p.X()
            self.jpsi_pyjpsi_py[0] = p.Y()
            self.jpsi_pzjpsi_pz[0] = p.Z()
            lepton = jpsi.getDaughters()[0]
            p = lepton.getMomentum()
            self.lepton_elepton_e[0] = lepton.getEnergy()
            self.lepton_pxlepton_px[0] = p.X()
            self.lepton_pylepton_py[0] = p.Y()
            self.lepton_pzlepton_pz[0] = p.Z()
            self.treetree.Fill()
 
    def terminate(self):
        """ Termination function. """
        self.output_fileoutput_file.cd()
        self.treetree.Write()
        self.output_fileoutput_file.Close()
 
 
# Input.
root_input = b2.register_module('RootInput')
root_input.param('inputFileName', 'PhokharaEvtgenData.root')
 
# Analysis.
phokhara_evtgen = PhokharaEvtgenAnalysisModule()
 
# Create main path.
main = b2.create_path()
 
# Add modules to main path
main.add_module(root_input)
main.add_module(phokhara_evtgen)
 
main.add_module('Progress')
# Run.
b2.process(main)
 
print(b2.statistics)