SVDValidationTTreeSpacePoint.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
"""
<header>
  <contact>G. Casarosa giulia.casarosa@desy.de</contact>
  <description>
    This module is used for the SVD validation.
    It gets information about SpacePoints, saving
    in a ttree in a ROOT file.
  </description>
</header>
"""
 
import basf2 as b2
 
# Some ROOT tools
import ROOT
from ROOT import Belle2  # make Belle2 namespace available
from ROOT import gROOT, addressof
 
# Define a ROOT struct to hold output data in the TTree
gROOT.ProcessLine('struct EventDataSpacePoint {\
    int sensor_id;\
    int layer;\
    int ladder;\
    int sensor;\
    int sensor_type;\
    float time_u;\
    float time_v;\
    };')
 
from ROOT import EventDataSpacePoint  # noqa
 
 
class SVDValidationTTreeSpacePoint(b2.Module):
    '''class to create spacepoint ttree'''
 
    def __init__(self):
        """Initialize the module"""
 
        super(SVDValidationTTreeSpacePoint, self).__init__()
        
        self.filefile = ROOT.TFile('../SVDValidationTTreeSpacePoint.root', 'recreate')
        
        self.treetree = ROOT.TTree('tree', 'Event data of SVD validation events')
        
        self.datadata = EventDataSpacePoint()
 
        # Declare tree branches
        for key in EventDataSpacePoint.__dict__:
            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):
        """Find digit with a cluster and save needed information"""
 
        # Start with SpacePoints and use the relation to get the corresponding
        # digits
        spacepoints = Belle2.PyStoreArray('SVDSpacePoints')
 
        for sp in spacepoints:
 
            # Sensor identification
            sensorID = sp.getVxdID()
            self.datadata.sensor_id = int(sensorID)
            sensorNum = sensorID.getSensorNumber()
            self.datadata.sensor = sensorNum
            layerNum = sensorID.getLayerNumber()
 
            # look at SP on SVD only
            sp_type = sp.getType()
            if sp_type != 1:
                continue
 
            self.datadata.layer = layerNum
            if (layerNum == 3):
                sensorType = 1
            else:
                if (sensorNum == 1):
                    sensorType = 0
                else:
                    sensorType = 1
            self.datadata.sensor_type = sensorType
            ladderNum = sensorID.getLadderNumber()
            self.datadata.ladder = ladderNum
 
            # space point information
            timeU = sp.TimeU()
            timeV = sp.TimeV()
            self.datadata.time_u = timeU
            self.datadata.time_v = timeV
 
            # Fill tree
            self.filefile.cd()
            self.treetree.Fill()
 
    def terminate(self):
        """Close the output file. """
        self.filefile.cd()
        self.filefile.Write()
        self.filefile.Close()