t0CalTime.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
# -------------------------------------------------------------------------------------------
# example of comparing time vs channel before and after t0 calibration
# input: "t0CalTime_slot<#>.root" from t0LaserCalData.py
# Usage: basf2 t0LaserCalData.py -i <TOPDigits.root> <slot#>
#        TOPDigits.root is the output file of unpackToTOPDigitsWithTBC.py
# -------------------------------------------------------------------------------------------
 
import basf2 as b2
from ROOT import Belle2
import ROOT
from ROOT import gStyle, gROOT, addressof, TH2F
import sys
 
 
gStyle.SetOptStat(0)
gROOT.ProcessLine('struct TreeStruct {\
   int slot;\
   int channel;\
   int pixel;\
   int pmt;\
   float time;\
   float caltime;\
   float height;\
   float width;\
   int quality;\
   float rawtime;\
   int firstwindow;\
   int pmtpixel;\
};')
 
from ROOT import TreeStruct  # noqa
 
args = sys.argv
 
 
class Ntuple(b2.Module):
    ''' t0const ntpule infomation '''
 
    global t0const
    
    t0const = {0: 0.0}
 
    
    f = ROOT.TFile.Open('t0const_slot' + str(args[1]) + '.root')
    for event in f.chT0:
        t0const[event.channel] = event.t0Const
 
    
    histTimeCh = TH2F('before T0Cal slot#' + str(args[1]), 'before T0Cal slot#' + str(args[1]), 512, 0, 511, 500, 50, 100)
    
    histCalTimeCh = TH2F('after T0Cal slot#' + str(args[1]), 'after T0Cal slot#' + str(args[1]), 512, 0, 511, 500, 50, 100)
 
    histTimeCh.GetXaxis().SetTitle('channel')
    histTimeCh.GetYaxis().SetTitle('time [ns]')
    histCalTimeCh.GetXaxis().SetTitle('channel')
    histCalTimeCh.GetYaxis().SetTitle('time [ns]')
 
    def initialize(self):
        ''' Initialize the Module: output root file '''
 
        
        self.filefile = ROOT.TFile('t0CalTime_slot' + str(args[1]) + '.root', 'recreate')
        
        self.treetree = ROOT.TTree('laser', '')
        
        self.datadata = TreeStruct()
 
        for key in TreeStruct.__dict__.keys():
            if '__' not in key:
                formstring = '/F'
                if isinstance(self.datadata.__getattribute__(key), int):
                    formstring = '/I'
                self.treetree.Branch(key, addressof(self.datadata, key), key + formstring)
 
    def event(self):
        ''' Event processor: fill the tree and scatter plots '''
 
        digits = Belle2.PyStoreArray('TOPDigits')
        for digit in digits:
            if digit.getModuleID() == int(args[1]):
                self.datadata.slot = digit.getModuleID()
                self.datadata.channel = digit.getChannel()
                self.datadata.pixel = digit.getPixelID()
                self.datadata.time = digit.getTime()
                self.datadata.caltime = digit.getTime() + t0const.get(digit.getChannel())
                self.datadata.height = digit.getPulseHeight()
                self.datadata.width = digit.getPulseWidth()
                self.datadata.quality = digit.getHitQuality()
                self.datadata.rawtime = digit.getRawTime()
                self.datadata.firstwindow = digit.getFirstWindow()
                self.datadata.pmtpixel = digit.getPMTPixel()
                self.datadata.pmt = digit.getPMTNumber()
 
                self.filefile.cd()
                self.treetree.Fill()
 
                self.histTimeChhistTimeCh.Fill(digit.getChannel(), digit.getTime())
                self.histCalTimeChhistCalTimeCh.Fill(digit.getChannel(), digit.getTime() - t0const.get(digit.getChannel()))
 
    def terminate(self):
        ''' Write the file '''
 
        self.filefile.cd()
        self.filefile.Write()
        self.histTimeChhistTimeCh.Write()
        self.histCalTimeChhistCalTimeCh.Write()
        self.filefile.Close()
 
 
# Create path
main = b2.create_path()
 
# input
main.add_module('RootInput')
 
# write to ntuple
main.add_module(Ntuple())
 
# Show progress of processing
main.add_module('Progress')
 
# Process events
b2.process(main)
 
# Print statistics
print(b2.statistics)