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     edm::Handle<reco::ForwardProtonCollection> hRecoProtonsMultiRP;
0064     iEvent.getByToken(tokenRecoProtonsMultiRP_, hRecoProtonsMultiRP);
0065
0066     edm::Handle<reco::ForwardProtonCollection> hRecoProtonsSingleRP;
0067     iEvent.getByToken(tokenRecoProtonsSingleRP_, hRecoProtonsSingleRP);
0068
0069     // Get PPS Local Track handle
0070     edm::Handle<std::vector<CTPPSLocalTrackLite>> ppsTracksLite;
0071     iEvent.getByToken(tokenTracksLite_, ppsTracksLite);
0072
0073     // book output variables for protons
0074     std::vector<int> multiRP_arm;
0075
0076     // book output variables for tracks
0077     std::vector<float> trackX, trackY, trackTime, trackTimeUnc;
0078     std::vector<int> multiRPProtonIdx, decRPId, rpType;
0079     std::vector<int> singleRPProtonIdx, singleRP_protonRPId;
0080
0081     if (storeSingleRPProtons_) {
0082       // process single-RP protons
0083       {
0084         const auto &num_proton = hRecoProtonsSingleRP->size();
0085         singleRP_protonRPId.reserve(num_proton);
0086
0087         for (const auto &proton : *hRecoProtonsSingleRP) {
0088           CTPPSDetId rpId((*proton.contributingLocalTracks().begin())->rpId());
0089           unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0090           singleRP_protonRPId.push_back(rpDecId);
0091         }
0092       }
0093
0094       // process multi-RP protons
0095       {
0096         const auto &num_proton = hRecoProtonsMultiRP->size();
0097         multiRP_arm.reserve(num_proton);
0098
0099         for (const auto &proton : *hRecoProtonsMultiRP) {
0100           multiRP_arm.push_back((proton.pz() < 0.) ? 1 : 0);
0101         }
0102       }
0103
0104       // process local tracks
0105       for (unsigned int tr_idx = 0; tr_idx < ppsTracksLite->size(); ++tr_idx) {
0106         const auto &tr = ppsTracksLite->at(tr_idx);
0107
0108         bool found = false;
0109
0110         CTPPSDetId rpId(tr.rpId());
0111         unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0112
0113         signed int singleRP_idx = -1;
0114         for (unsigned int p_idx = 0; p_idx < hRecoProtonsSingleRP->size(); ++p_idx) {
0115           const auto &proton = hRecoProtonsSingleRP->at(p_idx);
0116
0117           for (const auto &ref : proton.contributingLocalTracks()) {
0118             if (ref.key() == tr_idx) {
0119               singleRP_idx = p_idx;
0120               found = true;
0121             }
0122           }
0123         }
0124
0125         signed int multiRP_idx = -1;
0126         for (unsigned int p_idx = 0; p_idx < hRecoProtonsMultiRP->size(); ++p_idx) {
0127           const auto &proton = hRecoProtonsMultiRP->at(p_idx);
0128
0129           for (const auto &ref : proton.contributingLocalTracks()) {
0130             if (ref.key() == tr_idx) {
0131               multiRP_idx = p_idx;
0132               found = true;
0133             }
0134           }
0135         }
0136
0137         if (found) {
0138           singleRPProtonIdx.push_back(singleRP_idx);
0139           multiRPProtonIdx.push_back(multiRP_idx);
0140           decRPId.push_back(rpDecId);
0141           rpType.push_back(rpId.subdetId());
0142           trackX.push_back(tr.x());
0143           trackY.push_back(tr.y());
0144           trackTime.push_back(tr.time());
0145           trackTimeUnc.push_back(tr.timeUnc());
0146         }
0147       }
0148     }
0149
0150     else {
0151       for (unsigned int p_idx = 0; p_idx < hRecoProtonsMultiRP->size(); ++p_idx) {
0152         const auto &proton = hRecoProtonsMultiRP->at(p_idx);
0153         multiRP_arm.push_back((proton.pz() < 0.) ? 1 : 0);
0154
0155         for (const auto &ref : proton.contributingLocalTracks()) {
0156           multiRPProtonIdx.push_back(p_idx);
0157           CTPPSDetId rpId(ref->rpId());
0158           unsigned int rpDecId = rpId.arm() * 100 + rpId.station() * 10 + rpId.rp();
0159           decRPId.push_back(rpDecId);
0160           rpType.push_back(rpId.subdetId());
0161           trackX.push_back(ref->x());
0162           trackY.push_back(ref->y());
0163           trackTime.push_back(ref->time());
0164           trackTimeUnc.push_back(ref->timeUnc());
0165         }
0166       }
0167     }
0168
0169     // update proton tables
0170     std::unique_ptr<edm::ValueMap<int>> multiRP_armV(new edm::ValueMap<int>());
0171     edm::ValueMap<int>::Filler fillermultiArm(*multiRP_armV);
0172     fillermultiArm.insert(hRecoProtonsMultiRP, multiRP_arm.begin(), multiRP_arm.end());
0173     fillermultiArm.fill();
0174
0175     std::unique_ptr<edm::ValueMap<int>> protonRPIdV(new edm::ValueMap<int>());
0176     edm::ValueMap<int>::Filler fillerID(*protonRPIdV);
0177     if (storeSingleRPProtons_) {
0178       fillerID.insert(hRecoProtonsSingleRP, singleRP_protonRPId.begin(), singleRP_protonRPId.end());
0179       fillerID.fill();
0180     }
0181
0182     // build track table
0183     auto ppsTab = std::make_unique<nanoaod::FlatTable>(trackX.size(), "PPSLocalTrack", false);
0184     ppsTab->addColumn<int>("multiRPProtonIdx", multiRPProtonIdx, "local track - proton correspondence");
0185     if (storeSingleRPProtons_)
0186       ppsTab->addColumn<int>("singleRPProtonIdx", singleRPProtonIdx, "local track - proton correspondence");
0187     ppsTab->addColumn<float>("x", trackX, "local track x", 16);
0188     ppsTab->addColumn<float>("y", trackY, "local track y", 13);
0189     ppsTab->addColumn<float>("time", trackTime, "local track time", 16);
0190     ppsTab->addColumn<float>("timeUnc", trackTimeUnc, "local track time uncertainty", 13);
0191     ppsTab->addColumn<int>("decRPId", decRPId, "local track detector dec id");
0192     ppsTab->addColumn<int>("rpType", rpType, "strip=3, pixel=4, diamond=5, timing=6");
0193     ppsTab->setDoc("ppsLocalTrack variables");
0194
0195     // save output
0196     iEvent.put(std::move(multiRP_armV), "arm");
0197     iEvent.put(std::move(ppsTab), "ppsTrackTable");
0198     if (storeSingleRPProtons_)
0199       iEvent.put(std::move(protonRPIdV), "protonRPId");
0200   }
0201
0202   // ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
0203   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0204     edm::ParameterSetDescription desc;
0205     desc.add<edm::InputTag>("tagRecoProtonsMulti")->setComment("multiRP proton collection");
0206     desc.add<edm::InputTag>("tagRecoProtonsSingle")->setComment("singleRP proton collection");
0207     desc.add<edm::InputTag>("tagTrackLite")->setComment("pps local tracks lite collection");
0208     desc.add<bool>("storeSingleRPProtons")->setComment("flag to store singleRP protons and associated tracks");
0209     descriptions.add("ProtonProducer", desc);
0210   }
0211
0212 protected:
0213   const edm::EDGetTokenT<reco::ForwardProtonCollection> tokenRecoProtonsMultiRP_;
0214   const edm::EDGetTokenT<reco::ForwardProtonCollection> tokenRecoProtonsSingleRP_;
0215   const edm::EDGetTokenT<std::vector<CTPPSLocalTrackLite>> tokenTracksLite_;
0216   const bool storeSingleRPProtons_;
0217 };
0218
0219 DEFINE_FWK_MODULE(ProtonProducer);