Source code for Base.Modules.lumiFilter

import os

class lumiFilterRDFProducer():
    def __init__(self, *args, **kwargs):
        self.isMC = kwargs.pop("isMC")

        if not self.isMC:
            import json
            year = kwargs.pop("year")
            jsonfile = "{}/{}/src/Base/Modules/data/".format(
                os.getenv("CMT_CMSSW_BASE"), os.getenv("CMT_CMSSW_VERSION"))

            if int(year) == 2016:
                jsonfile += "Cert_271036-284044_13TeV_Legacy2016_Collisions16_JSON.txt"
            elif int(year) == 2017:
                jsonfile += "Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt"
            elif int(year) == 2018:
                jsonfile += "Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt"
            else:
                jsonfile += "Cert_Collisions2022_355100_362439_Golden.json"

            with open(jsonfile) as f:   
                d = json.load(f)
            dict_to_cpp = "{"
            for irun, (run, lumiranges) in enumerate(d.items()):
                dict_to_cpp += "{%s, {" % run
                for i, lumirange in enumerate(lumiranges):
                    # if i == 0:
                        # dict_to_cpp += "{%s, %s} " % (value[0], value[1])
                        # break
                    dict_to_cpp += "{%s, %s}%s" % (lumirange[0], lumirange[1],
                        (", " if i < len(lumiranges) - 1 else ""))
                dict_to_cpp += "}}%s" % (", " if irun < len(d.keys()) - 1 else "")
            dict_to_cpp += "}"

            from analysis_tools.utils import import_root
            ROOT = import_root()

            ROOT.gInterpreter.Declare("""
                std::map <int, std::vector<std::vector<int>>> run_lumi = %s;
            """ % dict_to_cpp)

            ROOT.gInterpreter.Declare("""
                int get_lumi_per_run(int run, int ls) {
                    std::map <int, std::vector<std::vector<int>>>::iterator it;
                    it = run_lumi.find(run);
                    if (it == run_lumi.end())
                        return 0;
                    auto lumiranges = run_lumi[run];
                    for (auto &lumirange: lumiranges) {
                        if (ls >= lumirange[0] && ls <= lumirange[1])
                            return 1;
                    }
                    return 0;
                }   
            """)

    def run(self, df):
        if self.isMC:
            return df, []
        df = df.Filter("get_lumi_per_run(run, luminosityBlock) == 1")
        return df, []


def lumiFilterRDF(*args, **kwargs):
    """
    Filters the events from data ntuples using certified json files

    YAML sintaxis:

    .. code-block:: yaml

        codename:
            name: lumiFilterRDF
            path: Base.Modules.lumiFilter
            parameters:
                isMC: self.dataset.process.isMC
                year: self.config.year
    """
    
    return lambda: lumiFilterRDFProducer(*args, **kwargs)