[JME,Nano] Refactor constituents table, add GloParT WvsQCD score

PhysicsTools/NanoAOD/plugins/SimpleJetConstituentTableProducer.cc

@@ -16,12 +16,16 @@
 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
 #include "DataFormats/NanoAOD/interface/FlatTable.h"
 
-template <typename T>
+template <typename T, typename C = std::vector<typename T::ConstituentTypeFwdPtr>>
 class SimpleJetConstituentTableProducer : public edm::stream::EDProducer<> {
 public:
+  using ConstituentsOutput = C;
+  using ConstituentValueType = typename C::value_type;
+
   explicit SimpleJetConstituentTableProducer(const edm::ParameterSet &);
   ~SimpleJetConstituentTableProducer() override;
 
+  ConstituentValueType const initptr(edm::Ptr<reco::Candidate> const &) const;
   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
 
 private:
@@ -35,41 +39,40 @@ class SimpleJetConstituentTableProducer : public edm::stream::EDProducer<> {
   edm::EDGetTokenT<reco::CandidateView> cand_token_;
 
   const StringCutObjectSelector<T> jetCut_;
-
-  edm::Handle<reco::CandidateView> cands_;
+  const StringCutObjectSelector<ConstituentValueType> jetConstCut_;
 };
 
 //
 // constructors and destructor
 //
-template <typename T>
-SimpleJetConstituentTableProducer<T>::SimpleJetConstituentTableProducer(const edm::ParameterSet &iConfig)
+template <typename T, typename C>
+SimpleJetConstituentTableProducer<T, C>::SimpleJetConstituentTableProducer(const edm::ParameterSet &iConfig)
     : name_(iConfig.getParameter<std::string>("name")),
       candIdxName_(iConfig.getParameter<std::string>("candIdxName")),
       candIdxDoc_(iConfig.getParameter<std::string>("candIdxDoc")),
       jet_token_(consumes<edm::View<T>>(iConfig.getParameter<edm::InputTag>("jets"))),
       cand_token_(consumes<reco::CandidateView>(iConfig.getParameter<edm::InputTag>("candidates"))),
-      jetCut_(iConfig.getParameter<std::string>("jetCut")) {
+      jetCut_(iConfig.getParameter<std::string>("jetCut")),
+      jetConstCut_(iConfig.getParameter<std::string>("jetConstCut")) {
   produces<nanoaod::FlatTable>(name_);
-  produces<std::vector<reco::CandidatePtr>>();
+  produces<ConstituentsOutput>();
 }
 
-template <typename T>
-SimpleJetConstituentTableProducer<T>::~SimpleJetConstituentTableProducer() {}
+template <typename T, typename C>
+SimpleJetConstituentTableProducer<T, C>::~SimpleJetConstituentTableProducer() {}
 
-template <typename T>
-void SimpleJetConstituentTableProducer<T>::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
+template <typename T, typename C>
+void SimpleJetConstituentTableProducer<T, C>::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
   // elements in all these collections must have the same order!
-  auto outCands = std::make_unique<std::vector<reco::CandidatePtr>>();
+  auto outCands = std::make_unique<ConstituentsOutput>();
 
   auto jets = iEvent.getHandle(jet_token_);
 
+  edm::Handle<reco::CandidateView> cands_;
   iEvent.getByToken(cand_token_, cands_);
   auto candPtrs = cands_->ptrs();
 
-  //
   // First, select jets
-  //
   std::vector<T> jetsPassCut;
   for (unsigned jetIdx = 0; jetIdx < jets->size(); ++jetIdx) {
     const auto &jet = jets->at(jetIdx);
@@ -78,23 +81,26 @@ void SimpleJetConstituentTableProducer<T>::produce(edm::Event &iEvent, const edm
     jetsPassCut.push_back(jets->at(jetIdx));
   }
 
-  //
   // Then loop over selected jets
-  //
   std::vector<int> parentJetIdx;
   std::vector<int> candIdx;
   for (unsigned jetIdx = 0; jetIdx < jetsPassCut.size(); ++jetIdx) {
     const auto &jet = jetsPassCut.at(jetIdx);
 
-    //
     // Loop over jet constituents
-    //
     std::vector<reco::CandidatePtr> const &daughters = jet.daughterPtrVector();
-    for (const auto &cand : daughters) {
-      auto candInNewList = std::find(candPtrs.begin(), candPtrs.end(), cand);
+    for (const auto &dauPtr : daughters) {
+      // Apply cut on jet constituent
+      typename C::value_type cand = initptr(dauPtr);
+      if (!jetConstCut_(cand))
+        continue;
+
+      // Find jet constituent in candidate collection
+      auto candInNewList = std::find(candPtrs.begin(), candPtrs.end(), dauPtr);
       if (candInNewList == candPtrs.end()) {
         continue;
       }
+
       outCands->push_back(cand);
       parentJetIdx.push_back(jetIdx);
       candIdx.push_back(candInNewList - candPtrs.begin());
@@ -103,34 +109,53 @@ void SimpleJetConstituentTableProducer<T>::produce(edm::Event &iEvent, const edm
 
   auto candTable = std::make_unique<nanoaod::FlatTable>(outCands->size(), name_, false);
   // We fill from here only stuff that cannot be created with the SimpleFlatTableProducer
-  candTable->addColumn<int>(candIdxName_, candIdx, candIdxDoc_);
+  candTable->template addColumn<int>(candIdxName_, candIdx, candIdxDoc_);
 
   std::string parentJetIdxName("jetIdx");
   std::string parentJetIdxDoc("Index of the parent jet");
   if constexpr (std::is_same<T, reco::GenJet>::value) {
     parentJetIdxName = "genJetIdx";
     parentJetIdxDoc = "Index of the parent gen jet";
   }
-  candTable->addColumn<int>(parentJetIdxName, parentJetIdx, parentJetIdxDoc);
+  candTable->template addColumn<int>(parentJetIdxName, parentJetIdx, parentJetIdxDoc);
 
   iEvent.put(std::move(candTable), name_);
   iEvent.put(std::move(outCands));
 }
 
-template <typename T>
-void SimpleJetConstituentTableProducer<T>::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
+template <>
+edm::Ptr<pat::PackedCandidate> const
+SimpleJetConstituentTableProducer<pat::Jet, std::vector<edm::Ptr<pat::PackedCandidate>>>::initptr(
+    edm::Ptr<reco::Candidate> const &dau) const {
+  edm::Ptr<pat::PackedCandidate> retval(dau);
+  return retval;
+}
+
+template <>
+edm::Ptr<pat::PackedGenParticle> const
+SimpleJetConstituentTableProducer<reco::GenJet, std::vector<edm::Ptr<pat::PackedGenParticle>>>::initptr(
+    edm::Ptr<reco::Candidate> const &dau) const {
+  edm::Ptr<pat::PackedGenParticle> retval(dau);
+  return retval;
+}
+
+template <typename T, typename C>
+void SimpleJetConstituentTableProducer<T, C>::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
   edm::ParameterSetDescription desc;
   desc.add<std::string>("name", "FatJetPFCand");
   desc.add<std::string>("candIdxName", "PFCandIdx");
   desc.add<std::string>("candIdxDoc", "Index in PFCand table");
   desc.add<edm::InputTag>("jets", edm::InputTag("finalJetsAK8"));
   desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"));
   desc.add<std::string>("jetCut", "");
+  desc.add<std::string>("jetConstCut", "");
   descriptions.addWithDefaultLabel(desc);
 }
 
-typedef SimpleJetConstituentTableProducer<pat::Jet> SimplePatJetConstituentTableProducer;
-typedef SimpleJetConstituentTableProducer<reco::GenJet> SimpleGenJetConstituentTableProducer;
+typedef SimpleJetConstituentTableProducer<pat::Jet, std::vector<edm::Ptr<pat::PackedCandidate>>>
+    SimplePatJetConstituentTableProducer;
+typedef SimpleJetConstituentTableProducer<reco::GenJet, std::vector<edm::Ptr<pat::PackedGenParticle>>>
+    SimpleGenJetConstituentTableProducer;
 
 DEFINE_FWK_MODULE(SimplePatJetConstituentTableProducer);
 DEFINE_FWK_MODULE(SimpleGenJetConstituentTableProducer);

PhysicsTools/NanoAOD/plugins/SimplePATFlatTableProducerPlugins.cc

@@ -24,6 +24,9 @@ typedef SimpleFlatTableProducer<pat::GenericParticle> SimplePATGenericParticleFl
 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
 typedef SimpleFlatTableProducer<pat::PackedCandidate> SimplePATCandidateFlatTableProducer;
 
+#include "DataFormats/PatCandidates/interface/PackedGenParticle.h"
+typedef SimpleFlatTableProducer<pat::PackedGenParticle> SimplePATGenParticleFlatTableProducer;
+
 #include "DataFormats/PatCandidates/interface/MET.h"
 typedef SimpleFlatTableProducer<pat::MET> SimplePATMETFlatTableProducer;
 
@@ -44,6 +47,7 @@ DEFINE_FWK_MODULE(SimplePATJetFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATIsolatedTrackFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATGenericParticleFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATCandidateFlatTableProducer);
+DEFINE_FWK_MODULE(SimplePATGenParticleFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATMETFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATElectron2TrackTimeLifeInfoFlatTableProducer);
 DEFINE_FWK_MODULE(SimplePATMuon2TrackTimeLifeInfoFlatTableProducer);

PhysicsTools/NanoAOD/python/jetConstituents_cff.py

@@ -0,0 +1,194 @@
+import FWCore.ParameterSet.Config as cms
+
+from PhysicsTools.NanoAOD.common_cff import *
+from PhysicsTools.NanoAOD.jetsAK8_cff import fatJetTable as _fatJetTable
+
+##############################################################
+# Take AK8 jets and collect their PF constituents
+###############################################################
+finalJetsAK8PFConstituents = cms.EDProducer("PatJetConstituentPtrSelector",
+    src = _fatJetTable.src,
+    cut = cms.string("abs(eta) <= 2.5")
+)
+
+selectedFinalJetsAK8PFConstituents = cms.EDFilter("PATPackedCandidatePtrSelector",
+    src = cms.InputTag("finalJetsAK8PFConstituents", "constituents"),
+    cut = cms.string("")
+)
+
+##############################################################
+# Setup PF candidates table
+##############################################################
+finalPFCandidates = cms.EDProducer("PackedCandidatePtrMerger",
+    src = cms.VInputTag(cms.InputTag("selectedFinalJetsAK8PFConstituents")),
+    skipNulls = cms.bool(True),
+    warnOnSkip = cms.bool(True)
+)
+
+pfCandidatesTable = cms.EDProducer("SimplePATCandidateFlatTableProducer",
+    src = cms.InputTag("finalPFCandidates"),
+    cut = cms.string(""),
+    name = cms.string("PFCand"),
+    doc = cms.string("PF candidate constituents of AK8 puppi jets (FatJet) with |eta| <= 2.5"),
+    singleton = cms.bool(False),
+    extension = cms.bool(False),
+    variables = cms.PSet(
+        pt = Var("pt * puppiWeight()", float, doc="Puppi-weighted pt", precision=10),
+        mass = Var("mass * puppiWeight()", float, doc="Puppi-weighted mass", precision=10),
+        eta = Var("eta", float, precision=12),
+        phi = Var("phi", float, precision=12),
+        pdgId = Var("pdgId", int, doc="PF candidate type (+/-211 = ChgHad, 130 = NeuHad, 22 = Photon, +/-11 = Electron, +/-13 = Muon, 1 = HFHad, 2 = HFEM)")
+  )
+)
+
+##############################################################
+# Setup AK8 jet constituents table
+##############################################################
+finalJetsAK8ConstituentsTable = cms.EDProducer("SimplePatJetConstituentTableProducer",
+  name = cms.string(_fatJetTable.name.value()+"PFCand"),
+  candIdxName = cms.string("pfCandIdx"),
+  candIdxDoc = cms.string("Index in the PFCand table"),
+  candidates = pfCandidatesTable.src,
+  jets = _fatJetTable.src,
+  jetCut = _fatJetTable.cut,
+  jetConstCut = selectedFinalJetsAK8PFConstituents.cut
+)
+
+jetConstituentsTask = cms.Task(finalJetsAK8PFConstituents,selectedFinalJetsAK8PFConstituents)
+jetConstituentsTablesTask = cms.Task(finalPFCandidates,pfCandidatesTable,finalJetsAK8ConstituentsTable)
+
+
+def SaveAK4JetConstituents(process, jetCut="", jetConstCut=""):
+    """
+    This function can be used as a cmsDriver customization
+    function to add AK4 jet constituents, on top of the AK8
+    jet constituents.
+    """
+    process.finalJetsPuppiPFConstituents = process.finalJetsAK8PFConstituents.clone(
+        src = process.jetPuppiTable.src,
+        cut = jetCut
+    )
+    process.jetConstituentsTask.add(process.finalJetsPuppiPFConstituents)
+
+    process.selectedFinalJetsPuppiPFConstituents = process.selectedFinalJetsAK8PFConstituents.clone(
+        src = cms.InputTag("finalJetsPuppiPFConstituents", "constituents"),
+        cut = jetConstCut
+    )
+    process.jetConstituentsTask.add(process.selectedFinalJetsPuppiPFConstituents)
+
+    process.finalPFCandidates.src += ["selectedFinalJetsPuppiPFConstituents"]
+    process.pfCandidatesTable.doc = pfCandidatesTable.doc.value()+" and AK4 puppi jets (Jet)"
+
+    process.finalJetsPuppiConstituentsTable = process.finalJetsAK8ConstituentsTable.clone(
+        name = process.jetPuppiTable.name.value()+"PFCand",
+        jets = process.jetPuppiTable.src,
+        jetCut = process.jetPuppiTable.cut,
+        jetConstCut = process.selectedFinalJetsPuppiPFConstituents.cut
+    )
+    process.jetConstituentsTablesTask.add(process.finalJetsPuppiConstituentsTable)
+
+    return process
+
+def SaveGenJetConstituents(process, addGenJetConst, addGenJetAK8Const, genJetConstCut="",genJetAK8ConstCut=""):
+    """
+    This function can be used as a cmsDriver
+    customization function to add gen jet
+    constituents.
+    """
+    process.genjetConstituentsTask = cms.Task()
+    process.genjetConstituentsTableTask = cms.Task()
+
+    if addGenJetConst:
+        process.genJetConstituents = cms.EDProducer("GenJetPackedConstituentPtrSelector",
+            src = process.genJetTable.src,
+            cut = process.genJetTable.cut,
+        )
+        process.genjetConstituentsTask.add(process.genJetConstituents)
+
+        process.selectedGenJetConstituents = cms.EDFilter("PATPackedGenParticlePtrSelector",
+            src = cms.InputTag("genJetConstituents", "constituents"),
+            cut = cms.string(genJetConstCut)
+        )
+        process.genjetConstituentsTask.add(process.selectedGenJetConstituents)
+
+    if addGenJetAK8Const:
+        process.genJetAK8Constituents = cms.EDProducer("GenJetPackedConstituentPtrSelector",
+            src = process.genJetAK8Table.src,
+            cut = process.genJetAK8Table.cut,
+        )
+        process.genjetConstituentsTask.add(process.genJetAK8Constituents)
+
+        process.selectedGenJetAK8Constituents = cms.EDFilter("PATPackedGenParticlePtrSelector",
+            src = cms.InputTag("genJetAK8Constituents", "constituents"),
+            cut = cms.string(genJetConstCut)
+        )
+        process.genjetConstituentsTask.add(process.selectedGenJetAK8Constituents)
+
+    if addGenJetConst or addGenJetConst:
+        process.finalGenPartCandidates = cms.EDProducer("PackedGenParticlePtrMerger",
+            src = cms.VInputTag(),
+            skipNulls = cms.bool(True),
+            warnOnSkip = cms.bool(True)
+        )
+        process.genjetConstituentsTableTask.add(process.finalGenPartCandidates)
+
+        process.genPartCandidatesTable = cms.EDProducer("SimplePATGenParticleFlatTableProducer",
+            src = cms.InputTag("finalGenPartCandidates"),
+            cut = cms.string(""),
+            name = cms.string("GenPartCand"),
+            doc = cms.string("Gen particle constituents:"),
+            singleton = cms.bool(False),
+            extension = cms.bool(False),
+            variables = cms.PSet(P4Vars,
+                pdgId = Var("pdgId", int, doc="pdgId")
+          )
+        )
+        process.genjetConstituentsTableTask.add(process.genPartCandidatesTable)
+        process.genPartCandidatesTable.variables.pt.precision=10
+        process.genPartCandidatesTable.variables.mass.precision=10
+
+        if addGenJetConst:
+            process.finalGenPartCandidates.src += ["selectedGenJetConstituents"]
+            process.genPartCandidatesTable.doc = process.genPartCandidatesTable.doc.value()+" AK4 Gen jets (GenJet) "
+
+            process.genJetConstituentsTable = cms.EDProducer("SimpleGenJetConstituentTableProducer",
+              name = cms.string(process.genJetTable.name.value()+"GenPartCand"),
+              candIdxName = cms.string("genPartCandIdx"),
+              candIdxDoc = cms.string("Index in the GenPartCand table"),
+              candidates = pfCandidatesTable.src,
+              jets = process.genJetTable.src,
+              jetCut = process.genJetTable.cut,
+              jetConstCut = process.selectedGenJetConstituents.cut
+            )
+            process.genjetConstituentsTableTask.add(process.genJetConstituentsTable)
+
+        if addGenJetAK8Const:
+            process.finalGenPartCandidates.src += ["selectedGenJetAK8Constituents"]
+            process.genPartCandidatesTable.doc = process.genPartCandidatesTable.doc.value()+" AK8 Gen jets (GenJetAK8)"
+
+            process.genJetAK8ConstituentsTable = cms.EDProducer("SimpleGenJetConstituentTableProducer",
+              name = cms.string(process.genJetAK8Table.name.value()+"GenPartCand"),
+              candIdxName = cms.string("genPartCandIdx"),
+              candIdxDoc = cms.string("Index in the GenPartCand table"),
+              candidates = pfCandidatesTable.src,
+              jets = process.genJetAK8Table.src,
+              jetCut = process.genJetAK8Table.cut,
+              jetConstCut = process.selectedGenJetConstituents.cut
+            )
+            process.genjetConstituentsTableTask.add(process.genJetAK8ConstituentsTable)
+
+        process.nanoTableTaskFS.add(process.genjetConstituentsTask)
+        process.nanoTableTaskFS.add(process.genjetConstituentsTableTask)
+
+    return process
+
+def SaveGenJetAK4Constituents(process):
+    process = SaveGenJetConstituents(process,True,False)
+    return process
+def SaveGenJetAK8Constituents(process):
+    process = SaveGenJetConstituents(process,True,False)
+    return process
+def SaveGenJetAK4AK8Constituents(process):
+    process = SaveGenJetConstituents(process,True,True)
+    return process
+