CoGCalibration_utils.py

# !/usr/bin/env python
# -*- coding: utf-8 -*-
 
#
# evaluates the CoG corrections, create a localDB
# with the corrections and a root file to check
# the corrections
#
# usage: basf2 SVDCoGTimeCalibratinWithErrorImporter localDB filename
# localDB = name of the local DB folder
# filename = single root file, or file with the list of reconstructed files
#
# this script can be launched with launch_calibration_cog.sh in the
# B2SVD project, svd_CoGHitTime_calibration repository
#
 
 
from basf2 import *
from svd import *
import ROOT
from ROOT import Belle2, TFile, TTree, TH1F, TH2F, TH2D, TGraph, TFitResultPtr
from ROOT import TROOT, gROOT, TF1, TMath, gStyle, gDirectory
import os
import numpy
import math
import random
from array import array
import basf2
import sys
from ROOT.Belle2 import SVDCoGCalibrationFunction
from ROOT.Belle2 import SVDCoGTimeCalibrations
 
import matplotlib.pyplot as plt
 
svd_recoDigits = "SVDRecoDigitsFromTracks"
cdc_Time0 = "EventT0"
svd_Clusters = "SVDClustersFromTracks"
 
gROOT.SetBatch(True)
 
# mode = True
 
 
class SVDCoGTimeCalibrationImporterModule(basf2.Module):
    """
    Python class used for evaluating the CoG corrections, create a localDB,
    creating a localDB with the corrections and a root file to check the corrections
    """
 
    def notApplyCorrectForCDCLatency(self, mode):
        """
        Function that allows to set if apply the CDC latency correction
 
        parameters:
             mode (bool):
             - if True -> not apply correction
             - if False -> apply correction
        """
        
        self.notApplyCDCLatencyCorrection = mode
        print("Not Correct for CDC latency: " + str(mode) + " " + str(self.notApplyCDCLatencyCorrection))
 
    def fillLists(self, mode_byte_object, svdClusters_rel_RecoTracks_cl):
        """
        Function that fill the lists needed for the CoG corrections
 
        parameters:
             mode_byte_object (modeByte): modeByte that contains the information about the TB
             svdClusters_rel_RecoTracks_cl (SVDCluster): cluster related to tracks
        """
 
        timeCluster = svdClusters_rel_RecoTracks_cl.getClsTime()
        snrCluster = svdClusters_rel_RecoTracks_cl.getSNR()
        layerCluster = svdClusters_rel_RecoTracks_cl.getSensorID().getLayerNumber()
        layerIndex = layerCluster - 3
        sensorCluster = svdClusters_rel_RecoTracks_cl.getSensorID().getSensorNumber()
        sensorIndex = sensorCluster - 1
        ladderCluster = svdClusters_rel_RecoTracks_cl.getSensorID().getLadderNumber()
        ladderIndex = ladderCluster - 1
        sideCluster = svdClusters_rel_RecoTracks_cl.isUCluster()
        if sideCluster:
            sideIndex = 1
        else:
            sideIndex = 0
 
        hasTimezero = self.cdcEventT0.hasEventT0()
        # print("Time: " + str(hasTimezero))
        if hasTimezero:
            TBClusters = mode_byte_object.getTriggerBin()
            TBIndex = ord(TBClusters)
            tZero = self.cdcEventT0.getEventT0()
            # tZero_err = self.cdcEventT0.getEventT0Uncertainty()
            tZero_err = 5.1
            tZeroSync = tZero - 4000./509 * (3 - TBIndex)
            et0 = self.EventT0Hist
            et0.Fill(tZeroSync)
            # print(str(tZero_err))
 
            resHist = self.resList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex]
            resHist.Fill(timeCluster - tZeroSync)
            spHist = self.spList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex]
            # spHist.Fill(timeCluster, 1.3*timeCluster - 50 + random.gauss(0,10))
            spHist.Fill(timeCluster, tZeroSync)
            cogHist = self.cogList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex]
            cogHist.Fill(timeCluster)
            cdcHist = self.cdcList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex]
            cdcHist.Fill(tZeroSync)
            snrHist = self.snrList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex]
            snrHist.Fill(snrCluster)
 
            self.nList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex] += 1
            self.sumCOGList[layerIndex][ladderIndex][sensorIndex][sideIndex][TBIndex] += timeCluster
 
            
            self.NTOT = self.NTOT + 1
 
    def set_localdb(self, localDB):
        """
        Function that allows to set the localDB
 
        parameters:
             localDB (str): Name of the localDB used
        """
        
        self.localdb = localDB
 
    def initialize(self):
        """
        Initialize object (histograms, lists, ...) used by the class
        """
 
        
        self.outputFileName = "CoGCorrectionMonitor_" + self.localdb + ".root"
 
        
        self.resList = []
        
        self.spList = []
        
        self.cogList = []
        
        self.cdcList = []
        
        self.snrList = []
        
        self.nList = []
        
        self.sumCOGList = []
 
        geoCache = Belle2.VXD.GeoCache.getInstance()
 
        self.Evt = 0
        for layer in geoCache.getLayers(Belle2.VXD.SensorInfoBase.SVD):
            layerList0 = []
            layerList1 = []
            layerList2 = []
            layerList3 = []
            layerList4 = []
            layerList5 = []
            layerList6 = []
            layerList7 = []
            layerList8 = []
            self.resList.append(layerList0)
            self.spList.append(layerList1)
            self.cogList.append(layerList2)
            self.cdcList.append(layerList3)
            self.snrList.append(layerList4)
            self.nList.append(layerList8)
            self.sumCOGList.append(layerList7)
            # layerNumber = layer.getLayerNumber()
            for ladder in geoCache.getLadders(layer):
                ladderList0 = []
                ladderList1 = []
                ladderList2 = []
                ladderList3 = []
                ladderList4 = []
                ladderList5 = []
                ladderList6 = []
                ladderList7 = []
                ladderList8 = []
                layerList0.append(ladderList0)
                layerList1.append(ladderList1)
                layerList2.append(ladderList2)
                layerList3.append(ladderList3)
                layerList4.append(ladderList4)
                layerList5.append(ladderList5)
                layerList6.append(ladderList6)
                layerList7.append(ladderList7)
                layerList8.append(ladderList8)
                # ladderNumber = ladder.getLadderNumber()
                for sensor in geoCache.getSensors(ladder):
                    sensorList0 = []
                    sensorList1 = []
                    sensorList2 = []
                    sensorList3 = []
                    sensorList4 = []
                    sensorList5 = []
                    sensorList6 = []
                    sensorList7 = []
                    sensorList8 = []
                    ladderList0.append(sensorList0)
                    ladderList1.append(sensorList1)
                    ladderList2.append(sensorList2)
                    ladderList3.append(sensorList3)
                    ladderList4.append(sensorList4)
                    ladderList5.append(sensorList5)
                    ladderList6.append(sensorList6)
                    ladderList7.append(sensorList7)
                    ladderList8.append(sensorList8)
                    # sensorNumber = sensor.getSensorNumber()
                    for side in range(2):
                        sideList0 = []
                        sideList1 = []
                        sideList2 = []
                        sideList3 = []
                        sideList4 = []
                        sideList5 = []
                        sideList6 = []
                        sideList7 = []
                        sideList8 = []
                        sensorList0.append(sideList0)
                        sensorList1.append(sideList1)
                        sensorList2.append(sideList2)
                        sensorList3.append(sideList3)
                        sensorList4.append(sideList4)
                        sensorList5.append(sideList5)
                        sensorList6.append(sideList6)
                        sensorList7.append(sideList7)
                        sensorList8.append(sideList8)
 
        for i in geoCache.getLayers(Belle2.VXD.SensorInfoBase.SVD):
            layerN = i.getLayerNumber()
            li = layerN - 3
            for j in geoCache.getLadders(i):
                ladderN = j.getLadderNumber()
                ldi = ladderN - 1
                for k in geoCache.getSensors(j):
                    sensorN = k.getSensorNumber()
                    si = sensorN - 1
                    for s in range(2):
                        for t in range(4):
                            self.resList[li][ldi][si][s].append(
                                TH1F("res" + "_" + str(k) + "." + str(s) + "." + str(t), " ", 200, -100, 100))
                            self.spList[li][ldi][si][s].append(
                                TH2D("sp" + "_" + str(k) + "." + str(s) + "." + str(t), " ", 300, -150, 150, 300, -150, 150))
                            self.cogList[li][ldi][si][s].append(
                                TH1F("cog" + "_" + str(k) + "." + str(s) + "." + str(t), " ", 200, -100, 100))
                            self.cdcList[li][ldi][si][s].append(
                                TH1F("cdc" + "_" + str(k) + "." + str(s) + "." + str(t), " ", 200, -100, 100))
                            self.snrList[li][ldi][si][s].append(
                                TH1F("snr" + "_" + str(k) + "." + str(s) + "." + str(t), " ", 100, 0, 100))
                            self.nList[li][ldi][si][s].append(0)
                            self.sumCOGList[li][ldi][si][s].append(0)
        
        self.EventT0Hist = TH1F("EventT0", " ", 200, -100, 100)
        
        self.AlphaUTB = TH2F("alphaVsTB_U", " ", 400, 0.5, 2, 4, 0, 4)
        self.AlphaUTB.GetXaxis().SetTitle("alpha")
        self.AlphaUTB.GetYaxis().SetTitle("trigger bin")
        
        self.AlphaVTB = TH2F("alphaVsTB_V", " ", 400, 0.5, 2, 4, 0, 4)
        self.AlphaVTB.GetXaxis().SetTitle("alpha")
        self.AlphaVTB.GetYaxis().SetTitle("trigger bin")
        
        self.BetaUTB = TH2F("betaVsTB_U", " ", 200, -100, 100, 4, 0, 4)
        self.BetaUTB.GetXaxis().SetTitle("beta (ns)")
        self.BetaUTB.GetYaxis().SetTitle("trigger bin")
        
        self.BetaVTB = TH2F("betaVsTB_V", " ", 200, -100, 100, 4, 0, 4)
        self.BetaVTB.GetXaxis().SetTitle("beta (ns)")
        self.BetaVTB.GetYaxis().SetTitle("trigger bin")
 
        
        self.MeanHistVTB = TH2F("meanHistVsTB_V", " ", 100, -10, 10, 4, 0, 4)
        self.MeanHistVTB.GetXaxis().SetTitle("distribution mean (ns)")
        self.MeanHistVTB.GetYaxis().SetTitle("trigger bin")
        
        self.MeanHistUTB = TH2F("meanHistVsTB_U", " ", 100, -10, 10, 4, 0, 4)
        self.MeanHistUTB.GetXaxis().SetTitle("distribution mean (ns)")
        self.MeanHistUTB.GetYaxis().SetTitle("trigger bin")
        
        self.RMSHistVTB = TH2F("rmsHistVsTB_V", " ", 100, 0, 10, 4, 0, 4)
        self.RMSHistVTB.GetXaxis().SetTitle("distribution RMS (ns)")
        self.RMSHistVTB.GetYaxis().SetTitle("trigger bin")
        
        self.RMSHistUTB = TH2F("rmsHistVsTB_U", " ", 100, 0, 10, 4, 0, 4)
        self.RMSHistUTB.GetXaxis().SetTitle("distribution RMS (ns)")
        self.RMSHistUTB.GetYaxis().SetTitle("trigger bin")
        
        self.MeanFitVTB = TH2F("meanFitVsTB_V", " ", 100, -10, 10, 4, 0, 4)
        self.MeanFitVTB.GetXaxis().SetTitle("fit mean (ns)")
        self.MeanFitVTB.GetYaxis().SetTitle("trigger bin")
        
        self.MeanFitUTB = TH2F("meanFitVsTB_U", " ", 100, -10, 10, 4, 0, 4)
        self.MeanFitUTB.GetXaxis().SetTitle("fit mean (ns)")
        self.MeanFitUTB.GetYaxis().SetTitle("trigger bin")
        
        self.RMSFitUTB = TH2F("rmsFitVsTB_U", " ", 100, 0, 10, 4, 0, 4)
        self.RMSFitUTB.GetXaxis().SetTitle("fit sigma (ns)")
        self.RMSFitUTB.GetYaxis().SetTitle("trigger bin")
        
        self.RMSFitVTB = TH2F("rmsFitVsTB_V", " ", 100, 0, 10, 4, 0, 4)
        self.RMSFitVTB.GetXaxis().SetTitle("fit sigma (ns)")
        self.RMSFitVTB.GetYaxis().SetTitle("trigger bin")
 
        
        self.gaus = TF1("gaus", 'gaus(0)', -150, 100)
 
        self.NTOT = 0
 
    def event(self):
        """
        Function that allows to cicle on the events
        """
        svd_evt_info = Belle2.PyStoreObj('SVDEventInfo')
        mode_byte = svd_evt_info.getModeByte()
        timeClusterU = 0
        timeClusterV = 0
        sideIndex = 0
        TBIndexU = 0
        TBIndexV = 0
        
        self.Evt = self.Evt + 1
 
        
        self.cdcEventT0 = Belle2.PyStoreObj(cdc_Time0)
        svdCluster_list = Belle2.PyStoreArray(svd_Clusters)
        svdRecoDigit_list = Belle2.PyStoreArray(svd_recoDigits)
 
        for svdCluster in svdCluster_list:
            svdRecoDigit = svdCluster.getRelatedTo(svd_recoDigits)
            self.fillLists(mode_byte, svdCluster)
 
    def terminate(self):
        """
        Terminates te class and produces the output rootfile
        """
 
        tfile = TFile(self.outputFileName, 'recreate')
        iov = Belle2.IntervalOfValidity.always()
 
        payload = Belle2.SVDCoGTimeCalibrations.t_payload()
 
        timeCal = SVDCoGCalibrationFunction()
        # Bias and Scale
        tbBias = [-50, -50, -50, -50]
        tbScale = [1, 1, 1, 1]
        tbBias_err = [1, 1, 1, 1]
        tbScale_err = [1, 1, 1, 1]
        tbCovScaleBias = [1, 1, 1, 1]
 
        TCOGMEAN = 0
        T0MEAN = 0
 
        '''
        geoCache = Belle2.VXD.GeoCache.getInstance()
        for layer in geoCache.getLayers(Belle2.VXD.SensorInfoBase.SVD):
            layerNumber = layer.getLayerNumber()
            li = layerNumber - 3
            for ladder in geoCache.getLadders(layer):
                ladderNumber = ladder.getLadderNumber()
                ldi = ladderNumber - 1
                for sensor in geoCache.getSensors(ladder):
                    sensorNumber = sensor.getSensorNumber()
                    si = sensorNumber - 1
                    for side in range(2):
                        for tb in range(4):
                            n = self.nList[li][ldi][si][side][tb]
                            NTOT += n
        '''
        geoCache = Belle2.VXD.GeoCache.getInstance()
        gDirectory.mkdir("plots")
        gDirectory.cd("plots")
        for layer in geoCache.getLayers(Belle2.VXD.SensorInfoBase.SVD):
            layerNumber = layer.getLayerNumber()
            li = layerNumber - 3
            gDirectory.mkdir("layer" + str(layer))
            gDirectory.cd("layer" + str(layer))
            for ladder in geoCache.getLadders(layer):
                ladderNumber = ladder.getLadderNumber()
                ldi = ladderNumber - 1
                for sensor in geoCache.getSensors(ladder):
                    sensorNumber = sensor.getSensorNumber()
                    si = sensorNumber - 1
                    for side in range(2):
                        for tb in range(4):
                            # Resolution distribution Histograms with Gaussian Fit
                            res = self.resList[li][ldi][si][side][tb]
                            fitResult = int(TFitResultPtr(res.Fit(self.gaus, "R")))
 
                            if res.GetEntries() > 5:
                                if side == 1:
                                    self.MeanHistUTB.Fill(res.GetMean(), tb)
                                    self.RMSHistUTB.Fill(res.GetRMS(), tb)
                                    if fitResult > -1:
                                        self.MeanFitUTB.Fill(self.gaus.GetParameter(1), tb)
                                        self.RMSFitUTB.Fill(self.gaus.GetParameter(2), tb)
                                else:
                                    self.MeanHistVTB.Fill(res.GetMean(), tb)
                                    self.RMSHistVTB.Fill(res.GetRMS(), tb)
                                    if fitResult > -1:
                                        self.MeanFitVTB.Fill(self.gaus.GetParameter(1), tb)
                                        self.RMSFitVTB.Fill(self.gaus.GetParameter(2), tb)
 
                            res.Write()
                            # COG Distribution Histograms
                            cog = self.cogList[li][ldi][si][side][tb]
                            cog.Write()
                            # CDC EventT0 Distribution Histograms
                            cdc = self.cdcList[li][ldi][si][side][tb]
                            cdc.Write()
                            # SNR Distribution Histograms
                            snr = self.snrList[li][ldi][si][side][tb]
                            snrMean = snr.GetMean()
                            snr.Write()
                            # ScatterPlot Histograms with Linear Fit
                            sp = self.spList[li][ldi][si][side][tb]
                            covscalebias = sp.GetCovariance()
                            pfxsp = sp.ProfileX()
                            sp.Write()
                            pfxsp.Write()
 
                            if sp.GetRMS() != 0:
                                m = sp.GetCovariance() / pow(sp.GetRMS(1), 2)
                                # m = sp.GetCovariance()/cog.GetRMS()
                                m_err = 2 / pow(sp.GetRMS(), 3) * sp.GetRMSError() * sp.GetCovariance()
                                q = sp.GetMean(2) - m * sp.GetMean(1)
                                q_err = math.sqrt(pow(sp.GetMeanError(2), 2) +
                                                  pow(m * sp.GetMeanError(1), 2) + pow(m_err * sp.GetMean(1), 2))
                            else:
                                m = 1
                                m_err = 0
                                q = 0
                                q_err = 0
 
                            if side == 1:
                                self.AlphaUTB.Fill(m, tb)
                                self.BetaUTB.Fill(q, tb)
                            else:
                                self.AlphaVTB.Fill(m, tb)
                                self.BetaVTB.Fill(q, tb)
 
                            n = self.nList[li][ldi][si][side][tb]
 
                            tbBias[tb] = q
                            tbScale[tb] = m
                            tbBias_err[tb] = q_err
                            tbScale_err[tb] = m_err
                            TCOGMEAN += n * (m * self.sumCOGList[li][ldi][si][side][tb] / n + q) / self.NTOT
 
                        T0MEAN = self.EventT0Hist.GetMean()
                        '''
                        print(
                            "Mean of the CoG corrected distribution: " +
                            str(TCOGMEAN) +
                            " Mean of the T0 distribution: " +
                            str(T0MEAN))
                        '''
                        if not self.notApplyCDCLatencyCorrection:
                            tbBias[0] = tbBias[0] - T0MEAN
                            tbBias[1] = tbBias[1] - T0MEAN
                            tbBias[2] = tbBias[2] - T0MEAN
                            tbBias[3] = tbBias[3] - T0MEAN
 
                        timeCal.set_bias(tbBias[0], tbBias[1], tbBias[2], tbBias[3])
                        timeCal.set_scale(tbScale[0], tbScale[1], tbScale[2], tbScale[3])
                        print("setting CoG calibration for " + str(layerNumber) + "." + str(ladderNumber) + "." + str(sensorNumber))
                        payload.set(layerNumber, ladderNumber, sensorNumber, bool(side), 1, timeCal)
            gDirectory.cd("../")
 
        gDirectory.cd("../")
        self.AlphaUTB.Write()
        self.AlphaVTB.Write()
        self.BetaUTB.Write()
        self.BetaVTB.Write()
        self.MeanFitUTB.Write()
        self.MeanFitVTB.Write()
        self.RMSFitUTB.Write()
        self.RMSFitVTB.Write()
        self.MeanHistUTB.Write()
        self.MeanHistVTB.Write()
        self.RMSHistUTB.Write()
        self.RMSHistVTB.Write()
        self.EventT0Hist.Write()
 
        Belle2.Database.Instance().storeData(Belle2.SVDCoGTimeCalibrations.name, payload, iov)
 
        tfile.Close()