From 6ccee5d97cb552ffbad85a773e833fc8e6ef658b Mon Sep 17 00:00:00 2001
From: Kenneth Long <kdlong@hep.wisc.edu>
Date: Fri, 17 Jan 2020 20:57:33 +0100
Subject: [PATCH] Working implementation producing new prod if needed

---
 .../Core/plugins/LHEWeightProductProducer.cc  | 24 ++++++++++-----
 .../plugins/LHEWeightsTableProducer.cc        | 29 ++++++++++++++-----
 PhysicsTools/NanoAOD/python/nanogen_cff.py    |  7 ++++-
 3 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/GeneratorInterface/Core/plugins/LHEWeightProductProducer.cc b/GeneratorInterface/Core/plugins/LHEWeightProductProducer.cc
index a4f0836314460..db74a3b25730b 100644
--- a/GeneratorInterface/Core/plugins/LHEWeightProductProducer.cc
+++ b/GeneratorInterface/Core/plugins/LHEWeightProductProducer.cc
@@ -36,6 +36,8 @@ class LHEWeightProductProducer : public edm::one::EDProducer<edm::BeginLuminosit
   std::string lheLabel_;
   edm::EDGetTokenT<LHERunInfoProduct> lheRunInfoToken_;
   edm::EDGetTokenT<LHEEventProduct> lheEventToken_;
+  const edm::EDGetTokenT<GenWeightInfoProduct> lheWeightInfoToken_;
+  bool foundWeightProduct_;
 
   void produce(edm::Event&, const edm::EventSetup&) override;
   void beginLuminosityBlockProduce(edm::LuminosityBlock& lumi, edm::EventSetup const& es) override;
@@ -48,11 +50,10 @@ class LHEWeightProductProducer : public edm::one::EDProducer<edm::BeginLuminosit
 // constructors and destructor
 //
 LHEWeightProductProducer::LHEWeightProductProducer(const edm::ParameterSet& iConfig) :
-    lheLabel_(iConfig.getUntrackedParameter<std::string>("lheSourceLabel")),
-    lheRunInfoToken_(consumes<LHERunInfoProduct, edm::InRun>(
-        iConfig.getUntrackedParameter<edm::InputTag>("lheSource", edm::InputTag("externalLHEProducer")))),
-    lheEventToken_(consumes<LHEEventProduct>(
-        iConfig.getUntrackedParameter<edm::InputTag>("lheSource", edm::InputTag("externalLHEProducer"))))
+    lheLabel_(iConfig.getParameter<std::string>("lheSourceLabel")),
+    lheRunInfoToken_(consumes<LHERunInfoProduct, edm::InRun>(lheLabel_)),
+    lheEventToken_(consumes<LHEEventProduct>(lheLabel_)),
+    lheWeightInfoToken_(consumes<GenWeightInfoProduct, edm::InLumi>(lheLabel_))
 {
   produces<GenWeightProduct>();
   produces<GenWeightInfoProduct, edm::Transition::BeginLuminosityBlock>();
@@ -67,6 +68,9 @@ LHEWeightProductProducer::~LHEWeightProductProducer()
 // ------------ method called to produce the data  ------------
 void
 LHEWeightProductProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  if (foundWeightProduct_)
+      return;
+  
   edm::Handle<LHEEventProduct> lheEventInfo;
   iEvent.getByToken(lheEventToken_, lheEventInfo);
   // Read weights from LHEEventProduct
@@ -74,12 +78,9 @@ LHEWeightProductProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
   iEvent.put(std::move(weightProduct));
 }
 
-// ------------ method called when starting to processes a run  ------------
 void 
 LHEWeightProductProducer::beginRun(edm::Run const& run, edm::EventSetup const& es) {
     edm::Handle<LHERunInfoProduct> lheRunInfoHandle;
-    //run.getByToken(lheRunInfoToken_, lheRunInfoHandle);
-    // get by token gives an error (the same one that's been in the ExternalLHEProducer for ages)
     run.getByLabel(lheLabel_, lheRunInfoHandle);
 
     typedef std::vector<LHERunInfoProduct::Header>::const_iterator header_cit;
@@ -97,6 +98,13 @@ LHEWeightProductProducer::endRun(edm::Run const& run, edm::EventSetup const& es)
 
 void
 LHEWeightProductProducer::beginLuminosityBlockProduce(edm::LuminosityBlock& lumi, edm::EventSetup const& es) {
+    edm::Handle<GenWeightInfoProduct> lheWeightInfoHandle;
+    lumi.getByToken(lheWeightInfoToken_, lheWeightInfoHandle);
+    if (lheWeightInfoHandle.isValid()) {
+        foundWeightProduct_ = true;
+        return;
+    }
+
     auto weightInfoProduct = std::make_unique<GenWeightInfoProduct>();
     for (auto& weightGroup : weightHelper_.weightGroups()) {
         weightInfoProduct->addWeightGroupInfo(weightGroup.clone());
diff --git a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
index ed26af86897bc..5587cbd3d31e6 100644
--- a/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
+++ b/PhysicsTools/NanoAOD/plugins/LHEWeightsTableProducer.cc
@@ -28,8 +28,10 @@ namespace {
 class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBlockCache<WeightGroupsToStore>> {
 public:
   LHEWeightsTableProducer(edm::ParameterSet const& params) :
-        lheWeightToken_(consumes<GenWeightProduct>(params.getParameter<edm::InputTag>("lheWeights"))),
-        lheWeightInfoToken_(consumes<GenWeightInfoProduct, edm::InLumi>(params.getParameter<edm::InputTag>("lheWeights"))),
+        lheWeightTokens_(edm::vector_transform(params.getParameter<std::vector<edm::InputTag>>("lheWeights"),
+                                [this](const edm::InputTag& tag) { return mayConsume<GenWeightProduct>(tag); })),
+        lheWeightInfoTokens_(edm::vector_transform(params.getParameter<std::vector<edm::InputTag>>("lheWeights"),
+                                [this](const edm::InputTag& tag) { return mayConsume<GenWeightInfoProduct, edm::InLumi>(tag); })),
         genWeightToken_(consumes<GenWeightProduct>(params.getParameter<edm::InputTag>("genWeights"))),
         genWeightInfoToken_(consumes<GenWeightInfoProduct, edm::InLumi>(params.getParameter<edm::InputTag>("genWeights"))),
         weightgroups_(edm::vector_transform(params.getParameter<std::vector<std::string>>("weightgroups"),
@@ -43,9 +45,14 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
   }
 
   void produce(edm::StreamID id, edm::Event& iEvent, const edm::EventSetup& iSetup) const override {
-    // tables for LHE weights, may not be filled
     edm::Handle<GenWeightProduct> lheWeightHandle;
-    iEvent.getByToken(lheWeightToken_, lheWeightHandle);
+    for (auto& token : lheWeightTokens_) {
+        iEvent.getByToken(token, lheWeightHandle);
+        if (lheWeightHandle.isValid()) {
+            break;
+        }
+    }
+
     const GenWeightProduct* lheWeightProduct = lheWeightHandle.product();
     WeightsContainer lheWeights = lheWeightProduct->weights();
 
@@ -113,7 +120,12 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
     // Set equal to the max number of groups
     // subtrack 1 for each weight group you find
     edm::Handle<GenWeightInfoProduct> lheWeightInfoHandle;
-    iLumi.getByToken(lheWeightInfoToken_, lheWeightInfoHandle);
+    for (auto& token : lheWeightInfoTokens_) {
+        iLumi.getByToken(token, lheWeightInfoHandle);
+        if (lheWeightInfoHandle.isValid()) {
+            break;
+        }
+    }
 
     edm::Handle<GenWeightInfoProduct> genWeightInfoHandle;
     iLumi.getByToken(genWeightInfoToken_, genWeightInfoHandle);
@@ -173,7 +185,8 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
 
   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
     edm::ParameterSetDescription desc;
-    desc.add<edm::InputTag>("lheWeights");
+    desc.add<std::vector<edm::InputTag>>("lheWeights");
+    //desc.add<std::vector<edm::InputTag>>("genWeights");
     desc.add<edm::InputTag>("genWeights");
     desc.add<std::vector<std::string>>("weightgroups");
     desc.add<std::vector<int>>("maxGroupsPerType");
@@ -184,8 +197,8 @@ class LHEWeightsTableProducer : public edm::global::EDProducer<edm::LuminosityBl
 
 protected:
   const edm::EDGetTokenT<LHEEventProduct> lheToken_;
-  const edm::EDGetTokenT<GenWeightProduct> lheWeightToken_;
-  const edm::EDGetTokenT<GenWeightInfoProduct> lheWeightInfoToken_;
+  const std::vector<edm::EDGetTokenT<GenWeightProduct>> lheWeightTokens_;
+  const std::vector<edm::EDGetTokenT<GenWeightInfoProduct>> lheWeightInfoTokens_;
   const edm::EDGetTokenT<GenWeightProduct> genWeightToken_;
   const edm::EDGetTokenT<GenWeightInfoProduct> genWeightInfoToken_;
   const std::vector<gen::WeightType> weightgroups_;
diff --git a/PhysicsTools/NanoAOD/python/nanogen_cff.py b/PhysicsTools/NanoAOD/python/nanogen_cff.py
index b79d0d61c0b39..88d8ebe03c18f 100644
--- a/PhysicsTools/NanoAOD/python/nanogen_cff.py
+++ b/PhysicsTools/NanoAOD/python/nanogen_cff.py
@@ -9,9 +9,13 @@
 genWeights = cms.EDProducer("GenWeightProductProducer",
     genInfo = cms.InputTag("generator"))
 
+lheWeights = cms.EDProducer("LHEWeightProductProducer",
+    lheSourceLabel = cms.string("externalLHEProducer"),
+    lheSource = cms.InputTag("externalLHEProducer"))
+
 lheWeightsTable = cms.EDProducer(
     "LHEWeightsTableProducer",
-    lheWeights = cms.InputTag("externalLHEProducer"),
+    lheWeights = cms.VInputTag(["externalLHEProducer", "lheWeights"]),
     genWeights = cms.InputTag("genWeights"),
     # Warning: you can use a full string, but only the first character is read.
     # Note also that the capitalization is important! For example, 'parton shower' 
@@ -44,6 +48,7 @@
 
 nanogenSequence = cms.Sequence(
     genWeights+
+    lheWeights+
     nanoMetadata+
     particleLevel+
     genJetTable+