svdCoGCalibrationAnalysisToolModule.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
 
 
 
import basf2 as b2
from ROOT import Belle2, TFile, TTree, gROOT
import numpy
 
 
import numpy as np
 
import sys
 
inputFile = sys.argv[1]
outputFile = sys.argv[2]
time0 = "EventT0"
svd_Clusters = "SVDClusters"
svd_ClustersFromTracks = "SVDClustersFromTracks"
svd_EventInfo = "SVDEventInfo"
svd_EventMD = "EventMetaData"
 
b2.set_random_seed(11)
 
gROOT.SetBatch(True)
 
 
class svdCoGCalibrationAnalysisTool(b2.Module):
    """
    Class used to check the goodness of the CoG time calibration
    """
 
    def initialize(self):
        """
        Initialize object (histograms, lists, ...) used by the class
        """
        
        self.outputFileNameoutputFileName = outputFile
        
        self.layerlayer = np.zeros(1, dtype=int)
        
        self.ladderladder = np.zeros(1, dtype=int)
        
        self.sensorsensor = np.zeros(1, dtype=int)
        
        self.sideside = np.zeros(1, dtype=int)
        
        self.TBTB = np.zeros(1, dtype=int)
        
        self.evtTimeevtTime = np.zeros(1, dtype=float)
        
        self.evtTimeSyncevtTimeSync = np.zeros(1, dtype=float)
        
        self.clsTimeclsTime = np.zeros(1, dtype=float)
        
        self.clsSizeclsSize = np.zeros(1, dtype=float)
        
        self.clsChargeclsCharge = np.zeros(1, dtype=float)
        
        self.clsSNRclsSNR = np.zeros(1, dtype=float)
        
        self.runrun = np.zeros(1, dtype=int)
        
        self.expexp = np.zeros(1, dtype=int)
 
        
        self.layerTrackslayerTracks = np.zeros(1, dtype=int)
        
        self.ladderTracksladderTracks = np.zeros(1, dtype=int)
        
        self.sensorTrackssensorTracks = np.zeros(1, dtype=int)
        
        self.sideTrackssideTracks = np.zeros(1, dtype=int)
        
        self.evtTimeTracksevtTimeTracks = np.zeros(1, dtype=float)
        
        self.evtTimeSyncTracksevtTimeSyncTracks = np.zeros(1, dtype=float)
        
        self.clsTimeTracksclsTimeTracks = np.zeros(1, dtype=float)
        
        self.clsSizeTracksclsSizeTracks = np.zeros(1, dtype=float)
        
        self.clsChargeTracksclsChargeTracks = np.zeros(1, dtype=float)
        
        self.clsSNRTracksclsSNRTracks = np.zeros(1, dtype=float)
 
        # self.tFile = TFile(self.outputFileName, 'recreate')
        # self.tFile.cd()
 
        
        self.treetree = TTree('cls', 'tree')
        
        self.treetree.Branch('layer', self.layerlayer, 'layer/I')
        
        self.treetree.Branch('ladder', self.ladderladder, 'ladder/I')
        
        self.treetree.Branch('sensor', self.sensorsensor, 'sensor/I')
        
        self.treetree.Branch('side', self.sideside, 'side/I')
        
        self.treetree.Branch('tb', self.TBTB, 'tb/I')
        
        self.treetree.Branch('evtT0', self.evtTimeevtTime, 'evtT0/D')
        
        self.treetree.Branch('evtT0Sync', self.evtTimeSyncevtTimeSync, 'evtT0Sync/D')
        
        self.treetree.Branch('clsTime', self.clsTimeclsTime, 'clsTime/D')
        
        self.treetree.Branch('clsSize', self.clsSizeclsSize, 'clsSize/D')
        
        self.treetree.Branch('clsCharge', self.clsChargeclsCharge, 'clsCharge/D')
        
        self.treetree.Branch('clsSNR', self.clsSNRclsSNR, 'clsSNR/D')
        
        self.treetree.Branch('run', self.runrun, 'run/I')
        
        self.treetree.Branch('exp', self.expexp, 'exp/I')
 
        
        self.treeTrackstreeTracks = TTree('clsTracks', 'treeTracks')
        
        self.treeTrackstreeTracks.Branch('layer', self.layerTrackslayerTracks, 'layer/I')
        
        self.treeTrackstreeTracks.Branch('ladder', self.ladderTracksladderTracks, 'ladder/I')
        
        self.treeTrackstreeTracks.Branch('sensor', self.sensorTrackssensorTracks, 'sensor/I')
        
        self.treeTrackstreeTracks.Branch('side', self.sideTrackssideTracks, 'side/I')
        
        self.treeTrackstreeTracks.Branch('tb', self.TBTB, 'tb/I')
        
        self.treeTrackstreeTracks.Branch('evtT0', self.evtTimeTracksevtTimeTracks, 'evtT0/D')
        
        self.treeTrackstreeTracks.Branch('evtT0Sync', self.evtTimeSyncTracksevtTimeSyncTracks, 'evtT0Sync/D')
        
        self.treeTrackstreeTracks.Branch('clsTime', self.clsTimeTracksclsTimeTracks, 'clsTime/D')
        
        self.treeTrackstreeTracks.Branch('clsSize', self.clsSizeTracksclsSizeTracks, 'clsSize/D')
        
        self.treeTrackstreeTracks.Branch('clsCharge', self.clsChargeTracksclsChargeTracks, 'clsCharge/D')
        
        self.treeTrackstreeTracks.Branch('clsSNR', self.clsSNRTracksclsSNRTracks, 'clsSNR/D')
        
        self.treeTrackstreeTracks.Branch('run', self.runrun, 'run/I')
        
        self.treeTrackstreeTracks.Branch('exp', self.expexp, 'exp/I')
 
    def event(self):
        """
        Function that allows to cicle on the events
        """
        svd_cls = Belle2.PyStoreArray(svd_Clusters)
        svd_clsFromTrks = Belle2.PyStoreArray(svd_ClustersFromTracks)
        et0 = Belle2.PyStoreObj(time0)
        svd_evtInfo = Belle2.PyStoreObj(svd_EventInfo)
        clsTB = svd_evtInfo.getModeByte().getTriggerBin()
        self.TBTB[0] = ord(clsTB)
        svd_evtMD = Belle2.PyStoreObj(svd_EventMD)
        self.runrun[0] = svd_evtMD.getRun()
        self.expexp[0] = svd_evtMD.getExperiment()
 
        for svdCluster in svd_cls:
            self.clsTimeclsTime[0] = svdCluster.getClsTime()
            self.clsSizeclsSize[0] = svdCluster.getSize()
            self.clsChargeclsCharge[0] = svdCluster.getCharge()
            self.clsSNRclsSNR[0] = svdCluster.getSNR()
            self.layerlayer[0] = svdCluster.getSensorID().getLayerNumber()
            self.ladderladder[0] = svdCluster.getSensorID().getLadderNumber()
            self.sensorsensor[0] = svdCluster.getSensorID().getSensorNumber()
            self.sideside[0] = svdCluster.isUCluster()
            
            self.ftswShiftftswShift = svd_evtInfo.getSVD2FTSWTimeShift(svdCluster.getFirstFrame())
            hasT0 = et0.hasEventT0()
            if hasT0:
                self.evtTimeevtTime[0] = et0.getEventT0()
                self.evtTimeSyncevtTimeSync[0] = et0.getEventT0() - self.ftswShiftftswShift
            else:
                self.evtTimeevtTime[0] = -999
                self.evtTimeSyncevtTimeSync[0] = -999
            self.treetree.Fill()
 
        for svdClusterTracks in svd_clsFromTrks:
            self.clsTimeTracksclsTimeTracks[0] = svdClusterTracks.getClsTime()
            self.clsSizeTracksclsSizeTracks[0] = svdClusterTracks.getSize()
            self.clsChargeTracksclsChargeTracks[0] = svdClusterTracks.getCharge()
            self.clsSNRTracksclsSNRTracks[0] = svdClusterTracks.getSNR()
            self.layerTrackslayerTracks[0] = svdClusterTracks.getSensorID().getLayerNumber()
            self.ladderTracksladderTracks[0] = svdClusterTracks.getSensorID().getLadderNumber()
            self.sensorTrackssensorTracks[0] = svdClusterTracks.getSensorID().getSensorNumber()
            self.sideTrackssideTracks[0] = svdClusterTracks.isUCluster()
            
            self.ftswShiftTracksftswShiftTracks = svd_evtInfo.getSVD2FTSWTimeShift(svdCluster.getFirstFrame())
            hasT0 = et0.hasEventT0()
            if hasT0:
                self.evtTimeTracksevtTimeTracks[0] = et0.getEventT0()
                self.evtTimeSyncTracksevtTimeSyncTracks[0] = et0.getEventT0() - self.ftswShiftTracksftswShiftTracks
            else:
                self.evtTimeTracksevtTimeTracks[0] = -999
                self.evtTimeSyncTracksevtTimeSyncTracks[0] = -999
            self.treeTrackstreeTracks.Fill()
 
    def terminate(self):
        """
        Terminates te class and produces the output rootfile
        """
        
        self.tFiletFile = TFile(self.outputFileNameoutputFileName, 'recreate')
        self.tFiletFile.cd()
        self.treetree.Write()
        self.treeTrackstreeTracks.Write()
        self.tFiletFile.Close()
        del self.tFiletFile
 
# conditions.override_globaltags()
# conditions.globaltags = ["online"]
 
 
main = b2.create_path()
 
rootinput = b2.register_module('RootInput')
rootinput.param('inputFileName', inputFile)
main.add_module(rootinput)
 
# main.add_module("Gearbox")
# main.add_module("Geometry", useDB = True)
 
main.add_module(svdCoGCalibrationAnalysisTool())
 
# Show progress of processing
progress = b2.register_module('Progress')
main.add_module(progress)
 
# Process events
b2.process(main)