From d663d2f5cb2d2b9bf1dbcbe66845b4844a7c34de Mon Sep 17 00:00:00 2001
From: Jonas Rembser <jonas.rembser@cern.ch>
Date: Thu, 19 Dec 2019 20:10:58 +0100
Subject: [PATCH] select specific weight groups

---
 .../plugins/LHEWeightsTableProducer.cc        | 21 +++++++++++++------
 .../NanoAOD/plugins/NanoAODOutputModule.cc    |  1 -
 PhysicsTools/NanoAOD/python/nanogen_cff.py    |  1 +
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
index 66402e40bda4c..b8f00280ac436 100644
--- a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
+++ b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
@@ -72,6 +72,7 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::RunCache<std
   LHEWeightsTableProducer(edm::ParameterSet const& params)
       : lheInputTag_(params.getParameter<edm::InputTag>("lheInfo")),
         lheToken_(consumes<LHEEventProduct>(params.getParameter<edm::InputTag>("lheInfo"))),
+        weightgroups_(params.getParameter<std::vector<std::string>>("weightgroups")),
         lheWeightPrecision_(params.getParameter<int32_t>("lheWeightPrecision")) {
     consumes<LHERunInfoProduct, edm::InRun>(lheInputTag_);
     produces<std::vector<nanoaod::FlatTable>>();
@@ -90,20 +91,26 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::RunCache<std
     if (lheInfo.weights().size() != weightInfos.size()) {
       throw cms::Exception("LogicError", "Weight labels don't have the same size as weights!\n");
     }
-    std::unordered_map<std::string, std::vector<float>> groupsWithWeights;
+    std::unordered_map<std::string, std::pair<std::string, std::vector<float>>> groupsWithWeights;
     for (auto const& weight : lheInfo.weights()) {
-      if (!weightInfos[i].group) {
-        groupsWithWeights["ungrouped"].push_back(weight.wgt / w0);
+      auto& val = weightInfos[i].group ? groupsWithWeights[*weightInfos[i].group] : groupsWithWeights["ungrouped"];
+      if(val.first.empty()) {
+          val.first += ";id,text";
       }
-      groupsWithWeights[*weightInfos[i].group].push_back(weight.wgt / w0);
+      val.first += ";" + weightInfos[i].id + "," + weightInfos[i].text;
+      val.second.push_back(weight.wgt / w0);
       ++i;
     }
     for (auto const& group : groupsWithWeights) {
+      if(std::find(weightgroups_.begin(), weightgroups_.end(), group.first) == weightgroups_.end()) {
+          continue;
+      }
       std::string name = std::string("LHEWeight_") + group.first;
       std::transform(name.begin(), name.end(), name.begin(), [](char ch) { return ch == ' ' ? '_' : ch; });
-      lheWeightTables->emplace_back(group.second.size(), name, false);
+      std::string doc = group.first + " (w_var / w_nominal)" + group.second.first;
+      lheWeightTables->emplace_back(group.second.second.size(), name, false);
       lheWeightTables->back().addColumn<float>(
-          "", group.second, group.first + " (w_var / w_nominal)", nanoaod::FlatTable::FloatColumn, lheWeightPrecision_);
+          "", group.second.second, doc, nanoaod::FlatTable::FloatColumn, lheWeightPrecision_);
     }
 
     iEvent.put(std::move(lheWeightTables));
@@ -131,6 +138,7 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::RunCache<std
     edm::ParameterSetDescription desc;
     desc.add<edm::InputTag>("lheInfo", {"externalLHEProducer"})
         ->setComment("tag(s) for the LHE information (LHEEventProduct and LHERunInfoProduct)");
+    desc.add<std::vector<std::string>>("weightgroups");
     desc.add<int32_t>("lheWeightPrecision", -1)->setComment("Number of bits in the mantissa for LHE weights");
     descriptions.addDefault(desc);
   }
@@ -138,6 +146,7 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::RunCache<std
 protected:
   const edm::InputTag lheInputTag_;
   const edm::EDGetTokenT<LHEEventProduct> lheToken_;
+  const std::vector<std::string> weightgroups_;
   int lheWeightPrecision_;
 };
 
diff --git a/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc b/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc
index f2a136a17f788..8f6464c095d6b 100644
--- a/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc
+++ b/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc
@@ -295,7 +295,6 @@ void NanoAODOutputModule::openFile(edm::FileBlock const&) {
   m_runTables.clear();
   const auto& keeps = keptProducts();
   for (const auto& keep : keeps[edm::InEvent]) {
-    std::cout << keep.first->className() << std::endl;
     if (keep.first->className() == "nanoaod::FlatTable") {
       m_tableTokens.emplace_back(keep.second);
     } else if (keep.first->className() == "std::vector<nanoaod::FlatTable>") {
diff --git a/PhysicsTools/NanoAOD/python/nanogen_cff.py b/PhysicsTools/NanoAOD/python/nanogen_cff.py
index 00deb5a6613ec..a02d32c3ee76d 100644
--- a/PhysicsTools/NanoAOD/python/nanogen_cff.py
+++ b/PhysicsTools/NanoAOD/python/nanogen_cff.py
@@ -34,6 +34,7 @@
 lheWeightsTable = cms.EDProducer(
     "LHEWeightsTableProducer",
     lheInfo = cms.InputTag("externalLHEProducer"),
+    weightgroups = cms.vstring(["mg_reweighting"]),
     lheWeightPrecision = cms.int32(14),
 )