[NanoAOD] Add ParticleNet-based Muon ID

DataFormats/BTauReco/interface/DeepBoostedJetFeatures.h

@@ -52,6 +52,8 @@ namespace btagbtvdeep {
       }
     }
 
+    const std::unordered_map<std::string, std::vector<float>>& get_all() const { return feature_map_; }
+
   private:
     bool is_empty_ = true;
     std::unordered_map<std::string, std::vector<float>> feature_map_;

PhysicsTools/NanoAOD/python/muons_cff.py

@@ -145,6 +145,66 @@
     variables = _legacy_muon_BDT_variable
 )
 
+pnetMuonVariables = cms.EDProducer("MuonInfoCollectionProducer",
+                                   src = cms.InputTag("linkedObjects","muons"),
+                                   secondary_vertices=cms.InputTag("slimmedSecondaryVertices"),
+                                   pvSrc = cms.InputTag("offlineSlimmedPrimaryVertices"),
+                                   pfCandidates=cms.InputTag("packedPFCandidates"),
+                                   leptonVars = cms.PSet(
+                                       MuonSelected_LepGood_pt = cms.string("pt"),
+                                       MuonSelected_LepGood_eta = cms.string("eta"),
+                                       MuonSelected_LepGood_jetNDauChargedMVASel = cms.string("?userCand('jetForLepJetVar').isNonnull()?userFloat('jetNDauChargedMVASel'):0"),
+                                       MuonSelected_LepGood_miniRelIsoCharged = cms.string("userFloat('miniIsoChg')/pt"),
+                                       MuonSelected_LepGood_miniRelIsoNeutral = cms.string("(userFloat('miniIsoAll')-userFloat('miniIsoChg'))/pt"),
+                                       MuonSelected_LepGood_jetPtRelv2 = cms.string("?userCand('jetForLepJetVar').isNonnull()?userFloat('ptRel'):0"),
+                                       MuonSelected_LepGood_jetDF = cms.string("?userCand('jetForLepJetVar').isNonnull()?max(userCand('jetForLepJetVar').bDiscriminator('pfDeepFlavourJetTags:probbb')+userCand('jetForLepJetVar').bDiscriminator('pfDeepFlavourJetTags:probb')+userCand('jetForLepJetVar').bDiscriminator('pfDeepFlavourJetTags:problepb'),0.0):0.0"),
+                                       MuonSelected_LepGood_jetPtRatio = cms.string("?userCand('jetForLepJetVar').isNonnull()?min(userFloat('ptRatio'),1.5):1.0/(1.0+(pfIsolationR04().sumChargedHadronPt + max(pfIsolationR04().sumNeutralHadronEt + pfIsolationR04().sumPhotonEt - pfIsolationR04().sumPUPt/2,0.0))/pt)"),
+                                       MuonSelected_dxy = cms.string("log(abs(dB('PV2D')))"),
+                                       MuonSelected_sip3d = cms.string("abs(dB('PV3D')/edB('PV3D'))"),
+                                       MuonSelected_dz = cms.string("log(abs(dB('PVDZ')))"),
+                                       MuonSelected_LepGood_dz = cms.string("log(abs(dB('PVDZ')))"),
+                                       MuonSelected_segmentComp = cms.string("segmentCompatibility"),
+                                       MuonSelected_global_muon = cms.string("isGlobalMuon"),
+                                       MuonSelected_validFraction = cms.string("?innerTrack.isNonnull?innerTrack().validFraction:-99"),
+                                       MuonSelected_local_chi2 = cms.string("combinedQuality().chi2LocalPosition"),
+                                       MuonSelected_kink = cms.string("combinedQuality().trkKink"),
+                                       MuonSelected_n_MatchedStations = cms.string("numberOfMatchedStations()"),
+                                       MuonSelected_Valid_pixel = cms.string("?innerTrack.isNonnull()?innerTrack().hitPattern().numberOfValidPixelHits():-99"),
+                                       MuonSelected_tracker_layers = cms.string("?innerTrack.isNonnull()?innerTrack().hitPattern().trackerLayersWithMeasurement():-99"),
+                                       MuonSelected_mvaId=cms.string("userFloat('mvaIDMuon')"),
+                                   ),
+                                   leptonVarsExt = cms.PSet(
+                                       MuonSelected_mvaTTH=cms.InputTag("muonMVATTH"),
+                                   ),
+                                   pfVars = cms.PSet(
+                                       PF_pt=cms.string("pt"),
+                                       PF_charge=cms.string("charge"),
+                                       PF_isElectron=cms.string("?abs(pdgId)==11?1:0"),
+                                       PF_isMuon=cms.string("?abs(pdgId)==13?1:0"),
+                                       PF_isNeutralHadron=cms.string("?abs(pdgId)==130?1:0"),
+                                       PF_isPhoton=cms.string("?abs(pdgId)==22?1:0"),
+                                       PF_isChargedHadron=cms.string("?abs(pdgId)==211?1:0"),
+                                       PF_puppiWeightNoLep=cms.string("puppiWeightNoLep"),
+                                       PF_fromPV=cms.string("fromPV"),
+                                       PF_numberOfPixelHits=cms.string("numberOfPixelHits"),
+                                       PF_dzSig=cms.string("?hasTrackDetails?dz/max(dzError,1.e-6):0"),
+                                       PF_dxySig=cms.string("?hasTrackDetails?dxy/max(dxyError,1.e-6):0"),
+                                       PF_hcalFraction=cms.string("hcalFraction"),
+                                       PF_trackerLayersWithMeasurement=cms.string("?hasTrackDetails?bestTrack().hitPattern().trackerLayersWithMeasurement:0"),
+                                       pf_mask=cms.string("1"),
+                                   ),
+                                   svVars = cms.PSet(
+                                       SV_eta=cms.string("eta"),
+                                       SV_phi=cms.string("phi"),
+                                       SV_pt=cms.string("pt"),
+                                       SV_ndof=cms.string("vertexNdof"),
+                                       SV_chi2=cms.string("vertexChi2"),
+                                       SV_nTracks=cms.string("numberOfDaughters"),
+                                       SV_mass=cms.string("mass"),
+                                       SV_mask=cms.string("1"),
+                                   ),
+)
+
 from TrackingTools.TransientTrack.TransientTrackBuilder_cfi import *
 muonBSConstrain = cms.EDProducer("MuonBeamspotConstraintValueMapProducer",
     src = cms.InputTag("linkedObjects","muons"),
@@ -208,13 +268,27 @@
     externalVariables = cms.PSet(
         promptMVA = ExtVar(cms.InputTag("muonPROMPTMVA"),float, doc="Prompt MVA lepton ID score. Corresponds to the previous mvaTTH",precision=14),
         mvaLowPt = ExtVar(cms.InputTag("muonMVALowPt"),float, doc="Low pt muon ID score",precision=14),
+        pnScore_light = ExtVar(cms.InputTag("muonPNnorwwithtaus:pnScorelight"),float, doc="Score of PNet Muon id",precision=14),
+        pnScore_prompt = ExtVar(cms.InputTag("muonPNnorwwithtaus:pnScoreprompt"),float, doc="Score of PNet Muon id",precision=14),
+        pnScore_heavy = ExtVar(cms.InputTag("muonPNnorwwithtaus:pnScoreheavy"),float, doc="Score of PNet Muon id",precision=14),
+        pnScore_tau   = ExtVar(cms.InputTag("muonPNnorwwithtaus:pnScoretau"),float, doc="Score of PNet Muon id",precision=14),
         fsrPhotonIdx = ExtVar(cms.InputTag("leptonFSRphotons:muFsrIndex"), "int16", doc="Index of the lowest-dR/ET2 among associated FSR photons"),
         bsConstrainedPt = ExtVar(cms.InputTag("muonBSConstrain:muonBSConstrainedPt"),float, doc="pT with beamspot constraint",precision=-1),
         bsConstrainedPtErr = ExtVar(cms.InputTag("muonBSConstrain:muonBSConstrainedPtErr"),float, doc="pT error with beamspot constraint ",precision=6),
         bsConstrainedChi2 = ExtVar(cms.InputTag("muonBSConstrain:muonBSConstrainedChi2"),float, doc="chi2 of beamspot constraint",precision=6),
     ),
 )
 
+muonPNnorwwithtaus = cms.EDProducer('PNETMuonProducer',
+                        src = cms.InputTag("pnetMuonVariables"),
+                        srcLeps = cms.InputTag("linkedObjects","muons"),
+                        model_path=cms.FileInPath('PhysicsTools/NanoAOD/data/PNetMuonId/model.onnx'),
+                        preprocess_json=cms.FileInPath('PhysicsTools/NanoAOD/data/PNetMuonId/preprocess.json'),
+                        name=cms.string("PN_lepton"),
+                        debugMode=cms.untracked.bool(True),
+                        flav_names = cms.vstring(["light","prompt","tau","heavy"]),
+)
+
 # Increase precision of eta and phi
 muonTable.variables.eta.precision = 16
 muonTable.variables.phi.precision = 16
@@ -249,5 +323,5 @@
 
 muonTask = cms.Task(slimmedMuonsUpdated,isoForMu,ptRatioRelForMu,slimmedMuonsWithUserData,finalMuons,finalLooseMuons )
 muonMCTask = cms.Task(muonsMCMatchForTable,muonMCTable)
-muonTablesTask = cms.Task(muonPROMPTMVA,muonMVALowPt,muonBSConstrain,muonTable,muonMVAID)
+muonTablesTask = cms.Task(muonPROMPTMVA,muonMVALowPt,muonBSConstrain,muonTable,muonMVAID,pnetMuonVariables,muonPNnorwwithtaus)
 

PhysicsTools/PatAlgos/BuildFile.xml

@@ -14,10 +14,13 @@
 <use name="FWCore/Utilities"/>
 <use name="PhysicsTools/IsolationAlgos"/>
 <use name="PhysicsTools/PatUtils"/>
+<use name="PhysicsTools/NanoAOD"/>
 <use name="PhysicsTools/TensorFlow"/>
 <use name="PhysicsTools/ONNXRuntime"/>
 <use name="PhysicsTools/XGBoost"/>
 <use name="TrackingTools/Records"/>
+<use name="RecoBTag/FeatureTools"/>
+<use name="DataFormats/BTauReco"/>
 <use name="clhep"/>
 <use name="xgboost"/>
 <export>

PhysicsTools/PatAlgos/interface/LeptonTagInfoCollectionProducer.h

@@ -0,0 +1,112 @@
+// -*- C++ -*-
+//
+// Package:    PhysicsTools/PatAlgos
+// Class:      LeptonTagInfoCollectionProducer
+//
+/**\class LeptonTagInfoCollectionProducer LeptonTagInfoCollectionProducer.cc PhysicsTools/PatAlgos/plugins/PNETLeptonProducer.cc
+
+
+*/
+//
+// Original Author:  Sergio Sanchez Cruz
+//         Created:  Mon, 15 May 2023 08:32:03 GMT
+//
+//
+
+#ifndef PhysicsTools_PatAlgos_LeptonTagInfoCollectionProducer_h
+#define PhysicsTools_PatAlgos_LeptonTagInfoCollectionProducer_h
+
+#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 "FWCore/Utilities/interface/ESGetToken.h"
+
+#include "DataFormats/Common/interface/CMS_CLASS_VERSION.h"
+#include "DataFormats/BTauReco/interface/DeepBoostedJetFeatures.h"
+#include "PhysicsTools/NanoAOD/interface/SimpleFlatTableProducer.h"
+#include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "DataFormats/BTauReco/interface/DeepBoostedJetTagInfo.h"  // this is flexible enough for our purposes
+#include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
+#include "DataFormats/VertexReco/interface/Vertex.h"
+
+namespace pat {
+  template <class Features, class T>
+
+  class FeaturesLepInfo {
+  public:
+    FeaturesLepInfo() {}
+
+    FeaturesLepInfo(const Features& features, const edm::RefToBase<T>& lep_ref)
+        : features_(features), lep_ref_(lep_ref) {}
+
+    edm::RefToBase<T> lep() const { return lep_ref_; }
+
+    const Features& features() const { return features_; }
+
+    ~FeaturesLepInfo() {}
+    FeaturesLepInfo* clone(void) const { return new FeaturesLepInfo(*this); }
+
+    CMS_CLASS_VERSION(3);
+
+  private:
+    Features features_;
+    edm::RefToBase<T> lep_ref_;
+  };
+
+  template <typename T>
+  using LeptonTagInfo = FeaturesLepInfo<btagbtvdeep::DeepBoostedJetFeatures, T>;
+
+  template <typename T>
+  using LeptonTagInfoCollection = std::vector<LeptonTagInfo<T>>;
+
+  template <typename T2>
+  using varWithName = std::pair<std::string, StringObjectFunction<T2, true>>;
+
+  template <typename T2>
+  using extVarWithName = std::pair<std::string, edm::EDGetTokenT<edm::ValueMap<T2>>>;
+
+  template <typename T>
+  class LeptonTagInfoCollectionProducer : public edm::stream::EDProducer<> {
+  public:
+    explicit LeptonTagInfoCollectionProducer(const edm::ParameterSet& iConfig);
+    ~LeptonTagInfoCollectionProducer() override {};
+
+  private:
+    void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
+    template <typename T2>
+    void parse_vars_into(const edm::ParameterSet&, std::vector<std::unique_ptr<varWithName<T2>>>&);
+    template <typename T2>
+    void parse_extvars_into(const edm::ParameterSet&, std::vector<std::unique_ptr<extVarWithName<T2>>>&);
+    void fill_lepton_features(const T&, btagbtvdeep::DeepBoostedJetFeatures&);
+    void fill_lepton_extfeatures(const edm::RefToBase<T>&, btagbtvdeep::DeepBoostedJetFeatures&, edm::Event&);
+    void fill_pf_features(const T&, btagbtvdeep::DeepBoostedJetFeatures&);
+    void fill_sv_features(const T&, btagbtvdeep::DeepBoostedJetFeatures&);
+
+    edm::EDGetTokenT<edm::View<T>> src_token_;
+    edm::EDGetTokenT<pat::PackedCandidateCollection> pf_token_;
+    edm::EDGetTokenT<reco::VertexCompositePtrCandidateCollection> sv_token_;
+    edm::EDGetTokenT<std::vector<reco::Vertex>> pv_token_;
+
+    edm::ParameterSet lepton_varsPSet_;
+    edm::ParameterSet lepton_varsExtPSet_;
+    edm::ParameterSet pf_varsPSet_;
+    edm::ParameterSet sv_varsPSet_;
+
+    std::vector<std::unique_ptr<varWithName<T>>> lepton_vars_;
+    std::vector<std::unique_ptr<varWithName<pat::PackedCandidate>>> pf_vars_;
+    std::vector<std::unique_ptr<varWithName<reco::VertexCompositePtrCandidate>>> sv_vars_;
+    edm::Handle<reco::VertexCompositePtrCandidateCollection> svs_;
+    edm::Handle<pat::PackedCandidateCollection> pfs_;
+    edm::Handle<std::vector<reco::Vertex>> pvs_;
+    std::vector<std::unique_ptr<extVarWithName<float>>> extLepton_vars_;
+  };
+
+}  // namespace pat
+
+#endif

PhysicsTools/PatAlgos/interface/PNETLeptonProducer.h

@@ -0,0 +1,89 @@
+#ifndef PhysicsTools_PatAlgos_PNETLeptonProducer
+#define PhysicsTools_PatAlgos_PNETLeptonProducer
+
+// -*- C++ -*-
+//
+// Package:    PhysicsTools/PatAlgos
+// Class:      PNETLeptonProducer
+//
+/**\class PNETLeptonProducer PNETLeptonProducer.cc PhysicsTools/PatAlgos/plugins/PNETLeptonProducer.cc
+
+
+*/
+//
+// Original Author:  Sergio Sanchez Cruz
+//         Created:  Mon, 15 May 2023 08:32:03 GMT
+//
+//
+
+#include <memory>
+
+#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 "CommonTools/Utils/interface/StringObjectFunction.h"
+#include "DataFormats/Common/interface/ValueMap.h"
+
+#include "DataFormats/PatCandidates/interface/Jet.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+#include "DataFormats/PatCandidates/interface/Electron.h"
+#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+#include "RecoBTag/FeatureTools/interface/deep_helpers.h"
+#include "DataFormats/BTauReco/interface/DeepBoostedJetTagInfo.h"  // this is flexible enough for our purposes
+//#include "DataFormats/BTauReco/interface/DeepBoostedJetFeatures.h" // this will need to be moved to the dedicated producer
+#include "PhysicsTools/PatAlgos/interface/LeptonTagInfoCollectionProducer.h"
+
+#include <string>
+//
+// class declaration
+//
+
+template <typename T>
+class PNETLeptonProducer : public edm::stream::EDProducer<edm::GlobalCache<cms::Ort::ONNXRuntime>> {
+public:
+  PNETLeptonProducer(const edm::ParameterSet&, const cms::Ort::ONNXRuntime*);
+  ~PNETLeptonProducer() override {}
+
+  /* void setValue(const std::string var, float val) { */
+  /*   if (positions_.find(var) != positions_.end()) */
+  /*     values_[positions_[var]] = val; */
+  /* } */
+
+  static std::unique_ptr<cms::Ort::ONNXRuntime> initializeGlobalCache(const edm::ParameterSet& cfg);
+  static void globalEndJob(const cms::Ort::ONNXRuntime* cache);
+
+  static edm::ParameterSetDescription getDescription();
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  void beginStream(edm::StreamID) override{};
+  void produce(edm::Event&, const edm::EventSetup&) override;
+  void endStream() override {};
+
+  ///to be implemented in derived classes, filling values for additional variables
+  virtual void readAdditionalCollections(edm::Event&, const edm::EventSetup&) {}
+  virtual void fillAdditionalVariables(const T&) {}
+
+  edm::EDGetTokenT<pat::LeptonTagInfoCollection<T>> src_;
+  edm::EDGetTokenT<std::vector<T>> leps_;
+  std::vector<std::string> flav_names_;
+  std::string name_;
+  std::vector<std::string> input_names_;            // names of each input group - the ordering is important!
+  std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
+  std::vector<unsigned> input_sizes_;               // total length of each input vector
+  std::unordered_map<std::string, btagbtvdeep::PreprocessParams>
+      prep_info_map_;  // preprocessing info for each input group
+
+  cms::Ort::FloatArrays data_;
+  bool debug_ = false;
+
+  void make_inputs(const pat::LeptonTagInfo<T>&);
+};
+
+#endif