PatUtils/interface/SmearedJetProducerT.h

0001 #ifndef PAT_SMEAREDJETPRODUCERT_H
0002 #define PAT_SMEAREDJETPRODUCERT_H
0003
0004 /** \class SmearedJetProducerT
0005  *
0006  * Produce collection of "smeared" jets.
0007  *
0008  * The aim of this correction is to account for the difference in jet energy resolution
0009  * between Monte Carlo simulation and Data.
0010  *
0011  * \author Sébastien Brochet
0012  *
0013  */
0014
0015 #include "CommonTools/Utils/interface/PtComparator.h"
0016
0017 #include "FWCore/Framework/interface/Frameworkfwd.h"
0018 #include "FWCore/Framework/interface/stream/EDProducer.h"
0019 #include "FWCore/Framework/interface/Event.h"
0020 #include "FWCore/Framework/interface/ConsumesCollector.h"
0021 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0022 #include "FWCore/Utilities/interface/InputTag.h"
0023 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0024 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0025
0026 #include "DataFormats/JetReco/interface/GenJet.h"
0027 #include "DataFormats/JetReco/interface/GenJetCollection.h"
0028 #include "DataFormats/Math/interface/deltaR.h"
0029
0030 #include "JetMETCorrections/Modules/interface/JetResolution.h"
0031
0032 #include <TFile.h>
0033 #include <TH1.h>
0034 #include <TH1F.h>
0035
0036 #include <memory>
0037 #include <random>
0038
0039 namespace pat {
0040   class GenJetMatcher {
0041   public:
0042     GenJetMatcher(const edm::ParameterSet& cfg, edm::ConsumesCollector&& collector)
0043         : m_genJetsToken(collector.consumes<reco::GenJetCollection>(cfg.getParameter<edm::InputTag>("genJets"))),
0044           m_dR_max(cfg.getParameter<double>("dRMax")),
0045           m_dPt_max_factor(cfg.getParameter<double>("dPtMaxFactor")) {
0046       // Empty
0047     }
0048
0049     static void fillDescriptions(edm::ParameterSetDescription& desc) {
0050       desc.add<edm::InputTag>("genJets");
0051       desc.add<double>("dRMax");
0052       desc.add<double>("dPtMaxFactor");
0053     }
0054
0055     void getTokens(const edm::Event& event) { event.getByToken(m_genJetsToken, m_genJets); }
0056
0057     template <class T>
0058     const reco::GenJet* match(const T& jet, double resolution) {
0059       const reco::GenJetCollection& genJets = *m_genJets;
0060
0061       // Try to find a gen jet matching
0062       // dR < m_dR_max
0063       // dPt < m_dPt_max_factor * resolution
0064
0065       double min_dR = std::numeric_limits<double>::infinity();
0066       const reco::GenJet* matched_genJet = nullptr;
0067
0068       for (const auto& genJet : genJets) {
0069         double dR = deltaR(genJet, jet);
0070
0071         if (dR > min_dR)
0072           continue;
0073
0074         if (dR < m_dR_max) {
0075           double dPt = std::abs(genJet.pt() - jet.pt());
0076           if (dPt > m_dPt_max_factor * resolution)
0077             continue;
0078
0079           min_dR = dR;
0080           matched_genJet = &genJet;
0081         }
0082       }
0083
0084       return matched_genJet;
0085     }
0086
0087   private:
0088     edm::EDGetTokenT<reco::GenJetCollection> m_genJetsToken;
0089     edm::Handle<reco::GenJetCollection> m_genJets;
0090
0091     double m_dR_max;
0092     double m_dPt_max_factor;
0093   };
0094 };  // namespace pat
0095
0096 template <typename T>
0097 class SmearedJetProducerT : public edm::stream::EDProducer<> {
0098   using JetCollection = std::vector<T>;
0099
0100 public:
0101   explicit SmearedJetProducerT(const edm::ParameterSet& cfg)
0102       : m_enabled(cfg.getParameter<bool>("enabled")),
0103         m_useDeterministicSeed(cfg.getParameter<bool>("useDeterministicSeed")),
0104         m_debug(cfg.getUntrackedParameter<bool>("debug", false)) {
0105     m_jets_token = consumes<JetCollection>(cfg.getParameter<edm::InputTag>("src"));
0106
0107     if (m_enabled) {
0108       m_rho_token = consumes<double>(cfg.getParameter<edm::InputTag>("rho"));
0109
0110       m_use_txt_files = cfg.exists("resolutionFile") && cfg.exists("scaleFactorFile");
0111
0112       if (m_use_txt_files) {
0113         std::string resolutionFile = cfg.getParameter<edm::FileInPath>("resolutionFile").fullPath();
0114         std::string scaleFactorFile = cfg.getParameter<edm::FileInPath>("scaleFactorFile").fullPath();
0115
0116         m_resolution_from_file = std::make_unique<JME::JetResolution>(resolutionFile);
0117         m_scale_factor_from_file = std::make_unique<JME::JetResolutionScaleFactor>(scaleFactorFile);
0118       } else {
0119         m_jets_algo = cfg.getParameter<std::string>("algo");
0120         m_jets_algo_pt = cfg.getParameter<std::string>("algopt");
0121         m_jets_algo_token = esConsumes(edm::ESInputTag("", m_jets_algo));
0122         m_jets_algo_pt_token = esConsumes(edm::ESInputTag("", m_jets_algo_pt));
0123       }
0124
0125       std::uint32_t seed = cfg.getParameter<std::uint32_t>("seed");
0126       m_random_generator = std::mt19937(seed);
0127
0128       bool skipGenMatching = cfg.getParameter<bool>("skipGenMatching");
0129       if (!skipGenMatching)
0130         m_genJetMatcher = std::make_shared<pat::GenJetMatcher>(cfg, consumesCollector());
0131
0132       std::int32_t variation = cfg.getParameter<std::int32_t>("variation");
0133       m_uncertaintySource = cfg.getParameter<std::string>("uncertaintySource");
0134       m_nomVar = 1;
0135       if (variation == 0)
0136         m_systematic_variation = Variation::NOMINAL;
0137       else if (variation == 1)
0138         m_systematic_variation = Variation::UP;
0139       else if (variation == -1)
0140         m_systematic_variation = Variation::DOWN;
0141       else if (variation == 101) {
0142         m_systematic_variation = Variation::NOMINAL;
0143         m_nomVar = 1;
0144       } else if (variation == -101) {
0145         m_systematic_variation = Variation::NOMINAL;
0146         m_nomVar = -1;
0147       } else
0148         throw edm::Exception(edm::errors::ConfigFileReadError,
0149                              "Invalid value for 'variation' parameter. Only -1, 0, 1 or 101, -101 are supported.");
0150     }
0151
0152     produces<JetCollection>();
0153   }
0154
0155   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0156     edm::ParameterSetDescription desc;
0157
0158     desc.add<edm::InputTag>("src");
0159     desc.add<bool>("enabled");
0160     desc.add<edm::InputTag>("rho");
0161     desc.add<std::int32_t>("variation", 0);
0162     desc.add<std::string>("uncertaintySource", "");
0163     desc.add<std::uint32_t>("seed", 37428479);
0164     desc.add<bool>("skipGenMatching", false);
0165     desc.add<bool>("useDeterministicSeed", true);
0166     desc.addUntracked<bool>("debug", false);
0167
0168     auto source = (edm::ParameterDescription<std::string>("algo", true) and
0169                    edm::ParameterDescription<std::string>("algopt", true)) xor
0170                   (edm::ParameterDescription<edm::FileInPath>("resolutionFile", true) and
0171                    edm::ParameterDescription<edm::FileInPath>("scaleFactorFile", true));
0172     desc.addNode(std::move(source));
0173
0174     pat::GenJetMatcher::fillDescriptions(desc);
0175
0176     descriptions.addDefault(desc);
0177   }
0178
0179   void produce(edm::Event& event, const edm::EventSetup& setup) override {
0180     edm::Handle<JetCollection> jets_collection;
0181     event.getByToken(m_jets_token, jets_collection);
0182
0183     // Disable the module when running on real data
0184     if (m_enabled && event.isRealData()) {
0185       m_enabled = false;
0186       m_genJetMatcher.reset();
0187     }
0188
0189     edm::Handle<double> rho;
0190     if (m_enabled)
0191       event.getByToken(m_rho_token, rho);
0192
0193     JME::JetResolution resolution;
0194     JME::JetResolutionScaleFactor resolution_sf;
0195
0196     const JetCollection& jets = *jets_collection;
0197
0198     if (m_enabled) {
0199       if (m_use_txt_files) {
0200         resolution = *m_resolution_from_file;
0201         resolution_sf = *m_scale_factor_from_file;
0202       } else {
0203         resolution = JME::JetResolution::get(setup, m_jets_algo_pt_token);
0204         resolution_sf = JME::JetResolutionScaleFactor::get(setup, m_jets_algo_token);
0205       }
0206
0207       if (m_useDeterministicSeed) {
0208         unsigned int runNum_uint = static_cast<unsigned int>(event.id().run());
0209         unsigned int lumiNum_uint = static_cast<unsigned int>(event.id().luminosityBlock());
0210         unsigned int evNum_uint = static_cast<unsigned int>(event.id().event());
0211         unsigned int jet0eta = uint32_t(jets.empty() ? 0 : jets[0].eta() / 0.01);
0212         std::uint32_t seed = jet0eta + m_nomVar + (lumiNum_uint << 10) + (runNum_uint << 20) + evNum_uint;
0213         m_random_generator.seed(seed);
0214       }
0215     }
0216
0217     if (m_genJetMatcher)
0218       m_genJetMatcher->getTokens(event);
0219
0220     auto smearedJets = std::make_unique<JetCollection>();
0221
0222     for (const auto& jet : jets) {
0223       if ((!m_enabled) || (jet.pt() == 0)) {
0224         // Module disabled or invalid p4. Simply copy the input jet.
0225         smearedJets->push_back(jet);
0226
0227         continue;
0228       }
0229
0230       double jet_resolution = resolution.getResolution(
0231           {{JME::Binning::JetPt, jet.pt()}, {JME::Binning::JetEta, jet.eta()}, {JME::Binning::Rho, *rho}});
0232       double jer_sf = resolution_sf.getScaleFactor({{JME::Binning::JetPt, jet.pt()}, {JME::Binning::JetEta, jet.eta()}},
0233                                                    m_systematic_variation,
0234                                                    m_uncertaintySource);
0235
0236       if (m_debug) {
0237         std::cout << "jet:  pt: " << jet.pt() << "  eta: " << jet.eta() << "  phi: " << jet.phi()
0238                   << "  e: " << jet.energy() << std::endl;
0239         std::cout << "resolution: " << jet_resolution << std::endl;
0240         std::cout << "resolution scale factor: " << jer_sf << std::endl;
0241       }
0242
0243       const reco::GenJet* genJet = nullptr;
0244       if (m_genJetMatcher)
0245         genJet = m_genJetMatcher->match(jet, jet.pt() * jet_resolution);
0246
0247       double smearFactor = 1.;
0248
0249       if (genJet) {
0250         /*
0251                      * Case 1: we have a "good" gen jet matched to the reco jet
0252                      */
0253
0254         if (m_debug) {
0255           std::cout << "gen jet:  pt: " << genJet->pt() << "  eta: " << genJet->eta() << "  phi: " << genJet->phi()
0256                     << "  e: " << genJet->energy() << std::endl;
0257         }
0258
0259         double dPt = jet.pt() - genJet->pt();
0260         smearFactor = 1 + m_nomVar * (jer_sf - 1.) * dPt / jet.pt();
0261
0262       } else if (jer_sf > 1) {
0263         /*
0264                      * Case 2: we don't have a gen jet. Smear jet pt using a random gaussian variation
0265                      */
0266
0267         double sigma = jet_resolution * std::sqrt(jer_sf * jer_sf - 1);
0268         if (m_debug) {
0269           std::cout << "gaussian width: " << sigma << std::endl;
0270         }
0271
0272         std::normal_distribution<> d(0, sigma);
0273         smearFactor = 1. + m_nomVar * d(m_random_generator);
0274       } else if (m_debug) {
0275         std::cout << "Impossible to smear this jet" << std::endl;
0276       }
0277
0278       if (jet.energy() * smearFactor < MIN_JET_ENERGY) {
0279         // Negative or too small smearFactor. We would change direction of the jet
0280         // and this is not what we want.
0281         // Recompute the smearing factor in order to have jet.energy() == MIN_JET_ENERGY
0282         double newSmearFactor = MIN_JET_ENERGY / jet.energy();
0283         if (m_debug) {
0284           std::cout << "The smearing factor (" << smearFactor << ") is either negative or too small. Fixing it to "
0285                     << newSmearFactor << " to avoid change of direction." << std::endl;
0286         }
0287         smearFactor = newSmearFactor;
0288       }
0289
0290       T smearedJet = jet;
0291       smearedJet.scaleEnergy(smearFactor);
0292
0293       if (m_debug) {
0294         std::cout << "smeared jet (" << smearFactor << "):  pt: " << smearedJet.pt() << "  eta: " << smearedJet.eta()
0295                   << "  phi: " << smearedJet.phi() << "  e: " << smearedJet.energy() << std::endl;
0296       }
0297
0298       smearedJets->push_back(smearedJet);
0299     }
0300
0301     // Sort jets by pt
0302     std::sort(smearedJets->begin(), smearedJets->end(), jetPtComparator);
0303
0304     event.put(std::move(smearedJets));
0305   }
0306
0307 private:
0308   static constexpr const double MIN_JET_ENERGY = 1e-2;
0309
0310   edm::EDGetTokenT<JetCollection> m_jets_token;
0311   edm::EDGetTokenT<double> m_rho_token;
0312   bool m_enabled;
0313   std::string m_jets_algo_pt;
0314   std::string m_jets_algo;
0315   JME::JetResolution::Token m_jets_algo_pt_token;
0316   JME::JetResolutionScaleFactor::Token m_jets_algo_token;
0317   Variation m_systematic_variation;
0318   std::string m_uncertaintySource;
0319   bool m_useDeterministicSeed;
0320   bool m_debug;
0321   std::shared_ptr<pat::GenJetMatcher> m_genJetMatcher;
0322
0323   bool m_use_txt_files;
0324   std::unique_ptr<JME::JetResolution> m_resolution_from_file;
0325   std::unique_ptr<JME::JetResolutionScaleFactor> m_scale_factor_from_file;
0326
0327   std::mt19937 m_random_generator;
0328
0329   GreaterByPt<T> jetPtComparator;
0330
0331   int m_nomVar;
0332 };
0333 #endif