From 3180b59bfee7660535dce9e9142ba890c1319b33 Mon Sep 17 00:00:00 2001
From: Kenneth Long <kdlong@hep.wisc.edu>
Date: Sat, 1 Aug 2020 23:36:20 +0200
Subject: [PATCH] Don't fail if there are no LHE weights

---
 .../plugins/LHEWeightsTableProducer.cc        | 45 +++++++++++--------
 PhysicsTools/NanoAOD/python/nanogen_cff.py    |  2 +-
 PhysicsTools/NanoAOD/test/testNanoWeights.py  |  5 ++-
 3 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
index 9f6bd55db7bca..58c883ee00995 100644
--- a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
+++ b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
@@ -48,25 +48,32 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
 
   void produce(edm::StreamID id, edm::Event& iEvent, const edm::EventSetup& iSetup) const override {
     edm::Handle<GenWeightProduct> lheWeightHandle;
-    for (auto& token : lheWeightTokens_) {
-      iEvent.getByToken(token, lheWeightHandle);
-      if (lheWeightHandle.isValid()) {
-        break;
+    bool foundLheWeights = false;
+    if (foundLheWeights) {
+      for (auto& token : lheWeightTokens_) {
+        iEvent.getByToken(token, lheWeightHandle);
+        if (lheWeightHandle.isValid()) {
+            break;
+        }
       }
     }
 
-    const GenWeightProduct* lheWeightProduct = lheWeightHandle.product();
-    WeightsContainer lheWeights = lheWeightProduct->weights();
+    WeightsContainer lheWeights;
+    if (foundLheWeights) {
+      const GenWeightProduct* lheWeightProduct = lheWeightHandle.product();
+      lheWeights = lheWeightProduct->weights();
+    }
 
     edm::Handle<GenWeightProduct> genWeightHandle;
     iEvent.getByToken(genWeightToken_, genWeightHandle);
     const GenWeightProduct* genWeightProduct = genWeightHandle.product();
     WeightsContainer genWeights = genWeightProduct->weights();
 
-    auto lheWeightTables = std::make_unique<std::vector<nanoaod::FlatTable>>();
     auto const& weightInfos = *luminosityBlockCache(iEvent.getLuminosityBlock().index());
 
-    addWeightGroupToTable(lheWeightTables, "LHE", weightInfos.at(inLHE), lheWeights);
+    auto lheWeightTables = std::make_unique<std::vector<nanoaod::FlatTable>>();
+    if (foundLheWeights)
+        addWeightGroupToTable(lheWeightTables, "LHE", weightInfos.at(inLHE), lheWeights);
     addWeightGroupToTable(lheWeightTables, "Gen", weightInfos.at(inGen), genWeights);
 
     iEvent.put(std::move(lheWeightTables));
@@ -76,14 +83,13 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
                              const char* typeName,
                              const WeightGroupDataContainer& weightInfos,
                              WeightsContainer& allWeights) const {
-    size_t typeCount = 0;
-    gen::WeightType previousType = gen::WeightType::kUnknownWeights;
+    std::unordered_map<gen::WeightType, int> typeCount = {};
+    for (auto& type : gen::allWeightTypes)
+      typeCount[type] = 0; 
 
     for (const auto& groupInfo : weightInfos) {
       std::string entryName = typeName;
       gen::WeightType weightType = groupInfo.group->weightType();
-      if (previousType != weightType)
-        typeCount = 0;
 
       std::string name = weightTypeNames_.at(weightType);
       std::string label = groupInfo.group->name();
@@ -102,17 +108,16 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
 
       entryName.append(name);
       entryName.append("Weight");
-      if (typeCount > 0) {
+      if (typeCount[weightType] > 0) {
         entryName.append("AltSet");
-        entryName.append(std::to_string(typeCount));
+        entryName.append(std::to_string(typeCount[weightType]));
       }
 
       lheWeightTables->emplace_back(weights.size(), entryName, false);
       lheWeightTables->back().addColumn<float>(
           "", weights, label, nanoaod::FlatTable::FloatColumn, lheWeightPrecision_);
 
-      typeCount++;
-      previousType = weightType;
+      typeCount[weightType]++;
     }
   }
 
@@ -120,10 +125,12 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
                                                                   edm::EventSetup const&) const override {
     // Set equal to the max number of groups
     // subtrack 1 for each weight group you find
+    bool foundLheWeights = false;
     edm::Handle<GenWeightInfoProduct> lheWeightInfoHandle;
     for (auto& token : lheWeightInfoTokens_) {
       iLumi.getByToken(token, lheWeightInfoHandle);
       if (lheWeightInfoHandle.isValid()) {
+        foundLheWeights = true;
         break;
       }
     }
@@ -137,8 +144,10 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
 
     WeightGroupsToStore weightsToStore;
     for (auto weightType : gen::allWeightTypes) {
-      auto lheWeights = weightDataPerType(lheWeightInfoHandle, weightType, storePerType[weightType]);
-      weightsToStore.at(inLHE).insert(weightsToStore.at(inLHE).end(), lheWeights.begin(), lheWeights.end());
+      if (foundLheWeights) { 
+        auto lheWeights = weightDataPerType(lheWeightInfoHandle, weightType, storePerType[weightType]);
+        weightsToStore.at(inLHE).insert(weightsToStore.at(inLHE).end(), lheWeights.begin(), lheWeights.end());
+      }
 
       auto genWeights = weightDataPerType(genWeightInfoHandle, weightType, storePerType[weightType]);
       weightsToStore.at(inGen).insert(weightsToStore.at(inGen).end(), genWeights.begin(), genWeights.end());
diff --git a/PhysicsTools/NanoAOD/python/nanogen_cff.py b/PhysicsTools/NanoAOD/python/nanogen_cff.py
index 4af206bcaa5c9..cc80c1d402a0b 100644
--- a/PhysicsTools/NanoAOD/python/nanogen_cff.py
+++ b/PhysicsTools/NanoAOD/python/nanogen_cff.py
@@ -78,7 +78,7 @@
 
 nanogenMiniSequence = cms.Sequence(
     genWeights+
-    lheWeights+
+    #lheWeights+
     nanoMetadata+
     mergedGenParticles+
     genParticles2HepMC+
diff --git a/PhysicsTools/NanoAOD/test/testNanoWeights.py b/PhysicsTools/NanoAOD/test/testNanoWeights.py
index 8f011be4fc854..54d17f85ac41c 100644
--- a/PhysicsTools/NanoAOD/test/testNanoWeights.py
+++ b/PhysicsTools/NanoAOD/test/testNanoWeights.py
@@ -19,7 +19,10 @@ def variableAndNumber(varName, tree):
 rtfile = ROOT.TFile(args.inputFile)
 tree = rtfile.Get("Events")
 tree.GetEntry(0)
-variables = ["LHEScaleWeight", "LHEPdfWeight", "LHEMEParamWeight", "GenPartonShowerWeight", "LHEUnknownWeight", ]
+types = ["ScaleWeight", "PdfWeight", "MEParamWeight", "UnknownWeight", ]
+variables = ["LHE"+t for t in types]
+variables.append("GenPartonShowerWeight")
+variables.extend(["Gen"+t for t in types])
 
 for varName in variables:
     variableAndNumber(varName, tree)