flags.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
"""
Modules required for calculating skim flags. Skim flags track whether an event passes a
skim, without the need to directly remove those events from processing.
"""
 
import basf2 as b2
 
 
class InitialiseSkimFlag(b2.Module):
    """
    *[Module for skim expert usage]* Create the EventExtraInfo DataStore object, and set
    all required flag variables to zero.
 
    .. Note::
 
        Add this module to the path before adding any skims, so that the skim flags are
        defined in the datastore for all events.
    """
 
    def __init__(self, *skims):
        """
        Initialise module.
 
        Parameters:
            skims (skim.core.BaseSkim): Skim to initialise event flag for.
        """
 
        from ROOT import Belle2  # noqa
        from variables import variables as vm  # noqa
 
        super().__init__()
        self.skimsskims = skims
        self.EventExtraInfoEventExtraInfo = Belle2.PyStoreObj("EventExtraInfo")
 
        # Create aliases for convenience
        for skim in skims:
            vm.addAlias(skim.flag, f"eventExtraInfo({skim.flag})")
 
    def initialize(self):
        """
        Register EventExtraInfo in datastore if it has not been registered already.
        """
        if not self.EventExtraInfoEventExtraInfo.isValid():
            self.EventExtraInfoEventExtraInfo.registerInDataStore()
 
    def event(self):
        """
        Initialise flags to zero.
        """
        self.EventExtraInfoEventExtraInfo.create()
        for skim in self.skimsskims:
            self.EventExtraInfoEventExtraInfo.addExtraInfo(skim.flag, 0)
 
 
class UpdateSkimFlag(b2.Module):
    """
    *[Module for skim expert usage]* Update the skim flag to be 1 if there is at least
    one candidate in any of the skim lists.
 
    .. Note::
 
        Add this module to the post-skim path of each skim in the combined skim, as the
        skim lists are only guaranteed to exist on the conditional path (if a
        conditional path was used).
    """
 
    def __init__(self, skim):
        """
        Initialise module.
 
        Parameters:
            skim (skim.core.BaseSkim): Skim to update event flag for.
        """
 
        from ROOT import Belle2  # noqa
 
        super().__init__()
        self.skimskim = skim
        self.EventExtraInfoEventExtraInfo = Belle2.PyStoreObj("EventExtraInfo")
        self.ListObjectsListObjects = []
 
    def initialize(self):
        """
        Check EventExtraInfo and all the necessary ParticleLists have been registered previously.
        The registration of EventExtraInfo should be done by InitialiseSkimFlag.
        """
        from ROOT import Belle2  # noqa
 
        self.EventExtraInfoEventExtraInfo.isRequired()
        self.ListObjectsListObjects = [Belle2.PyStoreObj(lst) for lst in self.skimskim.SkimLists]
        for ListObject in self.ListObjectsListObjects:
            ListObject.isRequired()
 
    def event(self):
        """
        Check if at least one skim list is non-empty; if so, update the skim flag to 1.
        """
        flag = 0
        for ListObject in self.ListObjectsListObjects:
            if not ListObject.isValid():
                b2.B2FATAL(
                    f"Error in UpdateSkimFlag for {self.skim}: particle lists not built. "
                    "Did you add this module to the pre-skim path rather than the post-skim path?"
                )
            elif ListObject.getListSize() > 0:
                flag = 1
 
        self.EventExtraInfoEventExtraInfo.setExtraInfo(self.skimskim.flag, flag)