TrackerSeedGenerator/plugins/MuonHLTSeedMVAClassifier.cc

0001
0002 // Package:    RecoMuon_TrackerSeedGenerator
0003 // Class:      MuonHLTSeedMVAClassifier
0004
0005 // Original Author:  Won Jun, OH Minseok
0006 //         Created:  Fri, 28 May 2021
0007
0008 // system include files
0009 #include <memory>
0010 #include <cmath>
0011
0012 // user include files
0013 #include "FWCore/Framework/interface/Frameworkfwd.h"
0014 #include "FWCore/Framework/interface/stream/EDProducer.h"
0015
0016 #include "FWCore/Framework/interface/Event.h"
0017 #include "FWCore/Framework/interface/MakerMacros.h"
0018
0019 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0020 #include "FWCore/Utilities/interface/StreamID.h"
0021
0022 // Geometry
0023 #include "Geometry/Records/interface/TrackerDigiGeometryRecord.h"
0024 #include "Geometry/TrackerGeometryBuilder/interface/TrackerGeometry.h"
0025
0026 // TrajectorySeed
0027 #include "DataFormats/TrajectorySeed/interface/TrajectorySeed.h"
0028 #include "DataFormats/TrajectorySeed/interface/TrajectorySeedCollection.h"
0029 #include "DataFormats/TrajectorySeed/interface/PropagationDirection.h"
0030 #include "DataFormats/TrajectoryState/interface/PTrajectoryStateOnDet.h"
0031 #include "DataFormats/TrajectoryState/interface/LocalTrajectoryParameters.h"
0032 #include "DataFormats/TrackingRecHit/interface/TrackingRecHit.h"
0033
0034 #include "RecoMuon/TrackerSeedGenerator/interface/SeedMvaEstimator.h"
0035
0036 // class declaration
0037 bool sortByMvaScore(const std::pair<unsigned, double>& A, const std::pair<unsigned, double>& B) {
0038   return (A.second > B.second);
0039 };
0040
0041 class MuonHLTSeedMVAClassifier : public edm::stream::EDProducer<> {
0042 public:
0043   explicit MuonHLTSeedMVAClassifier(const edm::ParameterSet&);
0044   ~MuonHLTSeedMVAClassifier() override = default;
0045
0046   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0047
0048 private:
0049   void produce(edm::Event&, const edm::EventSetup&) override;
0050
0051   // member data
0052   const edm::EDGetTokenT<TrajectorySeedCollection> seedToken_;
0053   const edm::EDGetTokenT<l1t::MuonBxCollection> l1MuonToken_;
0054   const edm::EDGetTokenT<reco::RecoChargedCandidateCollection> l2MuonToken_;
0055   const edm::ESGetToken<TrackerGeometry, TrackerDigiGeometryRecord> trackerGeometryToken_;
0056
0057   typedef std::pair<std::unique_ptr<const SeedMvaEstimator>, std::unique_ptr<const SeedMvaEstimator>>
0058       PairSeedMvaEstimator;
0059   PairSeedMvaEstimator mvaEstimator_;
0060
0061   const bool rejectAll_;
0062   const bool isFromL1_;
0063
0064   const edm::FileInPath mvaFileB_;
0065   const edm::FileInPath mvaFileE_;
0066
0067   const std::vector<double> mvaScaleMeanB_;
0068   const std::vector<double> mvaScaleStdB_;
0069   const std::vector<double> mvaScaleMeanE_;
0070   const std::vector<double> mvaScaleStdE_;
0071
0072   const bool doSort_;
0073   const int nSeedsMaxB_;
0074   const int nSeedsMaxE_;
0075
0076   const double etaEdge_;
0077   const double mvaCutB_;
0078   const double mvaCutE_;
0079
0080   const int minL1Qual_;
0081   const double baseScore_;
0082
0083   double getSeedMva(const PairSeedMvaEstimator& pairMvaEstimator,
0084                     const TrajectorySeed& seed,
0085                     const GlobalVector& global_p,
0086                     const l1t::MuonBxCollection& l1Muons,
0087                     const reco::RecoChargedCandidateCollection& l2Muons);
0088 };
0089
0090 MuonHLTSeedMVAClassifier::MuonHLTSeedMVAClassifier(const edm::ParameterSet& iConfig)
0091     : seedToken_(consumes<TrajectorySeedCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0092       l1MuonToken_(consumes<l1t::MuonBxCollection>(iConfig.getParameter<edm::InputTag>("L1Muon"))),
0093       l2MuonToken_(consumes<reco::RecoChargedCandidateCollection>(iConfig.getParameter<edm::InputTag>("L2Muon"))),
0094       trackerGeometryToken_(esConsumes<TrackerGeometry, TrackerDigiGeometryRecord>()),
0095
0096       rejectAll_(iConfig.getParameter<bool>("rejectAll")),
0097       isFromL1_(iConfig.getParameter<bool>("isFromL1")),
0098
0099       mvaFileB_(iConfig.getParameter<edm::FileInPath>(isFromL1_ ? "mvaFileBL1" : "mvaFileBL2")),
0100       mvaFileE_(iConfig.getParameter<edm::FileInPath>(isFromL1_ ? "mvaFileEL1" : "mvaFileEL2")),
0101
0102       mvaScaleMeanB_(iConfig.getParameter<std::vector<double>>(isFromL1_ ? "mvaScaleMeanBL1" : "mvaScaleMeanBL2")),
0103       mvaScaleStdB_(iConfig.getParameter<std::vector<double>>(isFromL1_ ? "mvaScaleStdBL1" : "mvaScaleStdBL2")),
0104       mvaScaleMeanE_(iConfig.getParameter<std::vector<double>>(isFromL1_ ? "mvaScaleMeanEL1" : "mvaScaleMeanEL2")),
0105       mvaScaleStdE_(iConfig.getParameter<std::vector<double>>(isFromL1_ ? "mvaScaleStdEL1" : "mvaScaleStdEL2")),
0106
0107       doSort_(iConfig.getParameter<bool>("doSort")),
0108       nSeedsMaxB_(iConfig.getParameter<int>("nSeedsMaxB")),
0109       nSeedsMaxE_(iConfig.getParameter<int>("nSeedsMaxE")),
0110
0111       etaEdge_(iConfig.getParameter<double>("etaEdge")),
0112       mvaCutB_(iConfig.getParameter<double>("mvaCutB")),
0113       mvaCutE_(iConfig.getParameter<double>("mvaCutE")),
0114
0115       minL1Qual_(iConfig.getParameter<int>("minL1Qual")),
0116       baseScore_(iConfig.getParameter<double>("baseScore")) {
0117   if (!rejectAll_) {
0118     mvaEstimator_ = std::make_pair(
0119         std::make_unique<SeedMvaEstimator>(mvaFileB_, mvaScaleMeanB_, mvaScaleStdB_, isFromL1_, minL1Qual_),
0120         std::make_unique<SeedMvaEstimator>(mvaFileE_, mvaScaleMeanE_, mvaScaleStdE_, isFromL1_, minL1Qual_));
0121   }
0122
0123   produces<TrajectorySeedCollection>();
0124 }
0125
0126 // -- method called on each new Event
0127 void MuonHLTSeedMVAClassifier::produce(edm::Event& iEvent, edm::EventSetup const& iEventSetup) {
0128   auto result = std::make_unique<TrajectorySeedCollection>();
0129
0130   if (rejectAll_) {
0131     iEvent.put(std::move(result));
0132     return;
0133   }
0134
0135   if (doSort_ && nSeedsMaxB_ <= 0 && nSeedsMaxE_ <= 0) {
0136     iEvent.put(std::move(result));
0137     return;
0138   }
0139
0140   if (!doSort_ && mvaCutB_ > 1. && mvaCutE_ > 1.) {
0141     iEvent.put(std::move(result));
0142     return;
0143   }
0144
0145   const TrajectorySeedCollection& seeds = iEvent.get(seedToken_);
0146   const l1t::MuonBxCollection& l1Muons = iEvent.get(l1MuonToken_);
0147   const reco::RecoChargedCandidateCollection& l2Muons = iEvent.get(l2MuonToken_);
0148   const TrackerGeometry& trkGeom = iEventSetup.getData(trackerGeometryToken_);
0149
0150   std::vector<std::pair<unsigned, double>> pairSeedIdxMvaScoreB = {};
0151   std::vector<std::pair<unsigned, double>> pairSeedIdxMvaScoreE = {};
0152   for (auto& seed : seeds) {
0153     const GlobalVector global_p =
0154         trkGeom.idToDet(seed.startingState().detId())->surface().toGlobal(seed.startingState().parameters().momentum());
0155
0156     bool isB = (std::abs(global_p.eta()) < etaEdge_);
0157
0158     if (doSort_) {
0159       if (isB) {
0160         if (nSeedsMaxB_ <= 0) {
0161           continue;
0162         }
0163       } else {
0164         if (nSeedsMaxE_ <= 0) {
0165           continue;
0166         }
0167       }
0168     } else {
0169       if (isB) {
0170         if (mvaCutB_ > 1.0) {
0171           continue;
0172         } else if (mvaCutB_ <= 0.) {
0173           result->emplace_back(seed);
0174           continue;
0175         }
0176       } else {
0177         if (mvaCutE_ > 1.0) {
0178           continue;
0179         } else if (mvaCutE_ <= 0.) {
0180           result->emplace_back(seed);
0181           continue;
0182         }
0183       }
0184     }
0185
0186     double mva = getSeedMva(mvaEstimator_, seed, global_p, l1Muons, l2Muons);
0187
0188     double score = 1. / (1. + std::exp(-1. * mva));
0189     bool passMva = isB ? score > mvaCutB_ : score > mvaCutE_;
0190     if (!passMva)
0191       continue;
0192
0193     if (doSort_) {
0194       if (isB)
0195         pairSeedIdxMvaScoreB.push_back(std::make_pair(&seed - &seeds.at(0), score));
0196       else
0197         pairSeedIdxMvaScoreE.push_back(std::make_pair(&seed - &seeds.at(0), score));
0198     } else {
0199       result->emplace_back(seed);
0200     }
0201   }
0202
0203   if (doSort_) {
0204     std::sort(pairSeedIdxMvaScoreB.begin(), pairSeedIdxMvaScoreB.end(), sortByMvaScore);
0205     std::sort(pairSeedIdxMvaScoreE.begin(), pairSeedIdxMvaScoreE.end(), sortByMvaScore);
0206
0207     for (auto i = 0U; i < pairSeedIdxMvaScoreB.size(); ++i) {
0208       if ((int)i == nSeedsMaxB_)
0209         break;
0210       const auto& seed(seeds.at(pairSeedIdxMvaScoreB.at(i).first));
0211       result->emplace_back(seed);
0212     }
0213
0214     for (auto i = 0U; i < pairSeedIdxMvaScoreE.size(); ++i) {
0215       if ((int)i == nSeedsMaxE_)
0216         break;
0217       const auto& seed(seeds.at(pairSeedIdxMvaScoreE.at(i).first));
0218       result->emplace_back(seed);
0219     }
0220   }
0221
0222   iEvent.put(std::move(result));
0223 }
0224
0225 double MuonHLTSeedMVAClassifier::getSeedMva(const PairSeedMvaEstimator& pairMvaEstimator,
0226                                             const TrajectorySeed& seed,
0227                                             const GlobalVector& global_p,
0228                                             const l1t::MuonBxCollection& l1Muons,
0229                                             const reco::RecoChargedCandidateCollection& l2Muons) {
0230   double mva = 0.;
0231   if (std::abs(global_p.eta()) < etaEdge_) {
0232     mva = pairMvaEstimator.first->computeMva(seed, global_p, l1Muons, l2Muons);
0233   } else {
0234     mva = pairMvaEstimator.second->computeMva(seed, global_p, l1Muons, l2Muons);
0235   }
0236
0237   return (mva + baseScore_);
0238 }
0239
0240 // -- method fills 'descriptions' with the allowed parameters for the module  ------------
0241 void MuonHLTSeedMVAClassifier::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0242   edm::ParameterSetDescription desc;
0243   desc.add<edm::InputTag>("src", edm::InputTag("hltIter2IterL3MuonPixelSeeds", ""));
0244   desc.add<edm::InputTag>("L1Muon", edm::InputTag("hltGtStage2Digis", "Muon"));
0245   desc.add<edm::InputTag>("L2Muon", edm::InputTag("hltL2MuonCandidates", ""));
0246
0247   desc.add<bool>("rejectAll", false);
0248   desc.add<bool>("isFromL1", false);
0249
0250   desc.add<edm::FileInPath>("mvaFileBL1",
0251                             edm::FileInPath("RecoMuon/TrackerSeedGenerator/data/xgb_Run3_Iter2FromL1Seeds_barrel.xml"));
0252   desc.add<edm::FileInPath>("mvaFileEL1",
0253                             edm::FileInPath("RecoMuon/TrackerSeedGenerator/data/xgb_Run3_Iter2FromL1Seeds_endcap.xml"));
0254   desc.add<edm::FileInPath>("mvaFileBL2",
0255                             edm::FileInPath("RecoMuon/TrackerSeedGenerator/data/xgb_Run3_Iter2Seeds_barrel.xml"));
0256   desc.add<edm::FileInPath>("mvaFileEL2",
0257                             edm::FileInPath("RecoMuon/TrackerSeedGenerator/data/xgb_Run3_Iter2Seeds_endcap.xml"));
0258   desc.add<std::vector<double>>("mvaScaleMeanBL1", {0., 0., 0., 0., 0., 0., 0., 0.});
0259   desc.add<std::vector<double>>("mvaScaleStdBL1", {1., 1., 1., 1., 1., 1., 1., 1.});
0260   desc.add<std::vector<double>>("mvaScaleMeanEL1", {0., 0., 0., 0., 0., 0., 0., 0.});
0261   desc.add<std::vector<double>>("mvaScaleStdEL1", {1., 1., 1., 1., 1., 1., 1., 1.});
0262   desc.add<std::vector<double>>("mvaScaleMeanBL2", {0., 0., 0., 0., 0., 0., 0., 0., 0., 0.});
0263   desc.add<std::vector<double>>("mvaScaleStdBL2", {1., 1., 1., 1., 1., 1., 1., 1., 1., 1.});
0264   desc.add<std::vector<double>>("mvaScaleMeanEL2", {0., 0., 0., 0., 0., 0., 0., 0., 0., 0.});
0265   desc.add<std::vector<double>>("mvaScaleStdEL2", {1., 1., 1., 1., 1., 1., 1., 1., 1., 1.});
0266
0267   desc.add<bool>("doSort", false);
0268   desc.add<int>("nSeedsMaxB", 1e6);
0269   desc.add<int>("nSeedsMaxE", 1e6);
0270
0271   desc.add<double>("etaEdge", 1.2);
0272   desc.add<double>("mvaCutB", -1.);
0273   desc.add<double>("mvaCutE", -1.);
0274
0275   desc.add<int>("minL1Qual", 7);
0276   desc.add<double>("baseScore", 0.5);
0277
0278   descriptions.add("MuonHLTSeedMVAClassifier", desc);
0279 }
0280
0281 //define this as a plug-in
0282 DEFINE_FWK_MODULE(MuonHLTSeedMVAClassifier);