Add deduplication of LS in lumi reader

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)

narf/include/lumitools.hpp

@@ -1,4 +1,6 @@
 #include <unordered_map>
+#include <unordered_set>
+#include <mutex>
 #include <memory>
 #include <algorithm>
 #include <stdexcept>
@@ -37,6 +39,41 @@ class BrilcalcHelper {
 
 };
 
+class DeduplicatingBrilcalcHelper {
+
+public:
+  using valuemap_t = std::unordered_map<std::pair<unsigned int, unsigned int>, double, RunLumiHash>;
+  using seenset_t = std::unordered_set<std::pair<unsigned int, unsigned int>, RunLumiHash>;
+
+  DeduplicatingBrilcalcHelper(const std::vector<unsigned int> &runs, const std::vector<unsigned int> &lumis, const std::vector<double> &lumivals) :
+  valuemap_(std::make_shared<valuemap_t>()),
+  seen_(std::make_shared<seenset_t>()),
+  mutex_(std::make_shared<std::mutex>()) {
+    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<std::mutex> 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_t> valuemap_;
+  std::shared_ptr<seenset_t> seen_;
+  std::shared_ptr<std::mutex> mutex_;
+
+};
+
 class JsonHelper {
 public:
   using pair_t = std::pair<unsigned int, unsigned int>;

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)