Merge pull request #32099 from intrepid42/TnP_112X

Add LHE HT and generator weights to TagAndProbe tree
cms-sw · Nov 19, 2020 · 7942362 · 7942362
2 parents fa3869f + f9ee5eb
commit 7942362
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 1 deletion.
diff --git a/PhysicsTools/TagAndProbe/interface/BaseTreeFiller.h b/PhysicsTools/TagAndProbe/interface/BaseTreeFiller.h
@@ -28,6 +28,9 @@
 
 #include "SimDataFormats/PileupSummaryInfo/interface/PileupSummaryInfo.h"
 #include "SimDataFormats/GeneratorProducts/interface/GenEventInfoProduct.h"
+#include "SimDataFormats/GeneratorProducts/interface/LHEEventProduct.h"
+#include "SimDataFormats/GeneratorProducts/interface/GenLumiInfoHeader.h"
+#include "SimDataFormats/GeneratorProducts/interface/LHERunInfoProduct.h"
 
 #include <TTree.h>
 
@@ -174,7 +177,12 @@ namespace tnp {
     /// How event weights are defined: 'None' = no weights, 'Fixed' = one value specified in cfg file, 'External' = read weight from the event (as double)
     enum WeightMode { None, Fixed, External };
     WeightMode weightMode_;
+    bool LHEinfo_;
     edm::EDGetTokenT<GenEventInfoProduct> weightSrcToken_;
+    edm::EDGetTokenT<LHEEventProduct> _LHECollection;
+    edm::EDGetTokenT<GenLumiInfoHeader> _genLumiInfoToken;
+    edm::EDGetTokenT<LHERunInfoProduct> _lheRunInfoToken;
+    edm::EDGetTokenT<reco::GenParticleCollection> genParticlesToken_;
     edm::EDGetTokenT<double> PUweightSrcToken_;
     edm::EDGetTokenT<double> rhoToken_;
     edm::EDGetTokenT<reco::VertexCollection> recVtxsToken_;
@@ -204,13 +212,16 @@ namespace tnp {
     //implementation notice: these two are 'mutable' because we will fill them from a 'const' method
     mutable TTree *tree_;
     mutable float weight_, PUweight_, totWeight_;
+    mutable float lheWeight_[9];
+    mutable float psWeight_[5];
     mutable uint32_t run_, lumi_, mNPV_;
     mutable uint64_t event_;
     mutable int truePU_;
 
     mutable float mPVx_, mPVy_, mPVz_, mBSx_, mBSy_, mBSz_;
     mutable float rho_;
-    mutable float mMET_, mSumET_, mMETSign_, mtcMET_, mtcSumET_, mtcMETSign_, mpfMET_, mpfSumET_, mpfMETSign_, mpfPhi_;
+    mutable float mMET_, mSumET_, mMETSign_, mtcMET_, mtcSumET_, mtcMETSign_, mpfMET_, mpfSumET_, mpfMETSign_, mpfPhi_,
+        lhe_ht_;
   };
 
 }  // namespace tnp

diff --git a/PhysicsTools/TagAndProbe/src/BaseTreeFiller.cc b/PhysicsTools/TagAndProbe/src/BaseTreeFiller.cc
@@ -45,6 +45,7 @@ tnp::BaseTreeFiller::BaseTreeFiller(const char *name, const edm::ParameterSet &i
   } else if (iConfig.existsAs<edm::InputTag>("eventWeight")) {
     weightMode_ = External;
     weightSrcToken_ = iC.consumes<GenEventInfoProduct>(iConfig.getParameter<edm::InputTag>("eventWeight"));
+    tree_->Branch("psWeight", &psWeight_, "psWeight[5]/F");
   } else {
     weightMode_ = None;
   }
@@ -53,6 +54,13 @@ tnp::BaseTreeFiller::BaseTreeFiller(const char *name, const edm::ParameterSet &i
     tree_->Branch("totWeight", &totWeight_, "totWeight/F");
   }
 
+  LHEinfo_ = iConfig.existsAs<edm::InputTag>("LHEWeightSrc");
+  if (LHEinfo_) {
+    _LHECollection = iC.consumes<LHEEventProduct>(iConfig.getParameter<edm::InputTag>("LHEWeightSrc"));
+    tree_->Branch("lheWeight", &lheWeight_, "lheWeight[9]/F");
+    tree_->Branch("lhe_ht", &lhe_ht_, "lhe_ht/F");
+  }
+
   storePUweight_ = iConfig.existsAs<edm::InputTag>("PUWeightSrc") ? true : false;
   if (storePUweight_) {
     PUweightSrcToken_ = iC.consumes<double>(iConfig.getParameter<edm::InputTag>("PUWeightSrc"));
@@ -194,6 +202,9 @@ void tnp::BaseTreeFiller::init(const edm::Event &iEvent) const {
   for (std::vector<tnp::ProbeFlag>::const_iterator it = flags_.begin(), ed = flags_.end(); it != ed; ++it) {
     it->init(iEvent);
   }
+  for (int i = 0; i < 5; i++) {
+    psWeight_[i] = 1.;  // init
+  }
   if (weightMode_ == External) {
     // edm::Handle<double> weight;
     //        iEvent.getByToken(weightSrcToken_, weight);
@@ -202,6 +213,35 @@ void tnp::BaseTreeFiller::init(const edm::Event &iEvent) const {
     iEvent.getByToken(weightSrcToken_, weight);
     weight_ = weight->weight();
     totWeight_ *= weight_;
+    if (weight->weights().size() >= 10) {
+      int k = 1;
+      for (int i = 6; i < 10; i++) {
+        // hardcoded Pythia 8 isrDefHi,fsrDefHi,isrDefLo,fsrDefLo
+        psWeight_[k] = weight->weights().at(i) / weight->weight();
+        k++;
+      }
+    }
+  }
+
+  for (unsigned int i = 0; i < 9; i++) {
+    lheWeight_[i] = 1.;  // init
+  }
+  lhe_ht_ = 0.;
+  if (LHEinfo_ and !_LHECollection.isUninitialized()) {
+    edm::Handle<LHEEventProduct> lheEventHandle;
+    iEvent.getByToken(_LHECollection, lheEventHandle);
+    for (unsigned int i = 0; i < 9; i++) {
+      lheWeight_[i] = lheEventHandle->weights().at(i).wgt / lheEventHandle->originalXWGTUP();
+    }
+    for (int i = 0; i < lheEventHandle->hepeup().NUP; i++) {
+      int id = lheEventHandle->hepeup().IDUP[i];
+      int st = lheEventHandle->hepeup().ISTUP[i];
+
+      // calculate HT at LHE level
+      if ((abs(id) < 6 || id == 21) && st > 0) {
+        lhe_ht_ += sqrt(pow(lheEventHandle->hepeup().PUP[i][0], 2) + pow(lheEventHandle->hepeup().PUP[i][1], 2));
+      }
+    }
   }
 
   ///// ********** Pileup weight: needed for MC re-weighting for PU *************