Update dEdx settings for HIN UPC miniAOD

PhysicsTools/PatAlgos/plugins/DeDxEstimatorRekeyer.cc

@@ -0,0 +1,128 @@
+#include "FWCore/Framework/interface/global/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "DataFormats/TrackReco/interface/DeDxData.h"
+#include "DataFormats/TrackReco/interface/DeDxHitInfo.h"
+#include "DataFormats/TrackReco/interface/TrackFwd.h"
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+
+//
+// class declaration
+//
+
+class DeDxEstimatorRekeyer : public edm::global::EDProducer<> {
+public:
+  explicit DeDxEstimatorRekeyer(const edm::ParameterSet&);
+  ~DeDxEstimatorRekeyer() override{};
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
+
+  template <typename T>
+  std::map<std::string, edm::EDGetTokenT<T>> getTokens(const std::vector<edm::InputTag>& tags) {
+    std::map<std::string, edm::EDGetTokenT<T>> tokens;
+    for (const auto& tag : tags)
+      tokens.emplace(tag.label(), consumes<T>(tag));
+    return tokens;
+  };
+
+  // ----------member data ---------------------------
+  const edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
+  const edm::EDGetTokenT<reco::DeDxHitInfoAss> dedxHitAssToken_;
+  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_;
+};
+
+void DeDxEstimatorRekeyer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("tracks", {"generalTracks"});
+  desc.add<edm::InputTag>("dedxHits", {"dedxHitInfo"});
+  desc.add<std::vector<edm::InputTag>>(
+      "packedCandidates",
+      {edm::InputTag("packedPFCandidates"), edm::InputTag("lostTracks"), edm::InputTag("lostTracks:eleTracks")});
+  desc.add<std::vector<edm::InputTag>>("dedxEstimators",
+                                       {edm::InputTag("dedxHarmonic2"), edm::InputTag("dedxPixelHarmonic2")});
+  descriptions.addWithDefaultLabel(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"))),
+      dedxEstimatorsTokens_(
+          getTokens<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<std::vector<edm::InputTag>>("dedxEstimators"))),
+      packedCandidatesTokens_(getTokens<pat::PackedCandidateCollection>(
+          iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))),
+      trk2pcTokens_(getTokens<edm::Association<pat::PackedCandidateCollection>>(
+          iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))) {
+  for (const auto& d : dedxEstimatorsTokens_)
+    produces<edm::ValueMap<reco::DeDxData>>(d.first);
+  produces<reco::DeDxHitInfoCollection>();
+  produces<reco::DeDxHitInfoAss>();
+}
+
+void DeDxEstimatorRekeyer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
+  // Get input collections
+  const auto& tracks = iEvent.getHandle(tracksToken_);
+  std::vector<reco::TrackRef> trackRefs;
+  trackRefs.reserve(tracks->size());
+  for (auto track = tracks->begin(); track != tracks->end(); track++)
+    trackRefs.emplace_back(tracks, track - tracks->begin());
+
+  typedef std::map<pat::PackedCandidateRef, reco::TrackRef> PCTrkMap;
+  std::vector<std::pair<edm::Handle<pat::PackedCandidateCollection>, PCTrkMap>> pcTrkMap;
+  pcTrkMap.reserve(packedCandidatesTokens_.size());
+  for (const auto& p : packedCandidatesTokens_) {
+    PCTrkMap map;
+    const auto& trk2pc = iEvent.get(trk2pcTokens_.at(p.first));
+    for (const auto& track : trackRefs) {
+      const auto& pc = trk2pc[track];
+      if (pc.isNonnull())
+        map.emplace(pc, track);
+    }
+    pcTrkMap.emplace_back(iEvent.getHandle(p.second), map);
+  }
+
+  // Rekey dEdx estimators
+  for (const auto& d : dedxEstimatorsTokens_) {
+    const auto& dedxEstimators = iEvent.get(d.second);
+    auto trackDeDxValueMap = std::make_unique<edm::ValueMap<reco::DeDxData>>();
+    edm::ValueMap<reco::DeDxData>::Filler filler(*trackDeDxValueMap);
+    // Loop over packed candidates
+    for (const auto& h : pcTrkMap) {
+      std::vector<reco::DeDxData> dedxEstimate(h.first->size());
+      for (const auto& p : h.second)
+        dedxEstimate[p.first.key()] = dedxEstimators[p.second];
+      filler.insert(h.first, dedxEstimate.begin(), dedxEstimate.end());
+    }
+    // Fill the value map and put it into the event
+    filler.fill();
+    iEvent.put(std::move(trackDeDxValueMap), d.first);
+  }
+
+  // Rekey dEdx hit info
+  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);
+  // 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]);
+    }
+    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));
+}
+
+//define this as a plug-in
+#include "FWCore/Framework/interface/MakerMacros.h"
+DEFINE_FWK_MODULE(DeDxEstimatorRekeyer);

PhysicsTools/PatAlgos/python/slimming/MicroEventContent_cff.py

@@ -166,6 +166,7 @@
     'keep recoCentrality_hiCentrality_*_*',
     'keep recoClusterCompatibility_hiClusterCompatibility_*_*',
     'keep QIE10DataFrameHcalDataFrameContainer_hcalDigis_ZDC_*',
+    'keep *_dedxEstimator_*_*',
 ]
 
 from Configuration.Eras.Modifier_run3_upc_cff import run3_upc

PhysicsTools/PatAlgos/python/slimming/isolatedTracks_cfi.py

@@ -76,9 +76,6 @@
 from Configuration.ProcessModifiers.pp_on_AA_cff import pp_on_AA
 pp_on_AA.toModify(isolatedTracks, useHighPurity = True)
 
-from Configuration.Eras.Modifier_run3_upc_cff import run3_upc
-run3_upc.toModify(isolatedTracks, pT_cut = 0.0, pT_cut_noIso = 0.0, saveDeDxHitInfoCut = "")
-
 def miniAOD_customizeIsolatedTracksFastSim(process):
     """Switch off dE/dx hit info on fast sim, as it's not available"""
     process.isolatedTracks.saveDeDxHitInfo = False

PhysicsTools/PatAlgos/python/slimming/slimming_cff.py

@@ -124,7 +124,11 @@
 _photonDRN.toReplaceWith(slimmingTask, cms.Task(slimmingTask.copy(), patPhotonsDRN))
 
 from Configuration.Eras.Modifier_run3_upc_cff import run3_upc
-run3_upc.toReplaceWith(slimmingTask, cms.Task(slimmingTask.copy(), hiPixelTracks, packedPFCandidateTrackChi2, lostTrackChi2))
+from PhysicsTools.PatAlgos.modules import DeDxEstimatorRekeyer
+dedxEstimator = DeDxEstimatorRekeyer()
+from Configuration.Eras.Modifier_run3_egamma_2023_cff import run3_egamma_2023
+(run3_upc & ~run3_egamma_2023).toModify(dedxEstimator, dedxEstimators = ["dedxHarmonic2", "dedxPixelHarmonic2", "dedxPixelLikelihood", "dedxStripLikelihood", "dedxAllLikelihood"])
+run3_upc.toReplaceWith(slimmingTask, cms.Task(slimmingTask.copy(), hiPixelTracks, packedPFCandidateTrackChi2, lostTrackChi2, dedxEstimator))
 
 from Configuration.Eras.Modifier_ppRef_2024_cff import ppRef_2024
 ppRef_2024.toReplaceWith(slimmingTask, cms.Task(slimmingTask.copy(), packedPFCandidateTrackChi2, lostTrackChi2))