JetCorrector/plugins/JetTracksAssociationToTrackRefs.cc

0001 #include "FWCore/Framework/interface/global/EDProducer.h"
0002 #include "FWCore/Framework/interface/MakerMacros.h"
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0005 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0007
0008 #include "DataFormats/Common/interface/EDProductfwd.h"
0009 #include "DataFormats/TrackReco/interface/Track.h"
0010 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0011 #include "DataFormats/JetReco/interface/Jet.h"
0012 #include "DataFormats/JetReco/interface/JetTracksAssociation.h"
0013
0014 #include "JetMETCorrections/JetCorrector/interface/JetCorrector.h"
0015
0016 #include <unordered_set>
0017
0018 /**
0019  * The purpose of this producer is to convert
0020  * AssociationVector<JetRefBaseProd, vector<RefVector<Track> > to a
0021  * RefVector<Track> of the (unique) values.
0022  */
0023 class JetTracksAssociationToTrackRefs : public edm::global::EDProducer<> {
0024 public:
0025   JetTracksAssociationToTrackRefs(const edm::ParameterSet& iConfig);
0026
0027   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0028
0029   void produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
0030
0031 private:
0032   edm::EDGetTokenT<reco::JetTracksAssociation::Container> associationToken_;
0033   edm::EDGetTokenT<edm::View<reco::Jet>> jetToken_;
0034   edm::EDGetTokenT<reco::JetCorrector> correctorToken_;
0035   const double ptMin_;
0036 };
0037
0038 JetTracksAssociationToTrackRefs::JetTracksAssociationToTrackRefs(const edm::ParameterSet& iConfig)
0039     : associationToken_(
0040           consumes<reco::JetTracksAssociation::Container>(iConfig.getParameter<edm::InputTag>("association"))),
0041       jetToken_(consumes<edm::View<reco::Jet>>(iConfig.getParameter<edm::InputTag>("jets"))),
0042       correctorToken_(consumes<reco::JetCorrector>(iConfig.getParameter<edm::InputTag>("corrector"))),
0043       ptMin_(iConfig.getParameter<double>("correctedPtMin")) {
0044   produces<reco::TrackRefVector>();
0045 }
0046
0047 void JetTracksAssociationToTrackRefs::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0048   edm::ParameterSetDescription desc;
0049   desc.add<edm::InputTag>("association", edm::InputTag("ak4JetTracksAssociatorAtVertexPF"));
0050   desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJetsCHS"));
0051   desc.add<edm::InputTag>("corrector", edm::InputTag("ak4PFL1FastL2L3Corrector"));
0052   desc.add<double>("correctedPtMin", 0);
0053   descriptions.add("jetTracksAssociationToTrackRefs", desc);
0054 }
0055
0056 void JetTracksAssociationToTrackRefs::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0057   edm::Handle<reco::JetTracksAssociation::Container> h_assoc;
0058   iEvent.getByToken(associationToken_, h_assoc);
0059   const reco::JetTracksAssociation::Container& association = *h_assoc;
0060
0061   edm::Handle<edm::View<reco::Jet>> h_jets;
0062   iEvent.getByToken(jetToken_, h_jets);
0063   const edm::View<reco::Jet>& jets = *h_jets;
0064
0065   edm::Handle<reco::JetCorrector> h_corrector;
0066   iEvent.getByToken(correctorToken_, h_corrector);
0067   const reco::JetCorrector& corrector = *h_corrector;
0068
0069   auto ret = std::make_unique<reco::TrackRefVector>();
0070   std::unordered_set<reco::TrackRefVector::key_type> alreadyAdded;
0071
0072   // Exctract tracks only for jets passing certain pT threshold after
0073   // correction
0074   for (size_t i = 0; i < jets.size(); ++i) {
0075     edm::RefToBase<reco::Jet> jetRef = jets.refAt(i);
0076     const reco::Jet& jet = *jetRef;
0077
0078     auto p4 = jet.p4();
0079
0080     // Energy correction in the most general way
0081     if (!corrector.vectorialCorrection()) {
0082       double scale = 1;
0083       if (!corrector.refRequired()) {
0084         scale = corrector.correction(jet);
0085       } else {
0086         scale = corrector.correction(jet, jetRef);
0087       }
0088       p4 = p4 * scale;
0089     } else {
0090       corrector.correction(jet, jetRef, p4);
0091     }
0092
0093     if (p4.pt() <= ptMin_)
0094       continue;
0095
0096     for (const auto& trackRef : association[jetRef]) {
0097       if (alreadyAdded.find(trackRef.key()) == alreadyAdded.end()) {
0098         ret->push_back(trackRef);
0099         alreadyAdded.insert(trackRef.key());
0100       }
0101     }
0102   }
0103
0104   iEvent.put(std::move(ret));
0105 }
0106
0107 DEFINE_FWK_MODULE(JetTracksAssociationToTrackRefs);