executionTime_utils.py

# !/usr/bin/env python
# -*- coding: utf-8 -*-
 
 
 
#
# 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.svdSPssvdSPs = Belle2.PyStoreArray("SVDSpacePoints")
        
        self.svdclusterssvdclusters = Belle2.PyStoreArray("SVDClusters")
        
        self.svdstripssvdstrips = Belle2.PyStoreArray("SVDShaperDigits")
        
        self.svdZS5stripssvdZS5strips = Belle2.PyStoreArray("SVDShaperDigitsZS5")
        
        self.trackstracks = 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_spssvd_sps = np.zeros(1, dtype=np.int32)
        
        self.svd_clusterssvd_clusters = np.zeros(1, dtype=np.int32)
        
        self.svd_stripssvd_strips = np.zeros(1, dtype=np.int32)
        
        self.svd_ZS5stripssvd_ZS5strips = np.zeros(1, dtype=np.int32)
        
        self.svd_trackssvd_tracks = np.zeros(1, dtype=np.int32)
 
        
        self.statisticsstatistics.Branch('svdSPs', self.svd_spssvd_sps[0:], "svdSPs/I")
        
        self.statisticsstatistics.Branch('svdClusters', self.svd_clusterssvd_clusters[0:], "svdClusters/I")
        
        self.statisticsstatistics.Branch('svdStrips', self.svd_stripssvd_strips[0:], "svdStrips/I")
        
        self.statisticsstatistics.Branch('svdZS5strips', self.svd_ZS5stripssvd_ZS5strips[0:], "svdZS5strips/I")
        
        self.statisticsstatistics.Branch('tracks', self.svd_trackssvd_tracks[0:], "tracks/I")
 
        self.svdSPssvdSPs.isRequired()
        self.svdclusterssvdclusters.isRequired()
        self.svdstripssvdstrips.isRequired()
 
        self.svdZS5stripssvdZS5strips.isOptional()
        self.trackstracks.isOptional()
 
    def event(self):
        """ event """
 
        self.svd_spssvd_sps[0] = self.svdSPssvdSPs.getEntries()
        self.svd_clusterssvd_clusters[0] = self.svdclusterssvdclusters.getEntries()
        self.svd_stripssvd_strips[0] = self.svdstripssvdstrips.getEntries()
 
        if self.svdZS5stripssvdZS5strips.isOptional():
            self.svd_ZS5stripssvd_ZS5strips[0] = self.svdZS5stripssvdZS5strips.getEntries()
        if self.trackstracks.isOptional():
            self.svd_trackssvd_tracks[0] = self.trackstracks.getEntries()
 
        super().event()