Store track momentum at hit for HIN UPC dEdX

Andre Govinda Stahl Leiton · Andre Govinda Stahl Leiton · commit 5f1453ac892f · 2024-07-13T21:10:34.000+02:00
diff --git a/PhysicsTools/PatAlgos/plugins/DeDxEstimatorRekeyer.cc b/PhysicsTools/PatAlgos/plugins/DeDxEstimatorRekeyer.cc
@@ -30,6 +30,7 @@ class DeDxEstimatorRekeyer : public edm::global::EDProducer<> {
   // ----------member data ---------------------------
   const edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
   const edm::EDGetTokenT<reco::DeDxHitInfoAss> dedxHitAssToken_;
+  const edm::EDGetTokenT<edm::ValueMap<std::vector<float>>> dedxHitMomToken_;
   const std::map<std::string, edm::EDGetTokenT<edm::ValueMap<reco::DeDxData>>> dedxEstimatorsTokens_;
   const std::map<std::string, edm::EDGetTokenT<pat::PackedCandidateCollection>> packedCandidatesTokens_;
   const std::map<std::string, edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>>> trk2pcTokens_;
@@ -39,6 +40,7 @@ void DeDxEstimatorRekeyer::fillDescriptions(edm::ConfigurationDescriptions& desc
   edm::ParameterSetDescription desc;
   desc.add<edm::InputTag>("tracks", {"generalTracks"});
   desc.add<edm::InputTag>("dedxHits", {"dedxHitInfo"});
+  desc.add<edm::InputTag>("dedxMomentum", {"dedxHitInfo:momentumAtHit"});
   desc.add<std::vector<edm::InputTag>>(
       "packedCandidates",
       {edm::InputTag("packedPFCandidates"), edm::InputTag("lostTracks"), edm::InputTag("lostTracks:eleTracks")});
@@ -50,6 +52,8 @@ void DeDxEstimatorRekeyer::fillDescriptions(edm::ConfigurationDescriptions& desc
 DeDxEstimatorRekeyer::DeDxEstimatorRekeyer(const edm::ParameterSet& iConfig)
     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracks"))),
       dedxHitAssToken_(consumes<reco::DeDxHitInfoAss>(iConfig.getParameter<edm::InputTag>("dedxHits"))),
+      dedxHitMomToken_(
+          consumes<edm::ValueMap<std::vector<float>>>(iConfig.getParameter<edm::InputTag>("dedxMomentum"))),
       dedxEstimatorsTokens_(
           getTokens<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<std::vector<edm::InputTag>>("dedxEstimators"))),
       packedCandidatesTokens_(getTokens<pat::PackedCandidateCollection>(
@@ -60,6 +64,7 @@ DeDxEstimatorRekeyer::DeDxEstimatorRekeyer(const edm::ParameterSet& iConfig)
     produces<edm::ValueMap<reco::DeDxData>>(d.first);
   produces<reco::DeDxHitInfoCollection>();
   produces<reco::DeDxHitInfoAss>();
+  produces<edm::ValueMap<std::vector<float>>>("momentumAtHit");
 }
 
 void DeDxEstimatorRekeyer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
@@ -102,25 +107,36 @@ void DeDxEstimatorRekeyer::produce(edm::StreamID, edm::Event& iEvent, const edm:
   }
 
   // Rekey dEdx hit info
+  const auto& dedxHitMom = iEvent.get(dedxHitMomToken_);
   const auto& dedxHitAss = iEvent.get(dedxHitAssToken_);
   const auto& dedxHitInfoHandle = iEvent.getRefBeforePut<reco::DeDxHitInfoCollection>();
   auto dedxHitInfoAssociation = std::make_unique<reco::DeDxHitInfoAss>(dedxHitInfoHandle);
   reco::DeDxHitInfoAss::Filler filler(*dedxHitInfoAssociation);
   auto resultdedxHitColl = std::make_unique<reco::DeDxHitInfoCollection>();
   resultdedxHitColl->reserve(pcTrkMap.size() > 0 ? pcTrkMap.size() * pcTrkMap[0].second.size() : 0);
+  std::vector<std::vector<float>> momenta;
+  momenta.reserve(resultdedxHitColl->capacity());
   // Loop over packed candidates
   for (const auto& h : pcTrkMap) {
     std::vector<int> indices(h.first->size(), -1);
     for (const auto& p : h.second) {
       indices[p.first.key()] = resultdedxHitColl->size();
-      resultdedxHitColl->emplace_back(*dedxHitAss[p.second]);
+      const auto& dedxHit = dedxHitAss[p.second];
+      resultdedxHitColl->emplace_back(*dedxHit);
+      momenta.emplace_back(dedxHitMom[dedxHit]);
     }
     filler.insert(h.first, indices.begin(), indices.end());
   }
   const auto& dedxHitCollHandle = iEvent.put(std::move(resultdedxHitColl));
   // Fill the association map and put it into the event
   filler.fill();
   iEvent.put(std::move(dedxHitInfoAssociation));
+  // Fill the value map and put it into the event
+  auto dedxMomenta = std::make_unique<edm::ValueMap<std::vector<float>>>();
+  edm::ValueMap<std::vector<float>>::Filler mfiller(*dedxMomenta);
+  mfiller.insert(dedxHitCollHandle, momenta.begin(), momenta.end());
+  mfiller.fill();
+  iEvent.put(std::move(dedxMomenta), "momentumAtHit");
 }
 
 //define this as a plug-in
diff --git a/RecoTracker/DeDx/plugins/DeDxEstimatorProducer.cc b/RecoTracker/DeDx/plugins/DeDxEstimatorProducer.cc
@@ -130,7 +130,6 @@ void DeDxEstimatorProducer::produce(edm::Event& iEvent, const edm::EventSetup& i
     int NClusterSaturating = 0;
     DeDxHitCollection dedxHits;
 
-    dedxHits.reserve(track->recHitsSize() / 2);
     if (useDeDxHits) {
       if (usePixel)
         dedxHits = (*pixelDeDxHits)[track];
diff --git a/RecoTracker/DeDx/plugins/DeDxHitInfoProducer.cc b/RecoTracker/DeDx/plugins/DeDxHitInfoProducer.cc
@@ -59,6 +59,7 @@ class DeDxHitInfoProducer : public edm::stream::EDProducer<> {
                   const float trackMomentum,
                   const LocalVector& trackDirection,
                   reco::DeDxHitInfo& hitDeDxInfo,
+                  std::vector<float>& hitMomentum,
                   const LocalPoint& hitLocalPos);
 
   // ----------member data ---------------------------
@@ -76,6 +77,7 @@ class DeDxHitInfoProducer : public edm::stream::EDProducer<> {
   const bool useCalibration_;
   const bool doShapeTest_;
   const bool usePixelShape_;
+  const bool storeMomentumAtHit_;
 
   const unsigned int lowPtTracksPrescalePass_, lowPtTracksPrescaleFail_;
   GenericTruncatedAverageDeDxEstimator lowPtTracksEstimator_;
@@ -116,6 +118,7 @@ DeDxHitInfoProducer::DeDxHitInfoProducer(const edm::ParameterSet& iConfig)
       useCalibration_(iConfig.getParameter<bool>("useCalibration")),
       doShapeTest_(iConfig.getParameter<bool>("shapeTest")),
       usePixelShape_(not iConfig.getParameter<edm::InputTag>("clusterShapeCache").label().empty()),
+      storeMomentumAtHit_(iConfig.getParameter<bool>("storeMomentumAtHit")),
       lowPtTracksPrescalePass_(iConfig.getParameter<uint32_t>("lowPtTracksPrescalePass")),
       lowPtTracksPrescaleFail_(iConfig.getParameter<uint32_t>("lowPtTracksPrescaleFail")),
       lowPtTracksEstimator_(iConfig.getParameter<edm::ParameterSet>("lowPtTracksEstimatorParameters")),
@@ -129,6 +132,8 @@ DeDxHitInfoProducer::DeDxHitInfoProducer(const edm::ParameterSet& iConfig)
   produces<reco::DeDxHitInfoCollection>();
   produces<reco::DeDxHitInfoAss>();
   produces<edm::ValueMap<int>>("prescale");
+  if (storeMomentumAtHit_)
+    produces<edm::ValueMap<std::vector<float>>>("momentumAtHit");
 
   if (!usePixel_ && !useStrip_)
     edm::LogError("DeDxHitsProducer") << "No Pixel Hits NOR Strip Hits will be saved.  Running this module is useless";
@@ -160,6 +165,7 @@ void DeDxHitInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
 
   std::vector<int> indices;
   std::vector<int> prescales;
+  std::vector<std::vector<float>> hitMomenta;
   uint64_t state[2] = {iEvent.id().event(), iEvent.id().luminosityBlock()};
   for (unsigned int j = 0; j < trackCollection.size(); j++) {
     const reco::Track& track = trackCollection[j];
@@ -181,15 +187,18 @@ void DeDxHitInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
     }
 
     reco::DeDxHitInfo hitDeDxInfo;
+    std::vector<float> hitMomentum;
     auto const& trajParams = track.extra()->trajParams();
     auto hb = track.recHitsBegin();
     for (unsigned int h = 0; h < track.recHitsSize(); h++) {
       auto recHit = *(hb + h);
       if (!trackerHitRTTI::isFromDet(*recHit))
         continue;
 
-      processHit(recHit, track.p(), trajParams[h].direction(), hitDeDxInfo, trajParams[h].position());
+      const auto& traj = trajParams[h];
+      processHit(recHit, traj.momentum().mag(), traj.direction(), hitDeDxInfo, hitMomentum, traj.position());
     }
+    assert(!storeMomentumAtHit_ || hitMomentum.size() == hitDeDxInfo.size());
 
     if (!passPt) {
       std::vector<DeDxHit> hits;
@@ -230,6 +239,8 @@ void DeDxHitInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
     }
     indices.push_back(resultdedxHitColl->size());
     resultdedxHitColl->push_back(hitDeDxInfo);
+    if (storeMomentumAtHit_)
+      hitMomenta.push_back(hitMomentum);
   }
   ///////////////////////////////////////
 
@@ -247,6 +258,14 @@ void DeDxHitInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
   pfiller.insert(dedxHitCollHandle, prescales.begin(), prescales.end());
   pfiller.fill();
   iEvent.put(std::move(dedxPrescale), "prescale");
+
+  if (storeMomentumAtHit_) {
+    auto dedxMomenta = std::make_unique<edm::ValueMap<std::vector<float>>>();
+    edm::ValueMap<std::vector<float>>::Filler mfiller(*dedxMomenta);
+    mfiller.insert(dedxHitCollHandle, hitMomenta.begin(), hitMomenta.end());
+    mfiller.fill();
+    iEvent.put(std::move(dedxMomenta), "momentumAtHit");
+  }
 }
 
 void DeDxHitInfoProducer::processRec(reco::DeDxHitInfo& hitDeDxInfo,
@@ -279,6 +298,7 @@ void DeDxHitInfoProducer::processHit(const TrackingRecHit* recHit,
                                      const float trackMomentum,
                                      const LocalVector& trackDirection,
                                      reco::DeDxHitInfo& hitDeDxInfo,
+                                     std::vector<float>& hitMomentum,
                                      const LocalPoint& hitLocalPos) {
   auto const& thit = static_cast<BaseTrackerRecHit const&>(*recHit);
   if (!thit.isValid())
@@ -327,10 +347,14 @@ void DeDxHitInfoProducer::processHit(const TrackingRecHit* recHit,
     const SiStripMatchedRecHit2D* matchedHit = dynamic_cast<const SiStripMatchedRecHit2D*>(recHit);
     if (!matchedHit)
       return;
+    if (storeMomentumAtHit_)
+      hitMomentum.push_back(trackMomentum);
 
     processRec(hitDeDxInfo, matchedHit->monoHit(), hitLocalPos, trackDirection, cosineAbs);
     processRec(hitDeDxInfo, matchedHit->stereoHit(), hitLocalPos, trackDirection, cosineAbs);
   }
+  if (storeMomentumAtHit_ && (clus.isPixel() || clus.isStrip()))
+    hitMomentum.push_back(trackMomentum);
 }
 
 //define this as a plug-in
diff --git a/RecoTracker/DeDx/python/dedxEstimators_cff.py b/RecoTracker/DeDx/python/dedxEstimators_cff.py
@@ -17,6 +17,7 @@
     calibrationPath    = cms.string("file:Gains.root"),
     shapeTest          = cms.bool(True),
     clusterShapeCache  = cms.InputTag("siPixelClusterShapeCache"),
+    storeMomentumAtHit = cms.bool(False),
 
     lowPtTracksPrescalePass = cms.uint32(100),   # prescale factor for low pt tracks above the dEdx cut
     lowPtTracksPrescaleFail = cms.uint32(2000), # prescale factor for low pt tracks below the dEdx cut
@@ -91,7 +92,7 @@
 
 # dEdx for Run-3 UPC
 from Configuration.Eras.Modifier_run3_upc_cff import run3_upc
-run3_upc.toModify(dedxHitInfo, minTrackPt = 0)
+run3_upc.toModify(dedxHitInfo, minTrackPt = 0, storeMomentumAtHit = True)
 
 from RecoTracker.DeDx.dedxHitCalibrator_cfi import dedxHitCalibrator as _dedxHitCalibrator
 from SimGeneral.MixingModule.SiStripSimParameters_cfi import SiStripSimBlock as _SiStripSimBlock
