[PPS][NanoAOD] Generator-level proton information

PhysicsTools/NanoAOD/plugins/GenProtonTableProducer.cc

@@ -0,0 +1,117 @@
+/****************************************************************************
+ *
+ * This is a part of PPS offline software.
+ * Authors:
+ *   Laurent Forthomme
+ *   Michael Pitt
+ *
+ ****************************************************************************/
+
+#include <memory>
+
+#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 "CommonTools/Utils/interface/StringCutObjectSelector.h"
+
+#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
+#include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
+
+#include "DataFormats/NanoAOD/interface/FlatTable.h"
+
+class GenProtonTableProducer : public edm::stream::EDProducer<> {
+public:
+  explicit GenProtonTableProducer(const edm::ParameterSet&);
+  ~GenProtonTableProducer() override = default;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions&);
+
+private:
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+  const edm::EDGetTokenT<reco::GenParticleCollection> prunedCandsToken_;
+  const edm::EDGetTokenT<reco::GenParticleCollection> puCandsToken_, puAltCandsToken_;
+  const StringCutObjectSelector<reco::Candidate> protonsCut_;
+  const std::string table_name_;
+  const double tolerance_;
+  bool use_alt_coll_{false};  ///< Are we using premix/mix collection name for PU protons?
+};
+
+GenProtonTableProducer::GenProtonTableProducer(const edm::ParameterSet& iConfig)
+    : prunedCandsToken_(consumes<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcPruned"))),
+      puCandsToken_(mayConsume<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcPUProtons"))),
+      puAltCandsToken_(mayConsume<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcAltPUProtons"))),
+      protonsCut_(iConfig.getParameter<std::string>("cut")),
+      table_name_(iConfig.getParameter<std::string>("name")),
+      tolerance_(iConfig.getParameter<double>("tolerance")) {
+  produces<nanoaod::FlatTable>();
+}
+
+void GenProtonTableProducer::produce(edm::Event& iEvent, const edm::EventSetup&) {
+  // define the variables
+  std::vector<float> pxs, pys, pzs, vzs;
+  std::vector<bool> isPUs;
+  // first loop over signal protons
+  for (const auto& pruned_cand : iEvent.get(prunedCandsToken_)) {
+    if (!protonsCut_(pruned_cand))
+      continue;
+    pxs.emplace_back(pruned_cand.px());
+    pys.emplace_back(pruned_cand.py());
+    pzs.emplace_back(pruned_cand.pz());
+    vzs.emplace_back(pruned_cand.vz());
+    isPUs.emplace_back(false);
+  }
+  // then loop over pruned candidates ; if already in signal protons, discard
+  edm::Handle<reco::GenParticleCollection> hPUCands;
+  if (use_alt_coll_ || !iEvent.getByToken(puCandsToken_, hPUCands))
+    use_alt_coll_ = iEvent.getByToken(puAltCandsToken_, hPUCands);
+  for (const auto& pu_cand : *hPUCands) {
+    if (!protonsCut_(pu_cand))
+      continue;
+    bool associated{false};
+    for (size_t i = 0; i < pzs.size(); ++i) {
+      if (fabs(1. - pxs.at(i) / pu_cand.px()) < tolerance_ && fabs(1. - pys.at(i) / pu_cand.py()) < tolerance_ &&
+          fabs(1. - pzs.at(i) / pu_cand.pz()) < tolerance_) {
+        associated = true;
+        break;
+      }
+    }
+    if (associated)
+      continue;
+    pxs.emplace_back(pu_cand.px());
+    pys.emplace_back(pu_cand.py());
+    pzs.emplace_back(pu_cand.pz());
+    vzs.emplace_back(pu_cand.vz());
+    isPUs.emplace_back(true);
+  }
+
+  auto protons_table = std::make_unique<nanoaod::FlatTable>(isPUs.size(), table_name_, false);
+  protons_table->addColumn<float>("px", pxs, "proton horizontal momentum", 8);
+  protons_table->addColumn<float>("py", pys, "proton vertical momentum", 8);
+  protons_table->addColumn<float>("pz", pzs, "proton longitudinal momentum", 8);
+  protons_table->addColumn<float>("vz", vzs, "proton vertex longitudinal coordinate", 8);
+  protons_table->addColumn<bool>("isPU", isPUs, "pileup proton?");
+  iEvent.put(std::move(protons_table));
+}
+
+void GenProtonTableProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("srcPruned", edm::InputTag("prunedGenParticles"))
+      ->setComment("input source for pruned gen-level particle candidates");
+  desc.add<edm::InputTag>("srcPUProtons", edm::InputTag("genPUProtons"))
+      ->setComment("input source for pileup protons collection");
+  desc.add<edm::InputTag>("srcAltPUProtons", edm::InputTag("genPUProtons", "genPUProtons"))
+      ->setComment("alternative input source for pileup protons collection (for premix-mix backward compatibility)");
+  desc.add<std::string>("cut", "")->setComment("proton kinematic selection");
+  desc.add<std::string>("name", "GenProton")->setComment("flat table name");
+  desc.add<std::string>("doc", "generator level information on (signal+PU) protons")
+      ->setComment("flat table description");
+  desc.add<double>("tolerance", 1.e-3)->setComment("relative difference between the signal and pileup protons momenta");
+  descriptions.add("genProtonTable", desc);
+}
+
+DEFINE_FWK_MODULE(GenProtonTableProducer);

PhysicsTools/NanoAOD/python/nano_cff.py

@@ -82,7 +82,7 @@
 nanoTableTaskFS = cms.Task(genParticleTask, particleLevelTask, jetMCTask, muonMCTask, electronMCTask, lowPtElectronMCTask, photonMCTask,
                             tauMCTask, boostedTauMCTask,
                             metMCTable, ttbarCatMCProducersTask, globalTablesMCTask, cms.Task(btagWeightTable), ttbarCategoryTableTask,
-                            genWeightsTableTask, genVertexTablesTask, genParticleTablesTask, particleLevelTablesTask)
+                            genWeightsTableTask, genVertexTablesTask, genParticleTablesTask, genProtonTablesTask, particleLevelTablesTask)
 
 nanoSequenceFS = cms.Sequence(nanoSequenceCommon + cms.Sequence(nanoTableTaskFS))
 
@@ -105,7 +105,7 @@ def nanoAOD_addTauIds(process):
 
 def nanoAOD_addBoostedTauIds(process):
     updatedBoostedTauName = "slimmedTausBoostedNewID"
-    boostedTauIdEmbedder = tauIdConfig.TauIDEmbedder(process, debug=False, 
+    boostedTauIdEmbedder = tauIdConfig.TauIDEmbedder(process, debug=False,
                                                      originalTauName = "slimmedTausBoosted",
                                                      updatedTauName = updatedBoostedTauName,
                                                      postfix="Boosted",
@@ -118,7 +118,7 @@ def nanoAOD_addBoostedTauIds(process):
     process.boostedTauTask = _boostedTauTask.copy()
 
     return process
- 
+
 
 from PhysicsTools.PatAlgos.tools.jetTools import updateJetCollection
 def nanoAOD_addDeepInfo(process,addDeepBTag,addDeepFlavour):
@@ -202,7 +202,7 @@ def nanoAOD_recalibrateMETs(process,isData):
         process.patJetsPuppi.addGenJetMatch = cms.bool(False)
 
         print("nanoAOD_PuppiV15_switch.reclusterJets is true")
-    
+
     runMetCorAndUncFromMiniAOD(process,isData=isData,metType="Puppi",postfix="Puppi",jetFlavor="AK4PFPuppi", recoMetFromPFCs=bool(nanoAOD_PuppiV15_switch.recoMetFromPFCs), reclusterJets=bool(nanoAOD_PuppiV15_switch.reclusterJets))
     process.nanoSequenceCommon.insert(2,cms.Sequence(process.puppiMETSequence+process.fullPatMetSequencePuppi))
 
@@ -393,10 +393,10 @@ def nanoAOD_customizeMC(process):
         modifier.toModify(process, lambda p: nanoAOD_runMETfixEE2017(p,isData=False))
     return process
 
-###increasing the precision of selected GenParticles.                                                                                                 
+###increasing the precision of selected GenParticles.
 def nanoWmassGenCustomize(process):
     pdgSelection="?(abs(pdgId) == 11|| abs(pdgId)==13 || abs(pdgId)==15 ||abs(pdgId)== 12 || abs(pdgId)== 14 || abs(pdgId)== 16|| abs(pdgId)== 24|| pdgId== 23)"
-    # Keep precision same as default RECO for selected particles                                                                                       
+    # Keep precision same as default RECO for selected particles
     ptPrecision="{}?{}:{}".format(pdgSelection, CandVars.pt.precision.value(),genParticleTable.variables.pt.precision.value())
     process.genParticleTable.variables.pt.precision=cms.string(ptPrecision)
     phiPrecision="{} ? {} : {}".format(pdgSelection, CandVars.phi.precision.value(), genParticleTable.variables.phi.precision.value())

PhysicsTools/NanoAOD/python/protons_cff.py

@@ -1,5 +1,6 @@
 import FWCore.ParameterSet.Config as cms
 from PhysicsTools.NanoAOD.common_cff import *
+from PhysicsTools.NanoAOD.genProtonTable_cfi import genProtonTable as _genproton
 from PhysicsTools.NanoAOD.nano_eras_cff import *
 from RecoPPS.ProtonReconstruction.ppsFilteredProtonProducer_cfi import *
 
@@ -53,12 +54,21 @@
         thetaY = Var("thetaY",float,doc="th y",precision=10),
     ),
     externalVariables = cms.PSet(
-        decRPId = ExtVar("protonTable:protonRPId",int,doc="Detector ID",precision=8), 
+        decRPId = ExtVar("protonTable:protonRPId",int,doc="Detector ID",precision=8),
     ),
 )
 
 protonTablesTask = cms.Task(filteredProtons,protonTable,multiRPTable)
 if singleRPProtons: protonTablesTask.add(singleRPTable)
 
+# GEN-level signal/PU protons collection
+genProtonTable = _genproton.clone(
+    cut = cms.string('(pdgId == 2212) && (abs(pz) > 5200) && (abs(pz) < 6467.5)') # xi in [0.015, 0.2]
+)
+
+genProtonTablesTask = cms.Task(genProtonTable)
+
 for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94XMiniAODv1, run2_nanoAOD_94XMiniAODv2, run2_nanoAOD_94X2016, run2_nanoAOD_102Xv1:
     modifier.toReplaceWith(protonTablesTask, cms.Task())
+    # input GEN-level PU protons collection only introduced for UL and 12_X_Y
+    modifier.toReplaceWith(genProtonTablesTask, cms.Task())