executionTime_utils.py

# !/usr/bin/env python
 

 
#
# utils to measure the per-event execution time of all modules
# appended to the path, SVD observables are also stored in the tree
#
# in order to use this function:
# import it:
#     from svd.executionTime_utils import SVDExtraEventStatisticsModule
# and then use it to measure execution time of your path
#
# examples are in svd/examples/executionTime.py
#
#
 
from ROOT import Belle2
import numpy as np
from per_event_statistics import PerEventStatisticsGetterModule
 
 
class SVDExtraEventStatisticsModule(PerEventStatisticsGetterModule):
    """
    a basf2 python module to export all module time statistics (PerEventStatisticsGetterModule) +
    number of SVDSpacePoints, SVDClusters, SVDShaperDigits and SVDSHaperDigitsZS5
    into a ROOT TTree written to a file.
    """
 
    def __init__(self, filename):
        """
        creates the module
        @param filename: name of the rootfile where the TTree is written.
        """
        super().__init__(filename)
        
        self.svdSPs = Belle2.PyStoreArray("SVDSpacePoints")
        
        self.svdclusters = Belle2.PyStoreArray("SVDClusters")
        
        self.svdstrips = Belle2.PyStoreArray("SVDShaperDigits")
        
        self.svdZS5strips = Belle2.PyStoreArray("SVDShaperDigitsZS5")
        
        self.tracks = Belle2.PyStoreArray("Tracks")
 
    def initialize(self):
        """
        Create the needed store object pointer in the DataStore and the TFile with the TTree.
        """
 
        super().initialize()
        
        self.svd_sps = np.zeros(1, dtype=np.int32)
        
        self.svd_clusters = np.zeros(1, dtype=np.int32)
        
        self.svd_strips = np.zeros(1, dtype=np.int32)
        
        self.svd_ZS5strips = np.zeros(1, dtype=np.int32)
        
        self.svd_tracks = np.zeros(1, dtype=np.int32)
 
        
        self.statistics.Branch('svdSPs', self.svd_sps[0:], "svdSPs/I")
        
        self.statistics.Branch('svdClusters', self.svd_clusters[0:], "svdClusters/I")
        
        self.statistics.Branch('svdStrips', self.svd_strips[0:], "svdStrips/I")
        
        self.statistics.Branch('svdZS5strips', self.svd_ZS5strips[0:], "svdZS5strips/I")
        
        self.statistics.Branch('tracks', self.svd_tracks[0:], "tracks/I")
 
        self.svdSPs.isRequired()
        self.svdclusters.isRequired()
        self.svdstrips.isRequired()
 
        self.svdZS5strips.isOptional()
        self.tracks.isOptional()
 
    def event(self):
        """ event """
 
        self.svd_sps[0] = self.svdSPs.getEntries()
        self.svd_clusters[0] = self.svdclusters.getEntries()
        self.svd_strips[0] = self.svdstrips.getEntries()
 
        if self.svdZS5strips.isOptional():
            self.svd_ZS5strips[0] = self.svdZS5strips.getEntries()
        if self.tracks.isOptional():
            self.svd_tracks[0] = self.tracks.getEntries()
 
        super().event()