File indexing completed on 2023-03-17 11:16:04
0001
0002
0003
0004
0005
0006
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};
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
0056 std::vector<float> pxs, pys, pzs, vzs;
0057 std::vector<bool> isPUs;
0058
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
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);