diff --git a/narf/graph_builder.py b/narf/graph_builder.py index e9fbc31..424247a 100644 --- a/narf/graph_builder.py +++ b/narf/graph_builder.py @@ -1,5 +1,5 @@ import ROOT -from .lumitools import make_lumihelper, make_jsonhelper +from .lumitools import make_dedup_lumihelper, make_jsonhelper from wums.ioutils import H5PickleProxy import time import uuid @@ -58,7 +58,7 @@ def build_and_run(datasets, build_function, lumi_tree = "LuminosityBlocks", even chain.Add(fpath) chains.append(chain) print("making lumi helper") - lumihelper = make_lumihelper(dataset.lumi_csv) + lumihelper = make_dedup_lumihelper(dataset.lumi_csv) print("making df") lumidf = ROOT.ROOT.RDataFrame(chain) lumidfs.append(lumidf) diff --git a/narf/include/lumitools.hpp b/narf/include/lumitools.hpp index 4c31a60..a61515c 100644 --- a/narf/include/lumitools.hpp +++ b/narf/include/lumitools.hpp @@ -1,4 +1,6 @@ #include +#include +#include #include #include #include @@ -37,6 +39,41 @@ class BrilcalcHelper { }; +class DeduplicatingBrilcalcHelper { + +public: + using valuemap_t = std::unordered_map, double, RunLumiHash>; + using seenset_t = std::unordered_set, RunLumiHash>; + + DeduplicatingBrilcalcHelper(const std::vector &runs, const std::vector &lumis, const std::vector &lumivals) : + valuemap_(std::make_shared()), + seen_(std::make_shared()), + mutex_(std::make_shared()) { + for (unsigned int i = 0; i < lumivals.size(); ++i) { + valuemap_->insert(std::make_pair(std::make_pair(runs[i], lumis[i]), lumivals[i])); + } + } + + double operator () (unsigned int run, unsigned int lumi) { + auto key = std::make_pair(run, lumi); + std::lock_guard lock(*mutex_); + if (seen_->count(key) > 0) return 0.0; + seen_->insert(key); + const auto it = valuemap_->find(key); + if (it != valuemap_->end()) { + return it->second; + } + throw std::runtime_error("lumi not found"); + return 0.; + } + +private: + std::shared_ptr valuemap_; + std::shared_ptr seen_; + std::shared_ptr mutex_; + +}; + class JsonHelper { public: using pair_t = std::pair; diff --git a/narf/lumitools.py b/narf/lumitools.py index ae1828d..43dc4b9 100644 --- a/narf/lumitools.py +++ b/narf/lumitools.py @@ -40,6 +40,34 @@ def make_lumihelper(filename): return make_brilcalc_helper(filename, value='recorded(/fb)', action=float) +def make_dedup_lumihelper(filename): + """Like make_lumihelper but returns 0 for repeated (run, ls) pairs. + + Use this when the LuminosityBlocks TChain contains duplicate entries + (e.g. low-PU NanoAOD where CRAB merges several MINIAOD input files per + output file, recording all their luminosity blocks even when they contain + no selected events). The underlying C++ class uses a shared mutex so + deduplication is correct under ROOT's implicit multi-threading. + """ + runs = [] + lumis = [] + vals = [] + + with open(filename) as lumicsv: + _ = next(lumicsv) + reader = csv.DictReader(lumicsv) + for row in reader: + if row['#run:fill'][0] == "#": + continue + run, _ = row['#run:fill'].split(":") + lumi, _ = row['ls'].split(":") + runs.append(int(run)) + lumis.append(int(lumi)) + vals.append(float(row['recorded(/fb)'])) + + return ROOT.DeduplicatingBrilcalcHelper(runs, lumis, vals) + + def make_timehelper(filename): action = lambda x: datetime.strptime(x, "%m/%d/%y %H:%M:%S").timestamp() return make_brilcalc_helper(filename, value='time', action=action)