Merge pull request #42039 from rbhattacharya04/Mu_POG_NanoAOD

Custom_Muon_POG_NanoAOD
cms-sw · Oct 30, 2023 · f932104 · f932104
2 parents 84789f5 + 3cca89c
commit f932104
Show file tree

Hide file tree

Showing 4 changed files with 425 additions and 0 deletions.
diff --git a/Configuration/PyReleaseValidation/python/relval_steps.py b/Configuration/PyReleaseValidation/python/relval_steps.py
@@ -4033,6 +4033,7 @@ def gen2023HiMix(fragment,howMuch):
 steps['NANOUP17'] = merge([{'--conditions':'auto:phase1_2017_realistic','--era': 'Run2_2017','-n':'10' ,'--filein':'file:step3_inMINIAODSIM.root', '--geometry':'DB:Extended', '--nThreads':'2'}, stepNanoEDMMCProd])
 steps['NANOUP15Had']=merge([{'--filein':'file:step4_inMINIAODSIM.root'},steps['NANOUP15']])
 steps['NANOUP15MC_PU25_JME']=merge([{'--customise':'PhysicsTools/NanoAOD/custom_jme_cff.PrepJMECustomNanoAOD_MC'},steps['NANOUP15']])
+steps['NANOUP15MC_PU25_Muon']=merge([{'--customise':'PhysicsTools/NanoAOD/custom_muon_cff.PrepMuonCustomNanoAOD'},steps['NANOUP15']])
 steps['NANOUP17Had']=merge([{'--filein':'file:step4_inMINIAODSIM.root'},steps['NANOUP17']])
 steps['NANOUP18'] = merge([{'--conditions': 'auto:phase1_2018_realistic', '--era': 'Run2_2018','-n':'10', '--filein':'file:step3_inMINIAODSIM.root', '--nThreads':'2'}, stepNanoEDMMCProd ])
 steps['NANOUP18Had']=merge([{'--filein':'file:step4_inMINIAODSIM.root'},steps['NANOUP18']])

diff --git a/PhysicsTools/NanoAOD/plugins/MuonSpecialVariables.cc b/PhysicsTools/NanoAOD/plugins/MuonSpecialVariables.cc
@@ -0,0 +1,120 @@
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "DataFormats/Common/interface/View.h"
+
+#include "DataFormats/PatCandidates/interface/UserData.h"
+#include "PhysicsTools/PatAlgos/interface/PATUserDataMerger.h"
+
+#include "DataFormats/MuonReco/interface/Muon.h"
+#include "DataFormats/MuonReco/interface/MuonSelectors.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+
+#include "DataFormats/VertexReco/interface/Vertex.h"
+#include "DataFormats/VertexReco/interface/VertexFwd.h"
+
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "DataFormats/TrackReco/interface/Track.h"
+#include "DataFormats/Math/interface/deltaR.h"
+#include "DataFormats/MuonReco/interface/MuonSimInfo.h"
+#include "MuonAnalysis/MuonAssociators/interface/PropagateToMuon.h"
+
+class MuonSpecialVariables : public edm::stream::EDProducer<> {
+public:
+  explicit MuonSpecialVariables(const edm::ParameterSet &iConfig)
+      : muonSrc_(consumes<std::vector<pat::Muon>>(iConfig.getParameter<edm::InputTag>("muonSrc"))),
+        vertexSrc_(consumes<std::vector<reco::Vertex>>(iConfig.getParameter<edm::InputTag>("vertexSrc"))),
+        trkSrc_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("trkSrc"))) {
+    produces<std::vector<pat::Muon>>();
+  }
+  ~MuonSpecialVariables() override {}
+
+  void setArbitratedTracker(pat::Muon &mu) const;
+  void setInnerTrackDxyz(pat::Muon &mu, const reco::Vertex &vtx) const;
+  void setTrkiso04(pat::Muon &mu, const std::vector<reco::Track> tracks) const;
+  void setNSegements(pat::Muon &mu) const;
+
+private:
+  void produce(edm::Event &iEvent, const edm::EventSetup &iSetup) override;
+
+  edm::EDGetToken muonSrc_;
+  edm::EDGetToken vertexSrc_;
+  edm::EDGetToken trkSrc_;
+};
+
+void MuonSpecialVariables::setArbitratedTracker(pat::Muon &mu) const {
+  mu.addUserInt("isArbitratedTracker", muon::isGoodMuon(mu, muon::TrackerMuonArbitrated));
+}
+
+void MuonSpecialVariables::setInnerTrackDxyz(pat::Muon &mu, const reco::Vertex &vtx) const {
+  float innerTrackDxy =
+      mu.innerTrack().isNonnull() ? mu.innerTrack()->dxy(reco::TrackBase::Point(vtx.x(), vtx.y(), vtx.z())) : -999.9;
+  float innerTrackDz =
+      mu.innerTrack().isNonnull() ? mu.innerTrack()->dz(reco::TrackBase::Point(vtx.x(), vtx.y(), vtx.z())) : -999.9;
+  mu.addUserFloat("innerTrackDxy", innerTrackDxy);
+  mu.addUserFloat("innerTrackDz", innerTrackDz);
+}
+
+void MuonSpecialVariables::setTrkiso04(pat::Muon &mu, const std::vector<reco::Track> tracks) const {
+  float energy = 0;
+  for (const auto &trk : tracks) {
+    if (deltaR(mu.eta(), mu.phi(), trk.eta(), trk.phi()) > 0.4)
+      continue;
+    energy += trk.pt();
+  }
+  float Trkiso04 = (energy - mu.pt()) / mu.pt();
+  float relTrkiso4 = (Trkiso04 > 0) ? Trkiso04 : 0;
+  mu.addUserFloat("relTrkiso4", relTrkiso4);
+}
+
+void MuonSpecialVariables::setNSegements(pat::Muon &mu) const {
+  int nsegments = 0;
+  for (auto &chamber : mu.matches()) {
+    if (chamber.id.det() != DetId::Muon)
+      continue;
+    if (chamber.id.subdetId() != MuonSubdetId::DT && chamber.id.subdetId() != MuonSubdetId::CSC)
+      continue;
+    nsegments += chamber.segmentMatches.size();
+  }
+  mu.addUserInt("nsegments", nsegments);
+}
+
+void MuonSpecialVariables::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
+  edm::Handle<std::vector<pat::Muon>> muonSrc;
+  iEvent.getByToken(muonSrc_, muonSrc);
+
+  edm::Handle<std::vector<reco::Vertex>> vertices;
+  iEvent.getByToken(vertexSrc_, vertices);
+
+  edm::Handle<reco::TrackCollection> trkCollection;
+  iEvent.getByToken(trkSrc_, trkCollection);
+
+  std::unique_ptr<std::vector<pat::Muon>> muonCollection(new std::vector<pat::Muon>(*muonSrc));
+
+  std::unique_ptr<std::vector<pat::Muon>> out(new std::vector<pat::Muon>());
+
+  out->reserve(muonCollection->size());
+
+  for (unsigned int i = 0; i < muonCollection->size(); i++) {
+    pat::Muon &mu = (*muonCollection).at(i);
+    setArbitratedTracker(mu);
+    int good_vertex = 0;
+    for (auto &vtx : *vertices) {
+      if (vtx.isFake() || !vtx.isValid())
+        continue;
+      setInnerTrackDxyz(mu, vtx);
+      good_vertex = 1;
+      break;
+    }
+    mu.addUserInt("isGoodVertex", good_vertex);
+    setTrkiso04(mu, *trkCollection);
+    setNSegements(mu);
+    out->push_back(mu);
+  }
+
+  iEvent.put(std::move(out));
+}
+
+#include "FWCore/Framework/interface/MakerMacros.h"
+DEFINE_FWK_MODULE(MuonSpecialVariables);
diff --git a/PhysicsTools/NanoAOD/plugins/pfTracksProducer.cc b/PhysicsTools/NanoAOD/plugins/pfTracksProducer.cc
@@ -0,0 +1,66 @@
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "DataFormats/Common/interface/View.h"
+
+#include "DataFormats/PatCandidates/interface/UserData.h"
+#include "PhysicsTools/PatAlgos/interface/PATUserDataMerger.h"
+#include "CommonTools/Utils/interface/StringCutObjectSelector.h"
+
+#include "DataFormats/VertexReco/interface/Vertex.h"
+#include "DataFormats/VertexReco/interface/VertexFwd.h"
+
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "DataFormats/TrackReco/interface/Track.h"
+#include "DataFormats/Math/interface/deltaR.h"
+#include "DataFormats/MuonReco/interface/MuonSimInfo.h"
+#include "MuonAnalysis/MuonAssociators/interface/PropagateToMuon.h"
+
+class pfTracksProducer : public edm::stream::EDProducer<> {
+public:
+  explicit pfTracksProducer(const edm::ParameterSet& iConfig)
+      : PFCands_(consumes<std::vector<pat::PackedCandidate>>(iConfig.getParameter<edm::InputTag>("PFCands"))),
+        LostTracks_(consumes<std::vector<pat::PackedCandidate>>(iConfig.getParameter<edm::InputTag>("lostTracks"))),
+        trkSelection_(iConfig.getParameter<std::string>("trkSelection")),
+        HighPurity_(iConfig.getParameter<bool>("TrkHPurity")) {
+    produces<reco::TrackCollection>();
+  }
+  ~pfTracksProducer() override {}
+
+private:
+  void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
+
+  edm::EDGetToken PFCands_;
+  edm::EDGetToken LostTracks_;
+  const StringCutObjectSelector<pat::PackedCandidate> trkSelection_;
+  const bool HighPurity_;
+};
+
+void pfTracksProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  edm::Handle<std::vector<pat::PackedCandidate>> pfcands;
+  iEvent.getByToken(PFCands_, pfcands);
+
+  edm::Handle<std::vector<pat::PackedCandidate>> lostTracks;
+  iEvent.getByToken(LostTracks_, lostTracks);
+
+  std::unique_ptr<std::vector<reco::Track>> tracks(new std::vector<reco::Track>());
+  tracks->reserve(pfcands->size() + lostTracks->size());
+
+  for (const auto& container : {pfcands, lostTracks}) {
+    for (const pat::PackedCandidate& trk : *container) {
+      if (!trk.hasTrackDetails())
+        continue;
+      if (!trkSelection_(trk))
+        continue;
+      if (HighPurity_ && !trk.trackHighPurity())
+        continue;
+      tracks->emplace_back(*trk.bestTrack());
+    }
+  }
+
+  iEvent.put(std::move(tracks));
+}
+
+#include "FWCore/Framework/interface/MakerMacros.h"
+DEFINE_FWK_MODULE(pfTracksProducer);