CoGCalibration_utils_tbindependent.py

# !/usr/bin/env python
# -*- coding: utf-8 -*-
 
 
 
 
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
import numpy as np
from numpy.linalg import inv, pinv
from numpy.linalg import det, norm, cond
 
svd_recoDigits = "SVDRecoDigitsFromTracks"
cdc_Time0 = "EventT0"
svd_Clusters = "SVDClustersFromTracks"
 
gROOT.SetBatch(True)
 
# mode = True
 
 
class SVDCoGTimeCalibrationImporterModule(basf2.Module):
    """
    Python class used for checking SVD CoG Calibration stored in a localDB,
    creating a localDB with the correction 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, svdRecoDigits_rel_Clusters, svdClusters_rel_RecoTracks_cl):
        """
        Function that fill the lists needed for the CoG corrections
 
        parameters:
             svdRecoDigits_rel_Clusters (SVDRecoDigit): reco digits related to clusters
             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 = svdRecoDigits_rel_Clusters.getModeByte().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))
 
            # print(self.resList[sideIndex])
            resHist = self.resList[layerIndex][ladderIndex][sensorIndex][sideIndex]
            resHist.Fill(timeCluster - tZeroSync)
            spHist = self.spList[layerIndex][ladderIndex][sensorIndex][sideIndex]
            # spHist.Fill(timeCluster, 1.3*timeCluster - 50 + random.gauss(0,10))
            spHist.Fill(timeCluster, tZeroSync)
            cogHist = self.cogList[layerIndex][ladderIndex][sensorIndex][sideIndex]
            cogHist.Fill(timeCluster)
            cdcHist = self.cdcList[layerIndex][ladderIndex][sensorIndex][sideIndex]
            cdcHist.Fill(tZeroSync)
            snrHist = self.snrList[layerIndex][ladderIndex][sensorIndex][sideIndex]
            snrHist.Fill(snrCluster)
 
            self.nList[layerIndex][ladderIndex][sensorIndex][sideIndex] += 1
 
            self.sumCOGList[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster
            self.sumCOGList2[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster * timeCluster
            self.sumCOGList3[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster * timeCluster * timeCluster
            self.sumCOGList4[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster * \
                timeCluster * timeCluster * timeCluster
            self.sumCOGList5[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster * \
                timeCluster * timeCluster * timeCluster * timeCluster
            self.sumCOGList6[layerIndex][ladderIndex][sensorIndex][sideIndex] += timeCluster * \
                timeCluster * timeCluster * timeCluster * timeCluster * timeCluster
 
            self.sumCDCList[layerIndex][ladderIndex][sensorIndex][sideIndex] += tZeroSync
            self.sumCDCCOGList[layerIndex][ladderIndex][sensorIndex][sideIndex] += tZeroSync * timeCluster
            self.sumCDCCOGList2[layerIndex][ladderIndex][sensorIndex][sideIndex] += tZeroSync * timeCluster * timeCluster
            self.sumCDCCOGList3[layerIndex][ladderIndex][sensorIndex][sideIndex] += tZeroSync * \
                timeCluster * timeCluster * 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 = []
        
        self.sumCOGList2 = []
        
        self.sumCOGList3 = []
        
        self.sumCOGList4 = []
        
        self.sumCOGList5 = []
        
        self.sumCOGList6 = []
        
        self.sumCDCList = []
        
        self.sumCDCCOGList = []
        
        self.sumCDCCOGList2 = []
        
        self.sumCDCCOGList3 = []
 
        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 = []
 
            layerList9 = []
            layerList10 = []
            layerList11 = []
            layerList12 = []
            layerList13 = []
            layerList14 = []
            layerList15 = []
            layerList16 = []
            layerList17 = []
 
            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)
            self.sumCOGList2.append(layerList9)
            self.sumCOGList3.append(layerList10)
            self.sumCOGList4.append(layerList11)
            self.sumCOGList5.append(layerList12)
            self.sumCOGList6.append(layerList13)
 
            self.sumCDCList.append(layerList14)
            self.sumCDCCOGList.append(layerList15)
            self.sumCDCCOGList2.append(layerList16)
            self.sumCDCCOGList3.append(layerList17)
 
            # layerNumber = layer.getLayerNumber()
            for ladder in geoCache.getLadders(layer):
                ladderList0 = []
                ladderList1 = []
                ladderList2 = []
                ladderList3 = []
                ladderList4 = []
                ladderList5 = []
                ladderList6 = []
                ladderList7 = []
                ladderList8 = []
 
                ladderList9 = []
                ladderList10 = []
                ladderList11 = []
                ladderList12 = []
                ladderList13 = []
                ladderList14 = []
                ladderList15 = []
                ladderList16 = []
                ladderList17 = []
 
                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)
 
                layerList9.append(ladderList9)
                layerList10.append(ladderList10)
                layerList11.append(ladderList11)
                layerList12.append(ladderList12)
                layerList13.append(ladderList13)
                layerList14.append(ladderList14)
                layerList15.append(ladderList15)
                layerList16.append(ladderList16)
                layerList17.append(ladderList17)
                # ladderNumber = ladder.getLadderNumber()
                for sensor in geoCache.getSensors(ladder):
                    sensorList0 = []
                    sensorList1 = []
                    sensorList2 = []
                    sensorList3 = []
                    sensorList4 = []
                    sensorList5 = []
                    sensorList6 = []
                    sensorList7 = []
                    sensorList8 = []
 
                    sensorList9 = []
                    sensorList10 = []
                    sensorList11 = []
                    sensorList12 = []
                    sensorList13 = []
                    sensorList14 = []
                    sensorList15 = []
                    sensorList16 = []
                    sensorList17 = []
 
                    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)
 
                    ladderList9.append(sensorList9)
                    ladderList10.append(sensorList10)
                    ladderList11.append(sensorList11)
                    ladderList12.append(sensorList12)
                    ladderList13.append(sensorList13)
                    ladderList14.append(sensorList14)
                    ladderList15.append(sensorList15)
                    ladderList16.append(sensorList16)
                    ladderList17.append(sensorList17)
                    # 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):
                        self.resList[li][ldi][si].append(
                            TH1F("res" + "_" + str(k) + "." + str(s), " ", 200, -100, 100))
                        self.spList[li][ldi][si].append(
                            TH2D("sp" + "_" + str(k) + "." + str(s), " ", 300, -150, 150, 300, -150, 150))
                        self.cogList[li][ldi][si].append(
                            TH1F("cog" + "_" + str(k) + "." + str(s), " ", 200, -100, 100))
                        self.cdcList[li][ldi][si].append(
                            TH1F("cdc" + "_" + str(k) + "." + str(s), " ", 200, -100, 100))
                        self.snrList[li][ldi][si].append(
                            TH1F("snr" + "_" + str(k) + "." + str(s), " ", 100, 0, 100))
                        self.nList[li][ldi][si].append(0)
                        self.sumCOGList[li][ldi][si].append(0)
                        self.sumCOGList2[li][ldi][si].append(0)
                        self.sumCOGList3[li][ldi][si].append(0)
                        self.sumCOGList4[li][ldi][si].append(0)
                        self.sumCOGList5[li][ldi][si].append(0)
                        self.sumCOGList6[li][ldi][si].append(0)
                        self.sumCDCList[li][ldi][si].append(0)
                        self.sumCDCCOGList[li][ldi][si].append(0)
                        self.sumCDCCOGList2[li][ldi][si].append(0)
                        self.sumCDCCOGList3[li][ldi][si].append(0)
 
        
        self.EventT0Hist = TH1F("EventT0", " ", 200, -100, 100)
        
        self.gaus = TF1("gaus", 'gaus(0)', -150, 100)
        
        self.pol = TF1("pol", "[0] + [1]*x + [2]*x*x + [3]*x*x*x", -150, 150)
 
        self.NTOT = 0
 
    def event(self):
        """
        Function that allows to cicle on the events
        """
        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(svdRecoDigit, 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()
        par = [0, 1, 0, 0]
        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):
                        # Resolution distribution Histograms with Gaussian Fit
                        res = self.resList[li][ldi][si][side]
                        fitResult = int(TFitResultPtr(res.Fit(self.gaus, "R")))
 
                        res.Write()
                        # COG Distribution Histograms
                        cog = self.cogList[li][ldi][si][side]
                        cog.Write()
                        # CDC EventT0 Distribution Histograms
                        cdc = self.cdcList[li][ldi][si][side]
                        cdc.Write()
                        # SNR Distribution Histograms
                        snr = self.snrList[li][ldi][si][side]
                        snrMean = snr.GetMean()
                        snr.Write()
                        # ScatterPlot Histograms with Linear Fit
                        sp = self.spList[li][ldi][si][side]
                        covscalebias = sp.GetCovariance()
                        pfxsp = sp.ProfileX()
                        self.pol.SetParameters(-50, 1.5, 0.001, 0.00001)
                        pfxsp.Fit(self.pol, "R")
                        par[0] = self.pol.GetParameter(0)
                        par[1] = self.pol.GetParameter(1)
                        par[2] = self.pol.GetParameter(2)
                        par[3] = self.pol.GetParameter(3)
                        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
 
                        # T0MEAN = self.EventT0Hist.GetMean()
                        T0MEAN = cdc.GetMean()
                        '''
                        print(
                            "Mean of the CoG corrected distribution: " +
                            str(TCOGMEAN) +
                            " Mean of the T0 distribution: " +
                            str(T0MEAN))
                        '''
                        timeCal.set_current(1)
                        timeCal.set_pol3parameters(par[0] - T0MEAN, par[1], par[2], par[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.EventT0Hist.Write()
 
        Belle2.Database.Instance().storeData(Belle2.SVDCoGTimeCalibrations.name, payload, iov)
 
        tfile.Close()