NanoAOD/plugins/GenProtonTableProducer.cc

0001 /****************************************************************************
0002  *
0003  * This is a part of PPS offline software.
0004  * Authors:
0005  *   Laurent Forthomme
0006  *   Michael Pitt
0007  *
0008  ****************************************************************************/
0009
0010 #include <memory>
0011
0012 #include "FWCore/Framework/interface/Frameworkfwd.h"
0013 #include "FWCore/Framework/interface/stream/EDProducer.h"
0014 #include "FWCore/Framework/interface/Event.h"
0015 #include "FWCore/Framework/interface/MakerMacros.h"
0016 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0017 #include "FWCore/Utilities/interface/StreamID.h"
0018
0019 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0020
0021 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0022 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0023
0024 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0025
0026 class GenProtonTableProducer : public edm::stream::EDProducer<> {
0027 public:
0028   explicit GenProtonTableProducer(const edm::ParameterSet&);
0029   ~GenProtonTableProducer() override = default;
0030
0031   static void fillDescriptions(edm::ConfigurationDescriptions&);
0032
0033 private:
0034   void produce(edm::Event&, const edm::EventSetup&) override;
0035
0036   const edm::EDGetTokenT<reco::GenParticleCollection> prunedCandsToken_;
0037   const edm::EDGetTokenT<reco::GenParticleCollection> puCandsToken_, puAltCandsToken_;
0038   const StringCutObjectSelector<reco::Candidate> protonsCut_;
0039   const std::string table_name_;
0040   const double tolerance_;
0041   bool use_alt_coll_{false};  ///< Are we using premix/mix collection name for PU protons?
0042 };
0043
0044 GenProtonTableProducer::GenProtonTableProducer(const edm::ParameterSet& iConfig)
0045     : prunedCandsToken_(consumes<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcPruned"))),
0046       puCandsToken_(mayConsume<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcPUProtons"))),
0047       puAltCandsToken_(mayConsume<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("srcAltPUProtons"))),
0048       protonsCut_(iConfig.getParameter<std::string>("cut")),
0049       table_name_(iConfig.getParameter<std::string>("name")),
0050       tolerance_(iConfig.getParameter<double>("tolerance")) {
0051   produces<nanoaod::FlatTable>();
0052 }
0053
0054 void GenProtonTableProducer::produce(edm::Event& iEvent, const edm::EventSetup&) {
0055   // define the variables
0056   std::vector<float> pxs, pys, pzs, vzs;
0057   std::vector<bool> isPUs;
0058   // first loop over signal protons
0059   for (const auto& pruned_cand : iEvent.get(prunedCandsToken_)) {
0060     if (!protonsCut_(pruned_cand))
0061       continue;
0062     pxs.emplace_back(pruned_cand.px());
0063     pys.emplace_back(pruned_cand.py());
0064     pzs.emplace_back(pruned_cand.pz());
0065     vzs.emplace_back(pruned_cand.vz());
0066     isPUs.emplace_back(false);
0067   }
0068   // then loop over pruned candidates ; if already in signal protons, discard
0069   edm::Handle<reco::GenParticleCollection> hPUCands;
0070   if (use_alt_coll_ || !iEvent.getByToken(puCandsToken_, hPUCands))
0071     use_alt_coll_ = iEvent.getByToken(puAltCandsToken_, hPUCands);
0072   for (const auto& pu_cand : *hPUCands) {
0073     if (!protonsCut_(pu_cand))
0074       continue;
0075     bool associated{false};
0076     for (size_t i = 0; i < pzs.size(); ++i) {
0077       if (fabs(1. - pxs.at(i) / pu_cand.px()) < tolerance_ && fabs(1. - pys.at(i) / pu_cand.py()) < tolerance_ &&
0078           fabs(1. - pzs.at(i) / pu_cand.pz()) < tolerance_) {
0079         associated = true;
0080         break;
0081       }
0082     }
0083     if (associated)
0084       continue;
0085     pxs.emplace_back(pu_cand.px());
0086     pys.emplace_back(pu_cand.py());
0087     pzs.emplace_back(pu_cand.pz());
0088     vzs.emplace_back(pu_cand.vz());
0089     isPUs.emplace_back(true);
0090   }
0091
0092   auto protons_table = std::make_unique<nanoaod::FlatTable>(isPUs.size(), table_name_, false);
0093   protons_table->addColumn<float>("px", pxs, "proton horizontal momentum", 8);
0094   protons_table->addColumn<float>("py", pys, "proton vertical momentum", 8);
0095   protons_table->addColumn<float>("pz", pzs, "proton longitudinal momentum", 8);
0096   protons_table->addColumn<float>("vz", vzs, "proton vertex longitudinal coordinate", 8);
0097   protons_table->addColumn<bool>("isPU", isPUs, "pileup proton?");
0098   iEvent.put(std::move(protons_table));
0099 }
0100
0101 void GenProtonTableProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0102   edm::ParameterSetDescription desc;
0103   desc.add<edm::InputTag>("srcPruned", edm::InputTag("prunedGenParticles"))
0104       ->setComment("input source for pruned gen-level particle candidates");
0105   desc.add<edm::InputTag>("srcPUProtons", edm::InputTag("genPUProtons"))
0106       ->setComment("input source for pileup protons collection");
0107   desc.add<edm::InputTag>("srcAltPUProtons", edm::InputTag("genPUProtons", "genPUProtons"))
0108       ->setComment("alternative input source for pileup protons collection (for premix-mix backward compatibility)");
0109   desc.add<std::string>("cut", "")->setComment("proton kinematic selection");
0110   desc.add<std::string>("name", "GenProton")->setComment("flat table name");
0111   desc.add<std::string>("doc", "generator level information on (signal+PU) protons")
0112       ->setComment("flat table description");
0113   desc.add<double>("tolerance", 1.e-3)->setComment("relative difference between the signal and pileup protons momenta");
0114   descriptions.add("genProtonTable", desc);
0115 }
0116
0117 DEFINE_FWK_MODULE(GenProtonTableProducer);