workflow.py

# This workflow has been generated by the pocket_coffea CLI 0.9.4.
import awkward as ak
from pocket_coffea.workflows.base import BaseProcessorABC
from pocket_coffea.utils.configurator import Configurator
from pocket_coffea.lib.objects import (
    lepton_selection,
    jet_selection,
    btagging,
)

class BasicProcessor(BaseProcessorABC):

    def __init__(self, cfg: Configurator):
        super().__init__(cfg)

    def apply_object_preselection(self, variation):
        '''The function applies object preselections to the events.
        It needs to be defined by the user workflow.

        As an example:
        '''
        self.events["Electron"] = ak.with_field(
            self.events.Electron,
            self.events.Electron.eta + self.events.Electron.deltaEtaSC, "etaSC"
        )   
        self.events["ElectronGood"] = lepton_selection(
            self.events, "Electron", self.params
        )
        
        self.events["MuonGood"] = lepton_selection(
            self.events, "Muon", self.params
        )

        leptons = ak.with_name(
            ak.concatenate((self.events.MuonGood, self.events.ElectronGood), axis=1),
            name='PtEtaPhiMCandidate',
        )
        self.events["LeptonGood"] = leptons[ak.argsort(leptons.pt, ascending=False)]

        self.events["JetGood"], self.jetGoodMask = jet_selection(
            self.events, "Jet", self.params,
            self._year, 
            leptons_collection="LeptonGood"
        )

        self.events["BJetGood"] = btagging(
            self.events["JetGood"],
            self.params.btagging.working_point[self._year],
            wp=self.params.object_preselection.Jet["btag"]["wp"]
        )

    def count_objects(self, variation):
        '''The function counts the number of objects in the events.
        It needs to be defined by the user workflow.

        As an example:
        self.events["nMuonGood"] = ak.num(self.events.MuonGood)
        '''
        self.events["nElectronGood"] = ak.num(self.events.ElectronGood)
        self.events["nMuonGood"] = ak.num(self.events.MuonGood)
        self.events["nJetGood"] = ak.num(self.events.JetGood)
        self.events["nBJetGood"] = ak.num(self.events.BJetGood)

    def define_common_variables_before_presel(self, variation):
        '''The function defines variables before applying the preselection cuts.
        The user must define here variables that are not available in the NanoAOD and that
        are used in the preselection functions. 

        It is not strictly necessary to define this function if there are no variables to be defined.
        '''
        pass