Adding MVA ID to nanoAOD code

PhysicsTools/NanoAOD/plugins/EvaluateMuonMVAID.cc

@@ -0,0 +1,63 @@
+//
+// 
+//
+//
+
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/StreamID.h"
+
+#include "DataFormats/PatCandidates/interface/Muon.h"
+#include "DataFormats/PatCandidates/interface/Electron.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+
+#include "DataFormats/VertexReco/interface/Vertex.h"
+#include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
+
+#include "RecoVertex/VertexTools/interface/VertexDistance3D.h"
+#include "RecoVertex/VertexPrimitives/interface/ConvertToFromReco.h"
+#include "RecoVertex/VertexPrimitives/interface/VertexState.h"
+
+#include "PhysicsTools/PatAlgos/interface/BaseMVAValueMapProducer.h"
+#include <vector>
+
+class EvaluateMuonMVAID : public BaseMVAValueMapProducer<pat::Muon> {
+public:
+  explicit EvaluateMuonMVAID(const edm::ParameterSet& iConfig, const BaseMVACache* cache)
+      : BaseMVAValueMapProducer<pat::Muon>(iConfig, cache) {}
+
+  void fillAdditionalVariables(const pat::Muon& m) override {
+    float norm_chi2;
+    float n_Valid_hits;
+    if (m.globalTrack().isNonnull()) {
+      norm_chi2 = m.globalTrack()->normalizedChi2();
+      n_Valid_hits = m.globalTrack()->hitPattern().numberOfValidMuonHits();
+    } else if (m.innerTrack().isNonnull()) {
+      norm_chi2 = m.innerTrack()->normalizedChi2();
+      n_Valid_hits = m.innerTrack()->hitPattern().numberOfValidMuonHits();
+    } else {
+      norm_chi2 = -99;
+      n_Valid_hits = -99;
+    }
+
+
+    this->setValue("Muon_norm_chi2_extended", norm_chi2);
+    this->setValue("Muon_n_Valid_hits_extended", n_Valid_hits);
+
+  }
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+    edm::ParameterSetDescription desc = BaseMVAValueMapProducer<pat::Muon>::getDescription();
+    descriptions.add("EvaluateMuonMVAID", desc);
+  }
+
+
+};
+
+//define this as a plug-in
+DEFINE_FWK_MODULE(EvaluateMuonMVAID);

PhysicsTools/NanoAOD/python/muons_cff.py

@@ -55,6 +55,32 @@
     cut = cms.string("pt > 3 && track.isNonnull && isLooseMuon")
 )
 
+muonMVAID= cms.EDProducer("EvaluateMuonMVAID",
+    src = cms.InputTag("linkedObjects","muons"),
+    weightFile =  cms.FileInPath("RecoMuon/MuonIdentification/data/mvaID.onnx"),
+    isClassifier = cms.bool(False),
+    backend = cms.string('ONNX'),
+    name = cms.string("muonMVAID"),
+    outputTensorName= cms.string("probabilities"),
+    inputTensorName= cms.string("float_input"),
+    outputNames = cms.vstring(["probBAD","probGOOD"]),
+    batch_eval =cms.bool(True),
+    outputFormulas = cms.vstring(["1.0*at(0)","at(1)"]),
+    variablesOrder = cms.vstring(["LepGood_global_muon","LepGood_validFraction","Muon_norm_chi2_extended","LepGood_local_chi2","LepGood_kink","LepGood_segmentComp","Muon_n_Valid_hits_extended","LepGood_n_MatchedStations","LepGood_Valid_pixel","LepGood_tracker_layers","LepGood_pt","LepGood_eta"]),
+    variables = cms.PSet(
+        LepGood_global_muon = cms.string("isGlobalMuon"),
+        LepGood_validFraction = cms.string("?innerTrack.isNonnull?innerTrack().validFraction:-99"),
+        LepGood_local_chi2 = cms.string("combinedQuality().chi2LocalPosition"),
+        LepGood_kink = cms.string("combinedQuality().trkKink"),
+        LepGood_segmentComp = cms.string("segmentCompatibility"),
+        LepGood_n_MatchedStations = cms.string("numberOfMatchedStations()"),
+        LepGood_Valid_pixel = cms.string("?innerTrack.isNonnull()?innerTrack().hitPattern().numberOfValidPixelHits():-99"),
+        LepGood_tracker_layers = cms.string("?innerTrack.isNonnull()?innerTrack().hitPattern().trackerLayersWithMeasurement():-99"),
+        LepGood_pt = cms.string("pt"),
+        LepGood_eta = cms.string("eta"),
+    ) 
+)
+
 muonMVATTH= cms.EDProducer("MuonBaseMVAValueMapProducer",
     src = cms.InputTag("linkedObjects","muons"),
     weightFile =  cms.FileInPath("PhysicsTools/NanoAOD/data/mu_BDTG_2017.weights.xml"),
@@ -135,6 +161,8 @@
         mvaId = Var("passed('MvaLoose')+passed('MvaMedium')+passed('MvaTight')+passed('MvaVTight')+passed('MvaVVTight')","uint8",doc="Mva ID from miniAOD selector (1=MvaLoose, 2=MvaMedium, 3=MvaTight, 4=MvaVTight, 5=MvaVVTight)"),
         mvaLowPtId = Var("passed('LowPtMvaLoose')+passed('LowPtMvaMedium')","uint8", doc="Low Pt Mva ID from miniAOD selector (1=LowPtMvaLoose, 2=LowPtMvaMedium)"),
         miniIsoId = Var("passed('MiniIsoLoose')+passed('MiniIsoMedium')+passed('MiniIsoTight')+passed('MiniIsoVeryTight')","uint8",doc="MiniIso ID from miniAOD selector (1=MiniIsoLoose, 2=MiniIsoMedium, 3=MiniIsoTight, 4=MiniIsoVeryTight)"),
+        mvaIDMuon = Var("mvaIDValue()",float,doc="MVA-based ID score (from miniAOD)",precision=14),
+        mvaIDMuon_WP = Var("passed('MvaIDwpMedium')+passed('MvaIDwpTight')","uint8",doc="MVA-based ID selector WPs (1=MVAIDwpMedium,2=MVAIDwpTight)"),
         multiIsoId = Var("?passed('MultiIsoMedium')?2:passed('MultiIsoLoose')","uint8",doc="MultiIsoId from miniAOD selector (1=MultiIsoLoose, 2=MultiIsoMedium)"),
         puppiIsoId = Var("passed('PuppiIsoLoose')+passed('PuppiIsoMedium')+passed('PuppiIsoTight')", "uint8", doc="PuppiIsoId from miniAOD selector (1=Loose, 2=Medium, 3=Tight)"),
         triggerIdLoose = Var("passed('TriggerIdLoose')",bool,doc="TriggerIdLoose ID"), 
@@ -145,6 +173,7 @@
         mvaTTH = ExtVar(cms.InputTag("muonMVATTH"),float, doc="TTH MVA lepton ID score",precision=14),
         mvaLowPt = ExtVar(cms.InputTag("muonMVALowPt"),float, doc="Low pt muon ID score",precision=14),
         fsrPhotonIdx = ExtVar(cms.InputTag("leptonFSRphotons:muFsrIndex"),int, doc="Index of the lowest-dR/ET2 among associated FSR photons"),
+        mvaIDMuon_nano = ExtVar(cms.InputTag("muonMVAID:probGOOD"),float, doc="MVA-based ID score (from nanoAOD)",precision=14)
     ),
 )
 
@@ -181,4 +210,4 @@
 
 muonTask = cms.Task(slimmedMuonsUpdated,isoForMu,ptRatioRelForMu,slimmedMuonsWithUserData,finalMuons,finalLooseMuons )
 muonMCTask = cms.Task(muonsMCMatchForTable,muonMCTable)
-muonTablesTask = cms.Task(muonMVATTH,muonMVALowPt,muonTable)
+muonTablesTask = cms.Task(muonMVATTH,muonMVALowPt,muonMVAID,muonTable)