PatAlgos/plugins/LeptonUpdater.cc

0001 #include "FWCore/Framework/interface/stream/EDProducer.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0006 #include "FWCore/Utilities/interface/InputTag.h"
0007 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0008 #include "PhysicsTools/PatUtils/interface/MiniIsolation.h"
0009 #include "PhysicsTools/PatAlgos/interface/SoftMuonMvaRun3Estimator.h"
0010 #include "PhysicsTools/XGBoost/interface/XGBooster.h"
0011 #include "DataFormats/Common/interface/View.h"
0012
0013 #include "DataFormats/PatCandidates/interface/Muon.h"
0014 #include "DataFormats/PatCandidates/interface/Electron.h"
0015 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0016 #include "DataFormats/MuonReco/interface/MuonSelectors.h"
0017
0018 #include <type_traits>
0019
0020 namespace pat {
0021
0022   template <typename T>
0023   class LeptonUpdater : public edm::stream::EDProducer<> {
0024   public:
0025     explicit LeptonUpdater(const edm::ParameterSet &iConfig)
0026         : src_(consumes<std::vector<T>>(iConfig.getParameter<edm::InputTag>("src"))),
0027           vertices_(consumes<std::vector<reco::Vertex>>(iConfig.getParameter<edm::InputTag>("vertices"))),
0028           beamLineToken_(consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("beamspot"))),
0029           computeMiniIso_(iConfig.getParameter<bool>("computeMiniIso")),
0030           fixDxySign_(iConfig.getParameter<bool>("fixDxySign")) {
0031       //for mini-isolation calculation
0032       if (computeMiniIso_) {
0033         readMiniIsoParams(iConfig);
0034         pcToken_ = consumes<pat::PackedCandidateCollection>(iConfig.getParameter<edm::InputTag>("pfCandsForMiniIso"));
0035       }
0036       recomputeMuonBasicSelectors_ = false;
0037       recomputeSoftMuonMvaRun3_ = false;
0038       if (typeid(T) == typeid(pat::Muon)) {
0039         recomputeMuonBasicSelectors_ = iConfig.getParameter<bool>("recomputeMuonBasicSelectors");
0040         recomputeSoftMuonMvaRun3_ = iConfig.getParameter<bool>("recomputeSoftMuonMvaRun3");
0041         if (recomputeSoftMuonMvaRun3_) {
0042           std::string softMvaRun3Model = iConfig.getParameter<std::string>("softMvaRun3Model");
0043           softMuonMvaRun3Booster_ =
0044               std::make_unique<pat::XGBooster>(edm::FileInPath(softMvaRun3Model + ".model").fullPath(),
0045                                                edm::FileInPath(softMvaRun3Model + ".features").fullPath());
0046         }
0047       }
0048       produces<std::vector<T>>();
0049     }
0050
0051     ~LeptonUpdater() override {}
0052
0053     void produce(edm::Event &, const edm::EventSetup &) override;
0054
0055     static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0056       edm::ParameterSetDescription desc;
0057       desc.add<edm::InputTag>("src")->setComment("Lepton collection");
0058       desc.add<edm::InputTag>("vertices")->setComment("Vertex collection");
0059       desc.add<edm::InputTag>("beamspot", edm::InputTag("offlineBeamSpot"))->setComment("Beam spot");
0060       desc.add<bool>("computeMiniIso", false)->setComment("Recompute miniIsolation");
0061       desc.add<bool>("fixDxySign", false)->setComment("Fix the IP sign");
0062       desc.addOptional<edm::InputTag>("pfCandsForMiniIso", edm::InputTag("packedPFCandidates"))
0063           ->setComment("PackedCandidate collection used for miniIso");
0064       if (typeid(T) == typeid(pat::Muon)) {
0065         desc.add<bool>("recomputeMuonBasicSelectors", false)
0066             ->setComment("Recompute basic cut-based muon selector flags");
0067         desc.add<bool>("recomputeSoftMuonMvaRun3", false)->setComment("Recompute Run3 soft muon MVA value");
0068         desc.add<std::string>("softMvaRun3Model", "RecoMuon/MuonIdentification/data/Run2022-20231030-1731-Event0")
0069             ->setComment("Run3 soft muon MVA model path");
0070         desc.addOptional<std::vector<double>>("miniIsoParams")
0071             ->setComment("Parameters used for miniIso (as in PATMuonProducer)");
0072         descriptions.add("muonsUpdated", desc);
0073       } else if (typeid(T) == typeid(pat::Electron)) {
0074         desc.addOptional<std::vector<double>>("miniIsoParamsB")
0075             ->setComment("Parameters used for miniIso in the barrel (as in PATElectronProducer)");
0076         desc.addOptional<std::vector<double>>("miniIsoParamsE")
0077             ->setComment("Parameters used for miniIso in the endcap (as in PATElectronProducer)");
0078         descriptions.add("electronsUpdated", desc);
0079       }
0080     }
0081
0082     void setDZ(T &lep, const reco::Vertex &pv) const {}
0083
0084     void readMiniIsoParams(const edm::ParameterSet &iConfig) {
0085       miniIsoParams_[0] = iConfig.getParameter<std::vector<double>>("miniIsoParams");
0086       if (miniIsoParams_[0].size() != 9)
0087         throw cms::Exception("ParameterError", "miniIsoParams must have exactly 9 elements.\n");
0088     }
0089     const std::vector<double> &miniIsoParams(const T &lep) const { return miniIsoParams_[0]; }
0090
0091     void recomputeMuonBasicSelectors(T &, const reco::Vertex &, const bool) const;
0092
0093     void recomputeSoftMuonMvaRun3(T &);
0094
0095   private:
0096     // configurables
0097     edm::EDGetTokenT<std::vector<T>> src_;
0098     edm::EDGetTokenT<std::vector<reco::Vertex>> vertices_;
0099     edm::EDGetTokenT<reco::BeamSpot> beamLineToken_;
0100     bool computeMiniIso_;
0101     bool fixDxySign_;
0102     bool recomputeMuonBasicSelectors_;
0103     bool recomputeSoftMuonMvaRun3_;
0104     std::vector<double> miniIsoParams_[2];
0105     edm::EDGetTokenT<pat::PackedCandidateCollection> pcToken_;
0106     std::unique_ptr<pat::XGBooster> softMuonMvaRun3Booster_;
0107   };
0108
0109   // must do the specialization within the namespace otherwise gcc complains
0110   //
0111   template <>
0112   void LeptonUpdater<pat::Electron>::setDZ(pat::Electron &anElectron, const reco::Vertex &pv) const {
0113     auto track = anElectron.gsfTrack();
0114     anElectron.setDB(track->dz(pv.position()), std::hypot(track->dzError(), pv.zError()), pat::Electron::PVDZ);
0115   }
0116
0117   template <>
0118   void LeptonUpdater<pat::Muon>::setDZ(pat::Muon &aMuon, const reco::Vertex &pv) const {
0119     auto track = aMuon.muonBestTrack();
0120     aMuon.setDB(track->dz(pv.position()), std::hypot(track->dzError(), pv.zError()), pat::Muon::PVDZ);
0121   }
0122
0123   template <>
0124   void LeptonUpdater<pat::Electron>::readMiniIsoParams(const edm::ParameterSet &iConfig) {
0125     miniIsoParams_[0] = iConfig.getParameter<std::vector<double>>("miniIsoParamsB");
0126     miniIsoParams_[1] = iConfig.getParameter<std::vector<double>>("miniIsoParamsE");
0127     if (miniIsoParams_[0].size() != 9)
0128       throw cms::Exception("ParameterError", "miniIsoParamsB must have exactly 9 elements.\n");
0129     if (miniIsoParams_[1].size() != 9)
0130       throw cms::Exception("ParameterError", "miniIsoParamsE must have exactly 9 elements.\n");
0131   }
0132   template <>
0133   const std::vector<double> &LeptonUpdater<pat::Electron>::miniIsoParams(const pat::Electron &lep) const {
0134     return miniIsoParams_[lep.isEE()];
0135   }
0136
0137   template <typename T>
0138   void LeptonUpdater<T>::recomputeMuonBasicSelectors(T &lep,
0139                                                      const reco::Vertex &pv,
0140                                                      const bool do_hip_mitigation_2016) const {}
0141
0142   template <>
0143   void LeptonUpdater<pat::Muon>::recomputeMuonBasicSelectors(pat::Muon &lep,
0144                                                              const reco::Vertex &pv,
0145                                                              const bool do_hip_mitigation_2016) const {
0146     lep.setSelectors(muon::makeSelectorBitset(lep, &pv, do_hip_mitigation_2016));
0147   }
0148
0149   template <typename T>
0150   void LeptonUpdater<T>::recomputeSoftMuonMvaRun3(T &lep) {}
0151
0152   template <>
0153   void LeptonUpdater<pat::Muon>::recomputeSoftMuonMvaRun3(pat::Muon &muon) {
0154     muon.setSoftMvaRun3Value(computeSoftMvaRun3(*softMuonMvaRun3Booster_, muon));
0155   }
0156
0157 }  // namespace pat
0158
0159 template <typename T>
0160 void pat::LeptonUpdater<T>::produce(edm::Event &iEvent, edm::EventSetup const &) {
0161   edm::Handle<std::vector<T>> src;
0162   iEvent.getByToken(src_, src);
0163
0164   edm::Handle<std::vector<reco::Vertex>> vertices;
0165   iEvent.getByToken(vertices_, vertices);
0166   const reco::Vertex &pv = vertices->front();
0167
0168   edm::Handle<pat::PackedCandidateCollection> pc;
0169   if (computeMiniIso_)
0170     iEvent.getByToken(pcToken_, pc);
0171
0172   edm::Handle<reco::BeamSpot> beamSpotHandle;
0173   iEvent.getByToken(beamLineToken_, beamSpotHandle);
0174   reco::BeamSpot beamSpot;
0175   bool beamSpotIsValid = false;
0176   if (beamSpotHandle.isValid()) {
0177     beamSpot = *beamSpotHandle;
0178     beamSpotIsValid = true;
0179   } else {
0180     edm::LogError("DataNotAvailable") << "No beam spot available  \n";
0181   }
0182
0183   std::unique_ptr<std::vector<T>> out(new std::vector<T>(*src));
0184
0185   const bool do_hip_mitigation_2016 =
0186       recomputeMuonBasicSelectors_ && (272728 <= iEvent.run() && iEvent.run() <= 278808);
0187
0188   for (unsigned int i = 0, n = src->size(); i < n; ++i) {
0189     T &lep = (*out)[i];
0190     setDZ(lep, pv);
0191     if (computeMiniIso_) {
0192       const auto &params = miniIsoParams(lep);
0193       pat::PFIsolation miniiso = pat::getMiniPFIsolation(pc.product(),
0194                                                          lep.polarP4(),
0195                                                          params[0],
0196                                                          params[1],
0197                                                          params[2],
0198                                                          params[3],
0199                                                          params[4],
0200                                                          params[5],
0201                                                          params[6],
0202                                                          params[7],
0203                                                          params[8]);
0204       lep.setMiniPFIsolation(miniiso);
0205     }
0206     if (recomputeMuonBasicSelectors_)
0207       recomputeMuonBasicSelectors(lep, pv, do_hip_mitigation_2016);
0208     if (recomputeSoftMuonMvaRun3_) {
0209       recomputeSoftMuonMvaRun3(lep);
0210     }
0211     //Fixing the sign of impact parameters
0212     if (fixDxySign_) {
0213       float signPV = 1.;
0214       float signBS = 1.;
0215       if (beamSpotIsValid) {
0216         if constexpr (std::is_same_v<T, pat::Electron>)
0217           signBS = copysign(1., lep.gsfTrack()->dxy(beamSpot));
0218         else
0219           signBS = copysign(1., lep.bestTrack()->dxy(beamSpot));
0220       }
0221       if constexpr (std::is_same_v<T, pat::Electron>)
0222         signPV = copysign(1., lep.gsfTrack()->dxy(pv.position()));
0223       else
0224         signPV = copysign(1., lep.bestTrack()->dxy(pv.position()));
0225       lep.setDB(abs(lep.dB(T::PV2D)) * signPV, lep.edB(T::PV2D), T::PV2D);
0226       lep.setDB(abs(lep.dB(T::BS2D)) * signBS, lep.edB(T::BS2D), T::BS2D);
0227     }
0228   }
0229
0230   iEvent.put(std::move(out));
0231 }
0232
0233 typedef pat::LeptonUpdater<pat::Electron> PATElectronUpdater;
0234 typedef pat::LeptonUpdater<pat::Muon> PATMuonUpdater;
0235
0236 #include "FWCore/Framework/interface/MakerMacros.h"
0237 DEFINE_FWK_MODULE(PATElectronUpdater);
0238 DEFINE_FWK_MODULE(PATMuonUpdater);