NanoAOD/plugins/MuonSpecialVariables.cc

0001 #include "FWCore/Framework/interface/stream/EDProducer.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/Utilities/interface/InputTag.h"
0005 #include "DataFormats/Common/interface/View.h"
0006
0007 #include "DataFormats/PatCandidates/interface/UserData.h"
0008 #include "PhysicsTools/PatAlgos/interface/PATUserDataMerger.h"
0009
0010 #include "DataFormats/MuonReco/interface/Muon.h"
0011 #include "DataFormats/MuonReco/interface/MuonSelectors.h"
0012 #include "DataFormats/PatCandidates/interface/Muon.h"
0013
0014 #include "DataFormats/VertexReco/interface/Vertex.h"
0015 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0016
0017 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0018 #include "DataFormats/TrackReco/interface/Track.h"
0019 #include "DataFormats/Math/interface/deltaR.h"
0020 #include "DataFormats/MuonReco/interface/MuonSimInfo.h"
0021 #include "MuonAnalysis/MuonAssociators/interface/PropagateToMuon.h"
0022
0023 class MuonSpecialVariables : public edm::stream::EDProducer<> {
0024 public:
0025   explicit MuonSpecialVariables(const edm::ParameterSet &iConfig)
0026       : muonSrc_(consumes<std::vector<pat::Muon>>(iConfig.getParameter<edm::InputTag>("muonSrc"))),
0027         vertexSrc_(consumes<std::vector<reco::Vertex>>(iConfig.getParameter<edm::InputTag>("vertexSrc"))),
0028         trkSrc_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("trkSrc"))) {
0029     produces<std::vector<pat::Muon>>();
0030   }
0031   ~MuonSpecialVariables() override {}
0032
0033   void setArbitratedTracker(pat::Muon &mu) const;
0034   void setInnerTrackDxyz(pat::Muon &mu, const reco::Vertex &vtx) const;
0035   void setTrkiso04(pat::Muon &mu, const std::vector<reco::Track> tracks) const;
0036   void setTrkiso03(pat::Muon &mu, const std::vector<reco::Track> tracks) const;
0037   void setNSegements(pat::Muon &mu) const;
0038
0039 private:
0040   void produce(edm::Event &iEvent, const edm::EventSetup &iSetup) override;
0041
0042   edm::EDGetToken muonSrc_;
0043   edm::EDGetToken vertexSrc_;
0044   edm::EDGetToken trkSrc_;
0045 };
0046
0047 void MuonSpecialVariables::setArbitratedTracker(pat::Muon &mu) const {
0048   mu.addUserInt("isArbitratedTracker", muon::isGoodMuon(mu, muon::TrackerMuonArbitrated));
0049 }
0050
0051 void MuonSpecialVariables::setInnerTrackDxyz(pat::Muon &mu, const reco::Vertex &vtx) const {
0052   float innerTrackDxy =
0053       mu.innerTrack().isNonnull() ? mu.innerTrack()->dxy(reco::TrackBase::Point(vtx.x(), vtx.y(), vtx.z())) : -999.9;
0054   float innerTrackDz =
0055       mu.innerTrack().isNonnull() ? mu.innerTrack()->dz(reco::TrackBase::Point(vtx.x(), vtx.y(), vtx.z())) : -999.9;
0056   mu.addUserFloat("innerTrackDxy", innerTrackDxy);
0057   mu.addUserFloat("innerTrackDz", innerTrackDz);
0058 }
0059
0060 void MuonSpecialVariables::setTrkiso04(pat::Muon &mu, const std::vector<reco::Track> tracks) const {
0061   float energy = 0;
0062   for (const auto &trk : tracks) {
0063     if (deltaR(mu.eta(), mu.phi(), trk.eta(), trk.phi()) > 0.4)
0064       continue;
0065     energy += trk.pt();
0066   }
0067   float Trkiso04 = (energy - mu.pt());
0068   float absTrkiso04 = (Trkiso04 > 0) ? Trkiso04 : 0;
0069   mu.addUserFloat("absTrkiso04", absTrkiso04);
0070 }
0071
0072 void MuonSpecialVariables::setTrkiso03(pat::Muon &mu, const std::vector<reco::Track> tracks) const {
0073   float energy = 0;
0074   for (const auto &trk : tracks) {
0075     if (deltaR(mu.eta(), mu.phi(), trk.eta(), trk.phi()) > 0.3)
0076       continue;
0077     energy += trk.pt();
0078   }
0079   float Trkiso03 = (energy - mu.pt());
0080   float absTrkiso03 = (Trkiso03 > 0) ? Trkiso03 : 0;
0081   mu.addUserFloat("absTrkiso03", absTrkiso03);
0082 }
0083
0084 void MuonSpecialVariables::setNSegements(pat::Muon &mu) const {
0085   int nsegments = 0;
0086   for (auto &chamber : mu.matches()) {
0087     if (chamber.id.det() != DetId::Muon)
0088       continue;
0089     if (chamber.id.subdetId() != MuonSubdetId::DT && chamber.id.subdetId() != MuonSubdetId::CSC)
0090       continue;
0091     nsegments += chamber.segmentMatches.size();
0092   }
0093   mu.addUserInt("nsegments", nsegments);
0094 }
0095
0096 void MuonSpecialVariables::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
0097   edm::Handle<std::vector<pat::Muon>> muonSrc;
0098   iEvent.getByToken(muonSrc_, muonSrc);
0099
0100   edm::Handle<std::vector<reco::Vertex>> vertices;
0101   iEvent.getByToken(vertexSrc_, vertices);
0102
0103   edm::Handle<reco::TrackCollection> trkCollection;
0104   iEvent.getByToken(trkSrc_, trkCollection);
0105
0106   std::unique_ptr<std::vector<pat::Muon>> muonCollection(new std::vector<pat::Muon>(*muonSrc));
0107
0108   std::unique_ptr<std::vector<pat::Muon>> out(new std::vector<pat::Muon>());
0109
0110   out->reserve(muonCollection->size());
0111
0112   for (unsigned int i = 0; i < muonCollection->size(); i++) {
0113     pat::Muon &mu = (*muonCollection).at(i);
0114     setArbitratedTracker(mu);
0115     int good_vertex = 0;
0116     for (auto &vtx : *vertices) {
0117       if (vtx.isFake() || !vtx.isValid())
0118         continue;
0119       setInnerTrackDxyz(mu, vtx);
0120       good_vertex = 1;
0121       break;
0122     }
0123     mu.addUserInt("isGoodVertex", good_vertex);
0124     setTrkiso04(mu, *trkCollection);
0125     setTrkiso03(mu, *trkCollection);
0126     setNSegements(mu);
0127     out->push_back(mu);
0128   }
0129
0130   iEvent.put(std::move(out));
0131 }
0132
0133 #include "FWCore/Framework/interface/MakerMacros.h"
0134 DEFINE_FWK_MODULE(MuonSpecialVariables);