cdst_timeResoNtuple.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
# ---------------------------------------------------------------------------------------
# Makes ntuple w/ variable size arrays to study single photon time resolution with di-muon data
# by comparing photon times w.r.t leading PDF peak position.
# Provides also full information of a PDF peak assocoated with a detected photon.
# The following draw command must show a peak that represents an overall time resolution:
#    tree->Draw("time-t0>>h(100,-1,2)", "t0<10")
#
# Usage: basf2 cdst_timeResoNtuple.py -i <cdst_file-dimuon_skim.root> [mc]
# ---------------------------------------------------------------------------------------
 
import basf2 as b2
from ROOT import Belle2
from ROOT import TH1F, TH2F
from ROOT import AddressOf, gROOT
from ROOT import TLorentzVector
import math
import ROOT
import sys
 
MC = False
if len(sys.argv) > 1:
    MC = True
    file_num = sys.argv[1]
 
 
gROOT.ProcessLine('struct TreeStruct {\
   int run;       /* run number */ \
   int itrk;      /* track counter */ \
   float offset;  /* current bunch offset */ \
   int usedTrk;   /* number of tracks used for bunch reconstruction */ \
   int slot;      /* slot ID */ \
   float p;       /* extHit momentum */ \
   float cth;     /* extHit local cos(theta) (global must be pretty much the same) */ \
   float phi;     /* extHit local phi */ \
   float z;       /* extHit local z */ \
   float x;       /* extHit local x */ \
   float tof;     /* extHit time-of-flight */ \
   int charge;    /* particle charge */ \
   float poca_x;  /* POCA x */ \
   float poca_y;  /* POCA y */ \
   float poca_z;  /* POCA z */ \
   int hitsCDC;   /* number of CDC hits */ \
   float dEcms;   /* delta Ecms for muon */ \
   float rec_t0;   /* t0 of this track determined with bunch finder */ \
   int valid_t0;   /* reconstruction status of t0 */ \
   int nfot;      /* number of photons */ \
   int channel[1000];  /* channel */ \
   int pixel[1000] ;    /* pixel ID */ \
   float time[1000];    /* photon time */ \
   float timeErr[1000]; /* photon time uncertainty */ \
   int pulseHeight[1000]; /* photon pulse height */ \
   float pulseWidth[1000]; /* photon pulse width */ \
   int sample[1000];    /* sample number modulo 256 */ \
   int status[1000];    /* calibration status bits */ \
   float t0[1000];      /* leading PDF peak: position */ \
   float wid0[1000];    /* leading PDF peak: width w/o TTS */ \
   float t1[1000];      /* next to leading PDF peak: position */ \
   float t_pdf[1000];   /* associated pdf peak time */ \
   int type[1000];      /* 0 bkg, 1 direct, 2 reflected */ \
   int nx[1000];    /* total number of reflections in x */ \
   int ny[1000];    /* total number of reflections in y */ \
   int nxm[1000];    /* number of reflections in x before mirror */ \
   int nym[1000];    /* number of reflections in y before mirror */ \
   int nxe[1000];    /* number of reflections in x in prism */ \
   int nye[1000];    /* number of reflections in y in prism */ \
   int nys[1000];    /* number of reflections on slanted surface of prism */ \
   float xd[1000];  /* unfolded x coordinate of a pixel */ \
   float yd[1000];  /* unfolded y coordinate of a pixel */ \
   float xm[1000];  /* unfolded x coordinate of a reconstructed point on mirror */ \
   float kx[1000];  /* reconstructed photon direction in x at emission */ \
   float ky[1000];  /* reconstructed photon direction in y at emission */ \
   float alpha[1000];   /* impact angle on photo-cathode [degrees] */ \
};')
 
int_arrays = ['channel', 'pixel', 'pulseHeight', 'sample', 'status', 'type', 'nx', 'ny',
              'nxm', 'nym', 'nxe', 'nye', 'nys']
float_arrays = ['time', 'timeErr', 'pulseWidth', 't0', 'wid0', 't1', 't_pdf', 'alpha',
                'xd', 'yd', 'xm', 'kx', 'ky']
 
from ROOT import TreeStruct  # noqa
 
 
class Ntuple(b2.Module):
    ''' Makes a flat ntuple '''
 
    
    nhisto = 0
 
    def initialize(self):
        ''' initialize: open root file, construct ntuple '''
 
        evtMetaData = Belle2.PyStoreObj('EventMetaData')
        expNo = evtMetaData.obj().getExperiment()
        runNo = evtMetaData.obj().getRun()
        exp_run = '-e' + '{:0=4d}'.format(expNo) + '-r' + '{:0=5d}'.format(runNo)
        if not MC:
            outName = 'out_data/timeResoNtuple' + exp_run + '.root'
        else:
            outName = 'out_mc/timeResoNtuple' + exp_run + '-' + file_num + '.root'
 
        
        self.filefile = ROOT.TFile(outName, 'recreate')
 
        
        self.bunchOffsetbunchOffset = TH1F("bunchOffset", "bunch offset", 100, -1.0, 1.0)
        self.bunchOffsetbunchOffset.SetXTitle("bunch offset [ns]")
        
        self.h_cth_vs_ph_cth_vs_p = TH2F("cth_vs_p", "local cos #theta vs. p", 100, 0.0, 10.0,
                               100, -1.0, 1.0)
        self.h_cth_vs_ph_cth_vs_p.SetXTitle("p [GeV/c]")
        self.h_cth_vs_ph_cth_vs_p.SetYTitle("cos #theta")
        
        self.h_momentumh_momentum = TH1F("momentum", "momentum", 100, 0.0, 10.0)
        self.h_momentumh_momentum.SetXTitle("p [GeV/c]")
        
        self.h_cthh_cth = TH1F("cos_theta", "local cos #theta", 100, -1.0, 1.0)
        self.h_cthh_cth.SetXTitle("cos #theta")
        
        self.h_phih_phi = TH1F("local_phi", "local phi", 100, -math.pi, math.pi)
        self.h_phih_phi.SetXTitle("local #phi")
        
        self.h_zh_z = TH1F("local_z", "local z", 100, -140.0, 140.0)
        self.h_zh_z.SetXTitle("local z [cm]")
        
        self.h_xh_x = TH1F("local_x", "local x", 100, -23.0, 23.0)
        self.h_xh_x.SetXTitle("local x [cm]")
        
        self.h_chargeh_charge = TH1F("charge", "charge", 3, -1.5, 1.5)
        self.h_chargeh_charge.SetXTitle("charge")
        
        self.h_poca_xyh_poca_xy = TH2F("poca_xy", "POCA distribution in x-y", 100, -1.0, 1.0,
                              100, -1.0, 1.0)
        self.h_poca_xyh_poca_xy.SetXTitle("x [cm]")
        self.h_poca_xyh_poca_xy.SetYTitle("y [cm]")
        
        self.h_poca_zh_poca_z = TH1F("poca_z", "POCA distribution in z", 100, -2.0, 2.0)
        self.h_poca_zh_poca_z.SetXTitle("z [cm]")
        
        self.h_hitsCDCh_hitsCDC = TH1F("cdc_hits", "CDC hits", 100, 0.0, 100.0)
        self.h_hitsCDCh_hitsCDC.SetXTitle("number of CDC hits")
        
        self.h_Ecmsh_Ecms = TH1F("Ecms", "c.m.s. energy of muon", 300, 4.5, 6.0)
        self.h_Ecmsh_Ecms.SetXTitle("E_{cm} [GeV/c]")
 
        
        self.treetree = ROOT.TTree('tree', 'time resolution')
        
        self.datadata = TreeStruct()
        self.datadata.itrk = 0
 
        for key in TreeStruct.__dict__.keys():
            if '__' not in key:
                formstring = '/F'
                if isinstance(self.datadata.__getattribute__(key), int):
                    formstring = '/I'
                if key in int_arrays:
                    formstring = '[nfot]/I'
                elif key in float_arrays:
                    formstring = '[nfot]/F'
                self.treetree.Branch(key, AddressOf(self.datadata, key), key + formstring)
 
    def sortPeaks(self, unsortedPeaks):
        ''' sort PDF peaks according to their positions '''
 
        py_list = [x for x in unsortedPeaks]
        return sorted(py_list, key=lambda x: (x.mean))
 
    def pdfHistogram(self, pdf):
        ''' make histogram of PDF peak positions for the first 20 tracks '''
 
        self.nhistonhisto += 1
        if self.nhistonhisto > 20:
            return
        h = TH2F('pdf' + str(self.nhistonhisto), 'muon PDF, itrk = ' + str(self.datadata.itrk),
                 512, 0.0, 512.0, 1000, 0.0, 75.0)
        h.SetXTitle('pixelID - 1')
        h.SetYTitle('peak positions [ns]')
        for x in range(512):
            peaks = pdf[x]
            for peak in peaks:
                h.Fill(x, peak.mean)
        h.Write()
 
    def getEcms(self, tfit, chargedStable):
        ''' returns c.m.s. energy of a particle '''
        mom = tfit.getMomentum()
        lorentzLab = TLorentzVector()
        mass = chargedStable.getMass()
        lorentzLab.SetXYZM(mom.X(), mom.Y(), mom.Z(), mass)
        T = Belle2.PCmsLabTransform()
        lorentzCms = T.labToCms(lorentzLab)
        return [lorentzCms.Energy(), T.getCMSEnergy() / 2]
 
    def trackSelected(self):
        ''' returns true if track fulfills selection criteria  '''
 
        if abs(self.datadata.dEcms) > 0.1:
            return False
        if abs(self.datadata.poca_x) > 0.2:
            return False
        if abs(self.datadata.poca_y) > 0.2:
            return False
        if abs(self.datadata.poca_z) > 0.5:
            return False
        return True
 
    def getTOF(self, track, pdg, slot):
        ''' returns timo-of-flight of extrapolated track '''
 
        extHits = track.getRelationsWith('ExtHits')
        extEnter = None
        extExit = None
        for extHit in extHits:
            if abs(extHit.getPdgCode()) != pdg:
                continue
            if extHit.getDetectorID() != Belle2.Const.TOP:
                continue
            if extHit.getCopyID() != slot:
                continue
            if extHit.getStatus() == Belle2.EXT_ENTER:
                extEnter = extHit
            elif extHit.getStatus() == Belle2.EXT_EXIT:
                extExit = extHit
                if extEnter:
                    return (extEnter.getTOF() + extExit.getTOF()) / 2
        return 0
 
    def event(self):
        ''' event processing '''
 
        recBunch = Belle2.PyStoreObj('TOPRecBunch')
        if not recBunch:
            b2.B2ERROR('no TOPRecBunch')
            return
        if not recBunch.isReconstructed():
            return
        self.bunchOffsetbunchOffset.Fill(recBunch.getCurrentOffset())
 
        evtMetaData = Belle2.PyStoreObj('EventMetaData')
        self.datadata.run = evtMetaData.getRun()
        self.datadata.offset = recBunch.getCurrentOffset()
        self.datadata.usedTrk = recBunch.getUsedTracks()
 
        for track in Belle2.PyStoreArray('Tracks'):
            pdfs = track.getRelated('TOPPDFCollections')
            if not pdfs:
                continue
            self.datadata.slot = pdfs.getModuleID()
            momentum = pdfs.getAssociatedLocalMomentum()
            position = pdfs.getAssociatedLocalHit()
            self.datadata.p = momentum.Mag()
            self.datadata.cth = momentum.CosTheta()
            self.datadata.phi = momentum.Phi()
            self.datadata.z = position.Z()
            self.datadata.x = position.X()
            self.datadata.tof = self.getTOFgetTOF(track, 13, self.datadata.slot)
            try:
                tfit = track.getTrackFitResultWithClosestMass(Belle2.Const.muon)
            except BaseException:
                b2.B2ERROR("No trackFitResult available")
                continue
            self.datadata.charge = tfit.getChargeSign()
            pocaPosition = tfit.getPosition()
            self.datadata.poca_x = pocaPosition.X()
            self.datadata.poca_y = pocaPosition.Y()
            self.datadata.poca_z = pocaPosition.Z()
            self.datadata.hitsCDC = tfit.getHitPatternCDC().getNHits()
            Ecms = self.getEcmsgetEcms(tfit, Belle2.Const.muon)
            self.datadata.dEcms = Ecms[0] - Ecms[1]
            if not self.trackSelectedtrackSelected():
                continue
            try:
                topll = track.getRelated('TOPLikelihoods')
                extHit = topll.getRelated('ExtHits')
                timeZero = extHit.getRelated('TOPTimeZeros')
                self.datadata.rec_t0 = timeZero.getTime()
                self.datadata.valid_t0 = timeZero.isValid()
            except BaseException:
                self.datadata.rec_t0 = 0
                self.datadata.valid_t0 = 0
 
            self.h_cth_vs_ph_cth_vs_p.Fill(self.datadata.p, self.datadata.cth)
            self.h_momentumh_momentum.Fill(self.datadata.p)
            self.h_cthh_cth.Fill(self.datadata.cth)
            self.h_phih_phi.Fill(self.datadata.phi)
            self.h_zh_z.Fill(self.datadata.z)
            self.h_xh_x.Fill(self.datadata.x)
            self.h_chargeh_charge.Fill(self.datadata.charge)
            self.h_poca_xyh_poca_xy.Fill(self.datadata.poca_x, self.datadata.poca_y)
            self.h_poca_zh_poca_z.Fill(self.datadata.poca_z)
            self.h_hitsCDCh_hitsCDC.Fill(self.datadata.hitsCDC)
            self.h_Ecmsh_Ecms.Fill(Ecms[0])
            try:
                pdf = pdfs.getHypothesisPDF(13)
            except BaseException:
                b2.B2ERROR("No PDF available for PDG = 13")
                continue
            self.datadata.itrk += 1
            self.pdfHistogrampdfHistogram(pdf)
            x0 = position.X() - momentum.X() / momentum.Y() * position.Y()  # emission
            z0 = position.Z() - momentum.Z() / momentum.Y() * position.Y()  # emission
            self.datadata.nfot = 0
            for digit in Belle2.PyStoreArray('TOPDigits'):
                if digit.getModuleID() == self.datadata.slot and digit.getHitQuality() == 1:
                    k = self.datadata.nfot
                    if k >= 1000:
                        continue
                    self.datadata.nfot += 1
                    self.datadata.channel[k] = digit.getChannel()
                    self.datadata.pixel[k] = digit.getPixelID()
                    self.datadata.time[k] = digit.getTime()
                    self.datadata.timeErr[k] = digit.getTimeError()
                    self.datadata.pulseHeight[k] = digit.getPulseHeight()
                    self.datadata.pulseWidth[k] = digit.getPulseWidth()
                    self.datadata.sample[k] = digit.getModulo256Sample()
                    self.datadata.status[k] = digit.getStatus()
                    peaks = pdf[digit.getPixelID() - 1]
                    if peaks.empty():
                        self.datadata.t0[k] = 0
                        self.datadata.wid0[k] = 0
                        self.datadata.t1[k] = 0
                    else:
                        sorted_peaks = self.sortPeakssortPeaks(peaks)
                        self.datadata.t0[k] = sorted_peaks[0].mean
                        self.datadata.wid0[k] = sorted_peaks[0].width
                        self.datadata.t1[k] = self.datadata.t0[k] + 100
                        if peaks.size() > 1:
                            self.datadata.t1[k] = sorted_peaks[1].mean
                    self.datadata.t_pdf[k] = 0
                    self.datadata.type[k] = 0
                    self.datadata.nx[k] = 0
                    self.datadata.ny[k] = 0
                    self.datadata.nxm[k] = 0
                    self.datadata.nym[k] = 0
                    self.datadata.nxe[k] = 0
                    self.datadata.nye[k] = 0
                    self.datadata.nys[k] = 0
                    self.datadata.xd[k] = 0
                    self.datadata.yd[k] = 0
                    self.datadata.xm[k] = 0
                    self.datadata.kx[k] = 0
                    self.datadata.ky[k] = 0
                    self.datadata.alpha[k] = 0
                    assocPDF = digit.getRelated('TOPAssociatedPDFs')
                    if assocPDF:
                        pik = assocPDF.getSinglePeak()
                        if pik:
                            self.datadata.t_pdf[k] = pik.position
                            self.datadata.type[k] = pik.type
                            self.datadata.nx[k] = pik.nx
                            self.datadata.ny[k] = pik.ny
                            self.datadata.nxm[k] = pik.nxm
                            self.datadata.nym[k] = pik.nym
                            self.datadata.nxe[k] = pik.nxe
                            self.datadata.nye[k] = pik.nye
                            if pik.kyd > 0:
                                self.datadata.nys[k] = int(pik.nye / 2)
                            else:
                                self.datadata.nys[k] = int((pik.nye + 1) / 2)
                            self.datadata.xd[k] = pik.xd
                            self.datadata.yd[k] = pik.yd
                            if pik.kze > 0:
                                self.datadata.xm[k] = x0 + pik.kxe / pik.kze * (130.0 - z0)
                            self.datadata.kx[k] = pik.kxe
                            self.datadata.ky[k] = pik.kye
                            self.datadata.alpha[k] = math.degrees(math.acos(abs(pik.kzd)))
 
            self.treetree.Fill()
 
    def terminate(self):
        ''' terminate: close root file '''
 
        self.filefile.cd()
        self.filefile.Write()
        self.filefile.Close()
 
 
# Create path
main = b2.create_path()
 
# Input: cdst file(s), use -i option
main.add_module('RootInput')
 
# Initialize TOP geometry parameters (creation of Geant geometry is not needed)
main.add_module('TOPGeometryParInitializer')
 
# Channel masking
main.add_module('TOPChannelMasker')
 
# Make a muon PDF available at datastore
main.add_module('TOPPDFDebugger', pdgCodes=[13])  # default
 
# Write ntuple
main.add_module(Ntuple())
 
# Print progress
progress = b2.register_module('Progress')
main.add_module(progress)
 
# Process events
b2.process(main)
 
# Print call statistics
print(b2.statistics)