Add mu beamspot constrained pT. Increase mu eta and phi precision to 16 bits

PhysicsTools/NanoAOD/python/muons_cff.py

@@ -119,6 +119,11 @@
     weightFile = "PhysicsTools/NanoAOD/data/mu_BDTG_2016.weights.xml",
 )
 
+from TrackingTools.TransientTrack.TransientTrackBuilder_cfi import *
+muonBSConstrain = cms.EDProducer("MuonBeamspotConstraintValueMapProducer",
+    src = cms.InputTag("linkedObjects","muons"),
+)
+
 muonTable = simpleCandidateFlatTableProducer.clone(
     src = cms.InputTag("linkedObjects","muons"),
     name = cms.string("Muon"),
@@ -175,9 +180,15 @@
         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"), "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),
     ),
 )
 
+# Increase precision of eta and phi
+muonTable.variables.eta.precision = 16
+muonTable.variables.phi.precision = 16
 
 (run2_nanoAOD_106Xv2 | run3_nanoAOD_122).toModify(muonTable.variables,mvaMuID=None).toModify(
      muonTable.variables, mvaMuID = Var("userFloat('mvaIDMuon')", float, doc="MVA-based ID score",precision=6))
@@ -212,8 +223,5 @@
 
 muonTask = cms.Task(slimmedMuonsUpdated,isoForMu,ptRatioRelForMu,slimmedMuonsWithUserData,finalMuons,finalLooseMuons )
 muonMCTask = cms.Task(muonsMCMatchForTable,muonMCTable)
-muonTablesTask = cms.Task(muonMVATTH,muonMVALowPt,muonTable,muonMVAID)
-
-
-
+muonTablesTask = cms.Task(muonMVATTH,muonMVALowPt,muonBSConstrain,muonTable,muonMVAID)
 

PhysicsTools/NanoAOD/python/nanoDQM_cfi.py

@@ -504,6 +504,9 @@
             ),
             plots = cms.VPSet(
                 Count1D('_size', 5, -0.5, 4.5, 'slimmedMuons after basic selection (pt > 3 && track.isNonnull && isLooseMuon)'),
+                Plot1D('bsConstrainedPt', '', 20, 0., 200., 'pt with beamspot constraint'),
+                Plot1D('bsConstrainedPtErr', '', 20, 0., 20., 'pt error w. beamspot constraint'),
+                Plot1D('bsConstrainedChi2', '', 20, 0., 40., 'chi2 of beamspot constraint'),
                 Plot1D('charge', 'charge', 3, -1.5, 1.5, 'electric charge'),
                 Plot1D('dxy', 'dxy', 20, -0.1, 0.1, 'dxy (with sign) wrt first PV, in cm'),
                 Plot1D('dxyErr', 'dxyErr', 20, 0, 0.1, 'dxy uncertainty, in cm'),

RecoMuon/GlobalTrackingTools/plugins/BuildFile.xml

@@ -8,6 +8,7 @@
   <use name="Geometry/Records"/>
   <use name="RecoMuon/TrackingTools"/>
   <use name="RecoMuon/GlobalTrackingTools"/>
+  <use name="RecoVertex/KalmanVertexFit"/>
   <use name="boost_regex"/>
   <flags EDM_PLUGIN="1"/>
 </library>

RecoMuon/GlobalTrackingTools/plugins/MuonBeamspotConstraintValueMapProducer.cc

@@ -0,0 +1,140 @@
+/** \class MuonBeamspotConstraintValueMapProducer
+ * Compute muon pt and ptErr after beamspot constraint.
+ *
+ */
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+#include "DataFormats/Common/interface/ValueMap.h"
+#include "FWCore/Framework/interface/global/EDProducer.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
+#include "TrackingTools/Records/interface/TransientTrackRecord.h"
+#include "TrackingTools/TransientTrack/interface/TransientTrack.h"
+#include "RecoVertex/KalmanVertexFit/interface/SingleTrackVertexConstraint.h"
+#include "DataFormats/VertexReco/interface/Vertex.h"
+
+class MuonBeamspotConstraintValueMapProducer : public edm::global::EDProducer<> {
+public:
+  explicit MuonBeamspotConstraintValueMapProducer(const edm::ParameterSet& config)
+      : muonToken_(consumes<pat::MuonCollection>(config.getParameter<edm::InputTag>("src"))),
+        beamSpotToken_(consumes<reco::BeamSpot>(config.getParameter<edm::InputTag>("beamspot"))),
+        PrimaryVertexToken_(consumes<reco::VertexCollection>(config.getParameter<edm::InputTag>("vertices"))),
+        ttbToken_(esConsumes(edm::ESInputTag("", "TransientTrackBuilder"))) {
+    produces<edm::ValueMap<float>>("muonBSConstrainedPt");
+    produces<edm::ValueMap<float>>("muonBSConstrainedPtErr");
+    produces<edm::ValueMap<float>>("muonBSConstrainedChi2");
+  }
+
+  ~MuonBeamspotConstraintValueMapProducer() override = default;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+    edm::ParameterSetDescription desc;
+    desc.add<edm::InputTag>("src", edm::InputTag("muons"))->setComment("Muon collection");
+    desc.add<edm::InputTag>("beamspot", edm::InputTag("offlineBeamSpot"))->setComment("Beam spot collection");
+    desc.add<edm::InputTag>("vertices", edm::InputTag("offlineSlimmedPrimaryVertices"))
+        ->setComment("Primary vertex collection");
+
+    descriptions.addWithDefaultLabel(desc);
+  }
+
+private:
+  void produce(edm::StreamID streamID, edm::Event& event, const edm::EventSetup& setup) const override {
+    edm::Handle<pat::MuonCollection> muons;
+    event.getByToken(muonToken_, muons);
+
+    edm::Handle<reco::BeamSpot> beamSpotHandle;
+    event.getByToken(beamSpotToken_, beamSpotHandle);
+
+    edm::ESHandle<TransientTrackBuilder> ttkb = setup.getHandle(ttbToken_);
+
+    std::vector<float> pts, ptErrs, chi2s;
+    pts.reserve(muons->size());
+    ptErrs.reserve(muons->size());
+    chi2s.reserve(muons->size());
+
+    for (const auto& muon : *muons) {
+      bool tbd = true;
+      if (beamSpotHandle.isValid()) {
+        double BeamWidthX = beamSpotHandle->BeamWidthX();
+        double BeamWidthXError = beamSpotHandle->BeamWidthXError();
+        double BeamWidthY = beamSpotHandle->BeamWidthY();
+        double BeamWidthYError = beamSpotHandle->BeamWidthYError();
+        // Protect for mis-reconstructed beamspots (note that
+        // SingleTrackVertexConstraint uses the width for the constraint,
+        // not the error)
+        if ((BeamWidthXError / BeamWidthX < 0.3) && (BeamWidthYError / BeamWidthY < 0.3)) {
+          SingleTrackVertexConstraint::BTFtuple btft =
+              stvc.constrain(ttkb->build(muon.muonBestTrack()), *beamSpotHandle);
+          if (std::get<0>(btft)) {
+            const reco::Track& trkBS = std::get<1>(btft).track();
+            pts.push_back(trkBS.pt());
+            ptErrs.push_back(trkBS.ptError());
+            chi2s.push_back(std::get<2>(btft));
+            tbd = false;
+          }
+        }
+      }
+
+      if (tbd) {
+        // Invalid BS; use PV instead
+        edm::Handle<reco::VertexCollection> pvHandle;
+        event.getByToken(PrimaryVertexToken_, pvHandle);
+
+        if (pvHandle.isValid() && !pvHandle->empty()) {
+          auto pv = pvHandle->at(0);
+          VertexState pvs = VertexState(GlobalPoint(Basic3DVector<float>(pv.position())), GlobalError(pv.covariance()));
+
+          SingleTrackVertexConstraint::BTFtuple btft = stvc.constrain(ttkb->build(muon.muonBestTrack()), pvs);
+          if (std::get<0>(btft)) {
+            const reco::Track& trkBS = std::get<1>(btft).track();
+            pts.push_back(trkBS.pt());
+            ptErrs.push_back(trkBS.ptError());
+            chi2s.push_back(std::get<2>(btft));
+            tbd = false;
+          }
+        }
+      }
+
+      if (tbd) {
+        // Fall-back case, keep the unconstrained values
+        pts.push_back(muon.pt());
+        ptErrs.push_back(muon.bestTrack()->ptError());
+        chi2s.push_back(-1.f);
+      }
+    }
+
+    {
+      std::unique_ptr<edm::ValueMap<float>> valueMap(new edm::ValueMap<float>());
+      edm::ValueMap<float>::Filler filler(*valueMap);
+      filler.insert(muons, pts.begin(), pts.end());
+      filler.fill();
+      event.put(std::move(valueMap), "muonBSConstrainedPt");
+    }
+
+    {
+      std::unique_ptr<edm::ValueMap<float>> valueMap(new edm::ValueMap<float>());
+      edm::ValueMap<float>::Filler filler(*valueMap);
+      filler.insert(muons, ptErrs.begin(), ptErrs.end());
+      filler.fill();
+      event.put(std::move(valueMap), "muonBSConstrainedPtErr");
+    }
+
+    {
+      std::unique_ptr<edm::ValueMap<float>> valueMap(new edm::ValueMap<float>());
+      edm::ValueMap<float>::Filler filler(*valueMap);
+      filler.insert(muons, chi2s.begin(), chi2s.end());
+      filler.fill();
+      event.put(std::move(valueMap), "muonBSConstrainedChi2");
+    }
+  }
+
+  edm::EDGetTokenT<pat::MuonCollection> muonToken_;
+  edm::EDGetTokenT<reco::BeamSpot> beamSpotToken_;
+  edm::EDGetTokenT<reco::VertexCollection> PrimaryVertexToken_;
+  edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> ttbToken_;
+  SingleTrackVertexConstraint stvc;
+};
+
+DEFINE_FWK_MODULE(MuonBeamspotConstraintValueMapProducer);