NanoAOD/plugins/ProtonProducer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    PhysicsTools/NanoAOD
0004 // Class:      ProtonProducer
0005 //
0006 /**\class ProtonProducer ProtonProducer.cc PhysicsTools/NanoAOD/plugins/ProtonProducer.cc
0007  Description: Realavent proton variables for analysis usage
0008  Implementation:
0009 */
0010 //
0011 // Original Author:  Justin Williams
0012 //         Created: 04 Jul 2019 15:27:53 GMT
0013 //
0014 //
0015
0016 // system include files
0017 #include <memory>
0018 #include <map>
0019 #include <string>
0020 #include <vector>
0021 #include <iostream>
0022
0023 // user include files
0024 #include "FWCore/Framework/interface/Frameworkfwd.h"
0025 #include "FWCore/Framework/interface/global/EDProducer.h"
0026
0027 #include "FWCore/Framework/interface/Event.h"
0028 #include "FWCore/Framework/interface/MakerMacros.h"
0029
0030 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0031 #include "FWCore/Utilities/interface/StreamID.h"
0032
0033 #include "CommonTools/Egamma/interface/EffectiveAreas.h"
0034
0035 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0036 #include "DataFormats/Common/interface/ValueMap.h"
0037
0038 #include "DataFormats/CTPPSDetId/interface/CTPPSDetId.h"
0039 #include "DataFormats/CTPPSReco/interface/CTPPSLocalTrackLite.h"
0040 #include "DataFormats/ProtonReco/interface/ForwardProton.h"
0041 #include "DataFormats/ProtonReco/interface/ForwardProtonFwd.h"
0042 #include "DataFormats/CTPPSReco/interface/CTPPSPixelLocalTrack.h"
0043
0044 class ProtonProducer : public edm::global::EDProducer<> {
0045 public:
0046   ProtonProducer(edm::ParameterSet const &ps)
0047       : tokenRecoProtonsMultiRP_(
0048             mayConsume<reco::ForwardProtonCollection>(ps.getParameter<edm::InputTag>("tagRecoProtonsMulti"))),
0049         tokenRecoProtonsSingleRP_(
0050             mayConsume<reco::ForwardProtonCollection>(ps.getParameter<edm::InputTag>("tagRecoProtonsSingle"))),
0051         tokenTracksLite_(mayConsume<std::vector<CTPPSLocalTrackLite>>(ps.getParameter<edm::InputTag>("tagTrackLite"))),
0052         storeSingleRPProtons_(ps.getParameter<bool>("storeSingleRPProtons")) {
0053     produces<edm::ValueMap<int>>("arm");
0054     produces<nanoaod::FlatTable>("ppsTrackTable");
0055     if (storeSingleRPProtons_)
0056       produces<edm::ValueMap<int>>("protonRPId");
0057   }
0058   ~ProtonProducer() override {}
0059
0060   // ------------ method called to produce the data  ------------
0061   void produce(edm::StreamID id, edm::Event &iEvent, const edm::EventSetup &iSetup) const override {
0062     // Get Forward Proton handle
0063     auto hRecoProtonsMultiRP = iEvent.getHandle(tokenRecoProtonsMultiRP_);
0064     auto hRecoProtonsSingleRP = iEvent.getHandle(tokenRecoProtonsSingleRP_);
0065
0066     // Get PPS Local Track handle
0067     auto ppsTracksLite = iEvent.getHandle(tokenTracksLite_);
0068
0069     // book output variables for protons
0070     std::vector<int> multiRP_arm;
0071
0072     // book output variables for tracks
0073     std::vector<float> trackX, trackY, trackTime, trackTimeUnc;
0074     std::vector<int> multiRPProtonIdx, decRPId, rpType;
0075     std::vector<int> singleRPProtonIdx, singleRP_protonRPId;
0076
0077     if (storeSingleRPProtons_) {
0078       // process single-RP protons
0079       {
0080         const auto &num_proton = hRecoProtonsSingleRP->size();
0081         singleRP_protonRPId.reserve(num_proton);
0082
0083         for (const auto &proton : *hRecoProtonsSingleRP) {
0084           CTPPSDetId rpId((*proton.contributingLocalTracks().begin())->rpId());
0085           unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0086           singleRP_protonRPId.push_back(rpDecId);
0087         }
0088       }
0089
0090       // process multi-RP protons
0091       {
0092         const auto &num_proton = hRecoProtonsMultiRP->size();
0093         multiRP_arm.reserve(num_proton);
0094
0095         for (const auto &proton : *hRecoProtonsMultiRP) {
0096           multiRP_arm.push_back((proton.pz() < 0.) ? 1 : 0);
0097         }
0098       }
0099
0100       // process local tracks
0101       for (unsigned int tr_idx = 0; tr_idx < ppsTracksLite->size(); ++tr_idx) {
0102         const auto &tr = ppsTracksLite->at(tr_idx);
0103
0104         bool found = false;
0105
0106         CTPPSDetId rpId(tr.rpId());
0107         unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0108
0109         signed int singleRP_idx = -1;
0110         for (unsigned int p_idx = 0; p_idx < hRecoProtonsSingleRP->size(); ++p_idx) {
0111           const auto &proton = hRecoProtonsSingleRP->at(p_idx);
0112
0113           for (const auto &ref : proton.contributingLocalTracks()) {
0114             if (ref.key() == tr_idx) {
0115               singleRP_idx = p_idx;
0116               found = true;
0117             }
0118           }
0119         }
0120
0121         signed int multiRP_idx = -1;
0122         for (unsigned int p_idx = 0; p_idx < hRecoProtonsMultiRP->size(); ++p_idx) {
0123           const auto &proton = hRecoProtonsMultiRP->at(p_idx);
0124
0125           for (const auto &ref : proton.contributingLocalTracks()) {
0126             if (ref.key() == tr_idx) {
0127               multiRP_idx = p_idx;
0128               found = true;
0129             }
0130           }
0131         }
0132
0133         if (found) {
0134           singleRPProtonIdx.push_back(singleRP_idx);
0135           multiRPProtonIdx.push_back(multiRP_idx);
0136           decRPId.push_back(rpDecId);
0137           rpType.push_back(rpId.subdetId());
0138           trackX.push_back(tr.x());
0139           trackY.push_back(tr.y());
0140           trackTime.push_back(tr.time());
0141           trackTimeUnc.push_back(tr.timeUnc());
0142         }
0143       }
0144     }
0145
0146     else {
0147       for (unsigned int p_idx = 0; p_idx < hRecoProtonsMultiRP->size(); ++p_idx) {
0148         const auto &proton = hRecoProtonsMultiRP->at(p_idx);
0149         multiRP_arm.push_back((proton.pz() < 0.) ? 1 : 0);
0150
0151         for (const auto &ref : proton.contributingLocalTracks()) {
0152           multiRPProtonIdx.push_back(p_idx);
0153           CTPPSDetId rpId(ref->rpId());
0154           unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0155           decRPId.push_back(rpDecId);
0156           rpType.push_back(rpId.subdetId());
0157           trackX.push_back(ref->x());
0158           trackY.push_back(ref->y());
0159           trackTime.push_back(ref->time());
0160           trackTimeUnc.push_back(ref->timeUnc());
0161         }
0162       }
0163     }
0164
0165     // update proton tables
0166     std::unique_ptr<edm::ValueMap<int>> multiRP_armV(new edm::ValueMap<int>());
0167     edm::ValueMap<int>::Filler fillermultiArm(*multiRP_armV);
0168     fillermultiArm.insert(hRecoProtonsMultiRP, multiRP_arm.begin(), multiRP_arm.end());
0169     fillermultiArm.fill();
0170
0171     std::unique_ptr<edm::ValueMap<int>> protonRPIdV(new edm::ValueMap<int>());
0172     edm::ValueMap<int>::Filler fillerID(*protonRPIdV);
0173     if (storeSingleRPProtons_) {
0174       fillerID.insert(hRecoProtonsSingleRP, singleRP_protonRPId.begin(), singleRP_protonRPId.end());
0175       fillerID.fill();
0176     }
0177
0178     // build track table
0179     auto ppsTab = std::make_unique<nanoaod::FlatTable>(trackX.size(), "PPSLocalTrack", false);
0180     ppsTab->addColumn<int>("multiRPProtonIdx", multiRPProtonIdx, "local track - proton correspondence");
0181     if (storeSingleRPProtons_)
0182       ppsTab->addColumn<int>("singleRPProtonIdx", singleRPProtonIdx, "local track - proton correspondence");
0183     ppsTab->addColumn<float>("x", trackX, "local track x", 16);
0184     ppsTab->addColumn<float>("y", trackY, "local track y", 13);
0185     ppsTab->addColumn<float>("time", trackTime, "local track time", 16);
0186     ppsTab->addColumn<float>("timeUnc", trackTimeUnc, "local track time uncertainty", 13);
0187     ppsTab->addColumn<int>("decRPId", decRPId, "local track detector dec id");
0188     ppsTab->addColumn<int>("rpType", rpType, "strip=3, pixel=4, diamond=5, timing=6");
0189     ppsTab->setDoc("ppsLocalTrack variables");
0190
0191     // save output
0192     iEvent.put(std::move(multiRP_armV), "arm");
0193     iEvent.put(std::move(ppsTab), "ppsTrackTable");
0194     if (storeSingleRPProtons_)
0195       iEvent.put(std::move(protonRPIdV), "protonRPId");
0196   }
0197
0198   // ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
0199   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0200     edm::ParameterSetDescription desc;
0201     desc.add<edm::InputTag>("tagRecoProtonsMulti")->setComment("multiRP proton collection");
0202     desc.add<edm::InputTag>("tagRecoProtonsSingle")->setComment("singleRP proton collection");
0203     desc.add<edm::InputTag>("tagTrackLite")->setComment("pps local tracks lite collection");
0204     desc.add<bool>("storeSingleRPProtons")->setComment("flag to store singleRP protons and associated tracks");
0205     descriptions.add("ProtonProducer", desc);
0206   }
0207
0208 protected:
0209   const edm::EDGetTokenT<reco::ForwardProtonCollection> tokenRecoProtonsMultiRP_;
0210   const edm::EDGetTokenT<reco::ForwardProtonCollection> tokenRecoProtonsSingleRP_;
0211   const edm::EDGetTokenT<std::vector<CTPPSLocalTrackLite>> tokenTracksLite_;
0212   const bool storeSingleRPProtons_;
0213 };
0214
0215 DEFINE_FWK_MODULE(ProtonProducer);