Extension of electron matching in nanoAOD

PhysicsTools/HepMCCandAlgos/plugins/MCTruthMatchers.cc

@@ -33,7 +33,16 @@ typedef reco::PhysObjectMatcher<
     reco::MatchByDR<reco::CandidateView::value_type, reco::CandidateView::value_type> >
     GenJetMatcher;
 
+// JET Match by deltaR and dPt, ranking by deltaR
+typedef reco::PhysObjectMatcher<
+    reco::CandidateView,
+    reco::GenJetCollection,
+    reco::MCMatchSelector<reco::CandidateView::value_type, reco::GenJetCollection::value_type>,
+    reco::MatchByDRDPt<reco::CandidateView::value_type, reco::GenJetCollection::value_type> >
+    GenJetMatcherDRPtByDR;
+
 #include "FWCore/Framework/interface/MakerMacros.h"
 DEFINE_FWK_MODULE(MCMatcher);
 DEFINE_FWK_MODULE(MCMatcherByPt);
 DEFINE_FWK_MODULE(GenJetMatcher);
+DEFINE_FWK_MODULE(GenJetMatcherDRPtByDR);

PhysicsTools/NanoAOD/plugins/CandMCMatchTableProducer.cc

@@ -7,6 +7,8 @@
 #include "DataFormats/Common/interface/View.h"
 #include "DataFormats/Candidate/interface/Candidate.h"
 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+#include <DataFormats/Math/interface/deltaR.h>
+#include "DataFormats/JetReco/interface/GenJetCollection.h"
 
 #include <vector>
 #include <iostream>
@@ -63,6 +65,13 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
       candMapVisTau_ =
           consumes<edm::Association<reco::GenParticleCollection>>(params.getParameter<edm::InputTag>("mcMapVisTau"));
     }
+
+    if (type_ == MElectron) {
+      candMapDressedLep_ =
+          consumes<edm::Association<reco::GenJetCollection>>(params.getParameter<edm::InputTag>("mcMapDressedLep"));
+      mapTauAnc_ = consumes<edm::ValueMap<bool>>(params.getParameter<edm::InputTag>("mapTauAnc"));
+      genPartsToken_ = consumes<reco::GenParticleCollection>(params.getParameter<edm::InputTag>("genparticles"));
+    }
   }
 
   ~CandMCMatchTableProducer() override {}
@@ -82,20 +91,38 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
       iEvent.getByToken(candMapVisTau_, mapVisTau);
     }
 
+    edm::Handle<edm::Association<reco::GenJetCollection>> mapDressedLep;
+    edm::Handle<edm::ValueMap<bool>> mapTauAnc;
+    edm::Handle<reco::GenParticleCollection> genParts;
+    if (type_ == MElectron) {
+      iEvent.getByToken(candMapDressedLep_, mapDressedLep);
+      iEvent.getByToken(mapTauAnc_, mapTauAnc);
+      iEvent.getByToken(genPartsToken_, genParts);
+    }
+
     std::vector<int> key(ncand, -1), flav(ncand, 0);
     for (unsigned int i = 0; i < ncand; ++i) {
       //std::cout << "cand #" << i << ": pT = " << cands->ptrAt(i)->pt() << ", eta = " << cands->ptrAt(i)->eta() << ", phi = " << cands->ptrAt(i)->phi() << std::endl;
       reco::GenParticleRef match = (*map)[cands->ptrAt(i)];
       reco::GenParticleRef matchVisTau;
+      reco::GenJetRef matchDressedLep;
+      bool hasTauAnc = false;
       if (type_ == MTau) {
         matchVisTau = (*mapVisTau)[cands->ptrAt(i)];
       }
+      if (type_ == MElectron) {
+        matchDressedLep = (*mapDressedLep)[cands->ptrAt(i)];
+        if (matchDressedLep.isNonnull()) {
+          hasTauAnc = (*mapTauAnc)[matchDressedLep];
+        }
+      }
       if (match.isNonnull())
         key[i] = match.key();
       else if (matchVisTau.isNonnull())
         key[i] = matchVisTau.key();
-      else
-        continue;
+      else if (type_ != MElectron)
+        continue;  // go ahead with electrons, as those may be matched to a dressed lepton
+
       switch (type_) {
         case MMuon:
           if (match->isPromptFinalState())
@@ -106,7 +133,35 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
             flav[i] = getParentHadronFlag(match);  // 3 = light, 4 = charm, 5 = b
           break;
         case MElectron:
-          if (match->isPromptFinalState())
+          if (matchDressedLep.isNonnull()) {
+            if (matchDressedLep->pdgId() == 22)
+              flav[i] = 22;
+            else
+              flav[i] = (hasTauAnc) ? 15 : 1;
+
+            float minpt = 0;
+            const reco::GenParticle* highestPtConstituent = nullptr;
+            for (auto& consti : matchDressedLep->getGenConstituents()) {
+              if (abs(consti->pdgId()) != 11)
+                continue;
+              if (consti->pt() < minpt)
+                continue;
+              minpt = consti->pt();
+              highestPtConstituent = consti;
+            }
+            if (highestPtConstituent) {
+              auto iter =
+                  std::find_if(genParts->begin(), genParts->end(), [highestPtConstituent](reco::GenParticle genp) {
+                    return (abs(genp.pdgId()) == 11) && (deltaR(genp, *highestPtConstituent) < 0.01) &&
+                           (abs(genp.pt() - highestPtConstituent->pt()) / highestPtConstituent->pt() < 0.01);
+                  });
+              if (iter != genParts->end()) {
+                key[i] = iter - genParts->begin();
+              }
+            }
+          } else if (!match.isNonnull())
+            flav[i] = 0;
+          else if (match->isPromptFinalState())
             flav[i] = (match->pdgId() == 22 ? 22 : 1);  // prompt electron or photon
           else if (match->isDirectPromptTauDecayProductFinalState())
             flav[i] = 15;  // tau
@@ -139,7 +194,9 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
     }
 
     tab->addColumn<int>(branchName_ + "Idx", key, "Index into genParticle list for " + doc_);
-    tab->addColumn<uint8_t>(branchName_ + "Flav", flav, "Flavour of genParticle for " + doc_ + ": " + flavDoc_);
+    tab->addColumn<uint8_t>(branchName_ + "Flav",
+                            flav,
+                            "Flavour of genParticle (DressedLeptons for electrons) for " + doc_ + ": " + flavDoc_);
 
     iEvent.put(std::move(tab));
   }
@@ -183,6 +240,11 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
         "type of object to match (Muon, Electron, Tau, Photon, Other), taylors what's in t Flav branch");
     desc.addOptional<edm::InputTag>("mcMapVisTau")
         ->setComment("as mcMap, but pointing to the visible gen taus (only if objType == Tau)");
+    desc.addOptional<edm::InputTag>("mcMapDressedLep")
+        ->setComment("as mcMap, but pointing to gen dressed leptons (only if objType == Electrons)");
+    desc.addOptional<edm::InputTag>("mapTauAnc")
+        ->setComment("Value map of matched gen electrons containing info on the tau ancestry");
+    desc.addOptional<edm::InputTag>("genparticles")->setComment("Collection of genParticles to be stored.");
     descriptions.add("candMcMatchTable", desc);
   }
 
@@ -191,6 +253,9 @@ class CandMCMatchTableProducer : public edm::global::EDProducer<> {
   const edm::EDGetTokenT<reco::CandidateView> src_;
   const edm::EDGetTokenT<edm::Association<reco::GenParticleCollection>> candMap_;
   edm::EDGetTokenT<edm::Association<reco::GenParticleCollection>> candMapVisTau_;
+  edm::EDGetTokenT<edm::Association<reco::GenJetCollection>> candMapDressedLep_;
+  edm::EDGetTokenT<edm::ValueMap<bool>> mapTauAnc_;
+  edm::EDGetTokenT<reco::GenParticleCollection> genPartsToken_;
   enum MatchType { MMuon, MElectron, MTau, MPhoton, MOther } type_;
   std::string flavDoc_;
 };

PhysicsTools/NanoAOD/plugins/GenJetGenPartMerger.cc

@@ -0,0 +1,113 @@
+
+// system include files
+#include <memory>
+
+// user include files
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/StreamID.h"
+
+#include "DataFormats/JetReco/interface/GenJetCollection.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+
+#include "DataFormats/Common/interface/ValueMap.h"
+
+//
+// class declaration
+//
+
+class GenJetGenPartMerger : public edm::stream::EDProducer<> {
+public:
+  explicit GenJetGenPartMerger(const edm::ParameterSet&);
+  ~GenJetGenPartMerger() override;
+
+private:
+  void beginStream(edm::StreamID) override;
+  void produce(edm::Event&, const edm::EventSetup&) override;
+  void endStream() override;
+
+  const edm::EDGetTokenT<reco::GenJetCollection> jetToken_;
+  const edm::EDGetTokenT<reco::GenParticleCollection> partToken_;
+  const edm::EDGetTokenT<edm::ValueMap<bool>> tauAncToken_;
+};
+
+//
+// constants, enums and typedefs
+//
+
+//
+// static data member definitions
+//
+
+//
+// constructors and destructor
+//
+GenJetGenPartMerger::GenJetGenPartMerger(const edm::ParameterSet& iConfig)
+    : jetToken_(consumes<reco::GenJetCollection>(iConfig.getParameter<edm::InputTag>("srcJet"))),
+      partToken_(consumes<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcPart"))),
+      tauAncToken_(consumes<edm::ValueMap<bool>>(iConfig.getParameter<edm::InputTag>("hasTauAnc"))) {
+  produces<reco::GenJetCollection>("merged");
+  produces<edm::ValueMap<bool>>("hasTauAnc");
+}
+
+GenJetGenPartMerger::~GenJetGenPartMerger() {}
+
+//
+// member functions
+//
+
+// ------------ method called to produce the data  ------------
+void GenJetGenPartMerger::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  using namespace edm;
+  std::unique_ptr<reco::GenJetCollection> merged(new reco::GenJetCollection);
+
+  std::vector<bool> hasTauAncValues;
+
+  edm::Handle<reco::GenJetCollection> jetHandle;
+  iEvent.getByToken(jetToken_, jetHandle);
+
+  edm::Handle<reco::GenParticleCollection> partHandle;
+  iEvent.getByToken(partToken_, partHandle);
+
+  edm::Handle<edm::ValueMap<bool>> tauAncHandle;
+  iEvent.getByToken(tauAncToken_, tauAncHandle);
+
+  for (unsigned int ijet = 0; ijet < jetHandle->size(); ++ijet) {
+    auto jet = jetHandle->at(ijet);
+    merged->push_back(reco::GenJet(jet));
+    reco::GenJetRef jetRef(jetHandle, ijet);
+    hasTauAncValues.push_back((*tauAncHandle)[jetRef]);
+  }
+
+  for (auto& part : *partHandle) {
+    reco::GenJet jet;
+    jet.setP4(part.p4());
+    jet.setPdgId(part.pdgId());
+    jet.setCharge(part.charge());
+    merged->push_back(jet);
+    hasTauAncValues.push_back(false);
+  }
+
+  auto newmerged = iEvent.put(std::move(merged), "merged");
+
+  std::unique_ptr<edm::ValueMap<bool>> out(new edm::ValueMap<bool>());
+  edm::ValueMap<bool>::Filler filler(*out);
+  filler.insert(newmerged, hasTauAncValues.begin(), hasTauAncValues.end());
+  filler.fill();
+
+  iEvent.put(std::move(out), "hasTauAnc");
+}
+
+// ------------ method called once each stream before processing any runs, lumis or events  ------------
+void GenJetGenPartMerger::beginStream(edm::StreamID) {}
+
+// ------------ method called once each stream after processing all runs, lumis and events  ------------
+void GenJetGenPartMerger::endStream() {}
+
+//define this as a plug-in
+DEFINE_FWK_MODULE(GenJetGenPartMerger);

PhysicsTools/NanoAOD/python/electrons_cff.py

@@ -503,6 +503,38 @@ def _get_bitmapVIDForEle_docstring(modules,WorkingPoints):
 
 )
 
+
+from PhysicsTools.NanoAOD.particlelevel_cff import particleLevel
+
+particleLevelForMatching = particleLevel.clone(
+    lepMinPt    = cms.double(3.),
+    phoMinPt = cms.double(3),
+)
+
+tautaggerForMatching = cms.EDProducer("GenJetTauTaggerProducer",
+                                      src = cms.InputTag('particleLevelForMatching:leptons')
+)
+
+
+matchingElecPhoton = cms.EDProducer("GenJetGenPartMerger",
+                                    srcJet =cms.InputTag("particleLevelForMatching:leptons"),
+                                    srcPart=cms.InputTag("particleLevelForMatching:photons"),
+                                    hasTauAnc=cms.InputTag("tautaggerForMatching"),
+)
+
+
+electronsMCMatchForTableAlt = cms.EDProducer("GenJetMatcherDRPtByDR",  # cut on deltaR, deltaPt/Pt; pick best by deltaR
+    src         = electronTable.src,                 # final reco collection
+    matched     = cms.InputTag("matchingElecPhoton:merged"), # final mc-truth particle collection
+    mcPdgId     = cms.vint32(11,22),                 # one or more PDG ID (11 = el, 22 = pho); absolute values (see below)
+    checkCharge = cms.bool(False),              # True = require RECO and MC objects to have the same charge
+    mcStatus    = cms.vint32(),
+    maxDeltaR   = cms.double(0.3),              # Minimum deltaR for the match
+    maxDPtRel   = cms.double(0.5),              # Minimum deltaPt/Pt for the match
+    resolveAmbiguities    = cms.bool(True),     # Forbid two RECO objects to match to the same GEN object
+    resolveByMatchQuality = cms.bool(True),    # False = just match input in order; True = pick lowest deltaR pair first
+) 
+
 electronsMCMatchForTable = cms.EDProducer("MCMatcher",  # cut on deltaR, deltaPt/Pt; pick best by deltaR
     src         = electronTable.src,                 # final reco collection
     matched     = cms.InputTag("finalGenParticles"), # final mc-truth particle collection
@@ -517,16 +549,19 @@ def _get_bitmapVIDForEle_docstring(modules,WorkingPoints):
 
 electronMCTable = cms.EDProducer("CandMCMatchTableProducer",
     src     = electronTable.src,
+    mcMapDressedLep = cms.InputTag("electronsMCMatchForTableAlt"),
     mcMap   = cms.InputTag("electronsMCMatchForTable"),
+    mapTauAnc = cms.InputTag("matchingElecPhoton:hasTauAnc"),
     objName = electronTable.name,
     objType = electronTable.name, #cms.string("Electron"),
     branchName = cms.string("genPart"),
     docString = cms.string("MC matching to status==1 electrons or photons"),
+    genparticles     = cms.InputTag("finalGenParticles"), 
 )
 
 electronSequence = cms.Sequence(bitmapVIDForEle + bitmapVIDForEleHEEP + isoForEle + ptRatioRelForEle + seedGainEle + slimmedElectronsWithUserData + finalElectrons)
 electronTables = cms.Sequence (electronMVATTH + electronTable)
-electronMC = cms.Sequence(electronsMCMatchForTable + electronMCTable)
+electronMC = cms.Sequence(particleLevelForMatching + tautaggerForMatching + matchingElecPhoton + electronsMCMatchForTable + electronsMCMatchForTableAlt + electronMCTable)
 from RecoEgamma.ElectronIdentification.heepIdVarValueMapProducer_cfi import heepIDVarValueMaps
 _updateTo106X_sequence =cms.Sequence(heepIDVarValueMaps + slimmedElectronsTo106X)
 heepIDVarValueMaps.dataFormat = 2