ElectronIdentification/plugins/ElectronIdMVABased.cc

0001 //
0002 //
0003 // Original Author:  Zablocki Jakub
0004 //         Created:  Thu Feb  9 10:47:50 CST 2012
0005 //
0006 //
0007
0008 // system include files
0009 #include <memory>
0010
0011 // user include files
0012 #include "FWCore/Framework/interface/Frameworkfwd.h"
0013 #include "FWCore/Framework/interface/global/EDProducer.h"
0014
0015 #include "FWCore/Framework/interface/Event.h"
0016 #include "FWCore/Framework/interface/MakerMacros.h"
0017
0018 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0019 #include "RecoEgamma/ElectronIdentification/interface/ElectronMVAEstimator.h"
0020 #include "DataFormats/EgammaCandidates/interface/GsfElectron.h"
0021 #include "DataFormats/EgammaCandidates/interface/GsfElectronFwd.h"
0022 #include "DataFormats/VertexReco/interface/Vertex.h"
0023 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0024 //
0025 // class declaration
0026 //
0027
0028 class ElectronIdMVABased : public edm::global::EDProducer<> {
0029 public:
0030   explicit ElectronIdMVABased(const edm::ParameterSet&);
0031   ~ElectronIdMVABased() override {}
0032
0033 private:
0034   void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0035
0036   // ----------member data ---------------------------
0037   const edm::EDGetTokenT<reco::VertexCollection> vertexToken;
0038   const edm::EDGetTokenT<reco::GsfElectronCollection> electronToken;
0039   const std::vector<std::string> mvaWeightFileEleID;
0040   const std::string path_mvaWeightFileEleID;
0041   const double thresholdBarrel;
0042   const double thresholdEndcap;
0043   const double thresholdIsoBarrel;
0044   const double thresholdIsoEndcap;
0045
0046   const std::unique_ptr<const ElectronMVAEstimator> mvaID_;
0047 };
0048
0049 // constructor
0050 ElectronIdMVABased::ElectronIdMVABased(const edm::ParameterSet& iConfig)
0051     : vertexToken(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("vertexTag"))),
0052       electronToken(consumes<reco::GsfElectronCollection>(iConfig.getParameter<edm::InputTag>("electronTag"))),
0053       thresholdBarrel(iConfig.getParameter<double>("thresholdBarrel")),
0054       thresholdEndcap(iConfig.getParameter<double>("thresholdEndcap")),
0055       thresholdIsoBarrel(iConfig.getParameter<double>("thresholdIsoDR03Barrel")),
0056       thresholdIsoEndcap(iConfig.getParameter<double>("thresholdIsoDR03Endcap")),
0057       mvaID_(new ElectronMVAEstimator(ElectronMVAEstimator::Configuration{
0058           .vweightsfiles = iConfig.getParameter<std::vector<std::string> >("HZZmvaWeightFile")})) {
0059   produces<reco::GsfElectronCollection>();
0060 }
0061
0062 // ------------ method called on each new Event  ------------
0063 void ElectronIdMVABased::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0064   constexpr double etaEBEE = 1.485;
0065
0066   auto mvaElectrons = std::make_unique<reco::GsfElectronCollection>();
0067
0068   edm::Handle<reco::VertexCollection> vertexCollection;
0069   iEvent.getByToken(vertexToken, vertexCollection);
0070   int nVtx = vertexCollection->size();
0071
0072   edm::Handle<reco::GsfElectronCollection> egCollection;
0073   iEvent.getByToken(electronToken, egCollection);
0074   const reco::GsfElectronCollection egCandidates = (*egCollection.product());
0075   for (reco::GsfElectronCollection::const_iterator egIter = egCandidates.begin(); egIter != egCandidates.end();
0076        ++egIter) {
0077     double mvaVal = mvaID_->mva(*egIter, nVtx);
0078     double isoDr03 = egIter->dr03TkSumPt() + egIter->dr03EcalRecHitSumEt() + egIter->dr03HcalTowerSumEt();
0079     double eleEta = fabs(egIter->eta());
0080     if (eleEta <= etaEBEE && mvaVal > thresholdBarrel && isoDr03 < thresholdIsoBarrel) {
0081       mvaElectrons->push_back(*egIter);
0082       reco::GsfElectron::MvaOutput myMvaOutput;
0083       myMvaOutput.mva_Isolated = mvaVal;
0084       mvaElectrons->back().setMvaOutput(myMvaOutput);
0085     } else if (eleEta > etaEBEE && mvaVal > thresholdEndcap && isoDr03 < thresholdIsoEndcap) {
0086       mvaElectrons->push_back(*egIter);
0087       reco::GsfElectron::MvaOutput myMvaOutput;
0088       myMvaOutput.mva_Isolated = mvaVal;
0089       mvaElectrons->back().setMvaOutput(myMvaOutput);
0090     }
0091   }
0092
0093   iEvent.put(std::move(mvaElectrons));
0094 }
0095
0096 //define this as a plug-in
0097 DEFINE_FWK_MODULE(ElectronIdMVABased);