Merge pull request #30566 from abaty/Chi2MapMaster

Addition of packedCandidates track chi2 information to pp_on_AA miniAOD

PhysicsTools/PatAlgos/plugins/PackedCandidateTrackChi2Producer.cc

@@ -0,0 +1,113 @@
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "DataFormats/Common/interface/ValueMap.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+#include "DataFormats/TrackReco/interface/Track.h"
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "DataFormats/Math/interface/libminifloat.h"
+
+namespace pat {
+
+  class PackedCandidateTrackChi2Producer : public edm::stream::EDProducer<> {
+    typedef edm::ValueMap<float> FloatMap;
+
+  public:
+    explicit PackedCandidateTrackChi2Producer(const edm::ParameterSet& iConfig)
+        : candidateToken_(consumes<pat::PackedCandidateCollection>(iConfig.getParameter<edm::InputTag>("candidates"))),
+          trackToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("trackCollection"))),
+          doLostTracks_(iConfig.getParameter<bool>("doLostTracks")) {
+      if (doLostTracks_) {
+        track2LostTrackToken_ = consumes<edm::Association<pat::PackedCandidateCollection>>(
+            iConfig.getParameter<edm::InputTag>("candidates"));
+      } else {
+        candidate2PFToken_ =
+            consumes<edm::Association<reco::PFCandidateCollection>>(iConfig.getParameter<edm::InputTag>("candidates"));
+      }
+
+      produces<FloatMap>();
+    }
+
+    void produce(edm::Event&, const edm::EventSetup&) override;
+
+    static void fillDescriptions(edm::ConfigurationDescriptions&);
+
+  private:
+    const edm::EDGetTokenT<pat::PackedCandidateCollection> candidateToken_;
+    edm::EDGetTokenT<edm::Association<reco::PFCandidateCollection>> candidate2PFToken_;
+    edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>> track2LostTrackToken_;
+    const edm::EDGetTokenT<reco::TrackCollection> trackToken_;
+    const bool doLostTracks_;
+    static const uint8_t roundingPrecision = 8;
+  };
+
+}  // namespace pat
+
+void pat::PackedCandidateTrackChi2Producer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  auto const candidates = iEvent.getHandle(candidateToken_);
+
+  const edm::Association<reco::PFCandidateCollection>* candidate2PF = nullptr;
+  if (!doLostTracks_) {
+    candidate2PF = &iEvent.get(candidate2PFToken_);
+  }
+
+  const edm::Association<pat::PackedCandidateCollection>* tracks2LT = nullptr;
+  edm::Handle<reco::TrackCollection> trks;
+  if (doLostTracks_) {
+    tracks2LT = &iEvent.get(track2LostTrackToken_);
+    iEvent.getByToken(trackToken_, trks);
+  }
+
+  const auto nCand = candidates->size();
+  std::vector<float> trkChi2Map(nCand, 0);
+
+  if (doLostTracks_) {  //for Lost tracks we don't have references to PFCands, so we must loop over tracks and check keys...
+    for (size_t i = 0; i < trks->size(); i++) {
+      const auto& trk = reco::TrackRef(trks, i);
+      const auto& lostTrack = (*tracks2LT)[trk];
+      if (lostTrack.isNonnull()) {
+        const float nChi2 = trk->normalizedChi2();
+        trkChi2Map.at(lostTrack.key()) = MiniFloatConverter::reduceMantissaToNbitsRounding<roundingPrecision>(nChi2);
+      }
+    }
+  } else {  //for the regular PackedPFCands we have direct references...
+    for (size_t i = 0; i < nCand; i++) {
+      const auto& cand = pat::PackedCandidateRef(candidates, i);
+
+      // ignore neutral candidates or without track
+      if (cand->charge() == 0 || !cand->hasTrackDetails())
+        continue;
+
+      const auto& candTrack = (*candidate2PF)[cand]->trackRef();
+
+      if (candTrack.isNonnull()) {
+        const float nChi2 = candTrack->normalizedChi2();
+
+        trkChi2Map.at(i) = MiniFloatConverter::reduceMantissaToNbitsRounding<roundingPrecision>(nChi2);
+      }
+    }
+  }
+
+  // fill the value maps
+  std::unique_ptr<FloatMap> valueMap = std::make_unique<FloatMap>();
+  FloatMap::Filler filler(*valueMap);
+  filler.insert(candidates, trkChi2Map.begin(), trkChi2Map.end());
+  filler.fill();
+  iEvent.put(std::move(valueMap), "");
+}
+
+// ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
+void pat::PackedCandidateTrackChi2Producer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"))
+      ->setComment("packed candidate input collection");
+  desc.add<edm::InputTag>("trackCollection", edm::InputTag("generalTracks"))->setComment("track input collection");
+  desc.add<bool>("doLostTracks", false);
+  descriptions.add("packedPFCandidateTrackChi2", desc);
+}
+
+#include "FWCore/Framework/interface/MakerMacros.h"
+using namespace pat;
+DEFINE_FWK_MODULE(PackedCandidateTrackChi2Producer);

PhysicsTools/PatAlgos/python/slimming/MicroEventContent_cff.py

@@ -135,7 +135,9 @@
 
 _pp_on_AA_extraCommands = [
     'keep patPackedCandidates_hiPixelTracks_*_*',
-    'keep *_packedCandidateMuonID_*_*' 
+	  'keep *_packedCandidateMuonID_*_*',
+	  'keep floatedmValueMap_packedPFCandidateTrackChi2_*_*',
+	  'keep floatedmValueMap_lostTrackChi2_*_*'
 ]
 from Configuration.Eras.Modifier_pp_on_AA_2018_cff import pp_on_AA_2018
 from Configuration.Eras.Modifier_pp_on_PbPb_run3_cff import pp_on_PbPb_run3

PhysicsTools/PatAlgos/python/slimming/slimming_cff.py

@@ -67,7 +67,13 @@
 
 from Configuration.Eras.Modifier_pp_on_PbPb_run3_cff import pp_on_PbPb_run3
 from PhysicsTools.PatAlgos.packedCandidateMuonID_cfi import packedCandidateMuonID
-(pp_on_AA_2018 | pp_on_PbPb_run3).toReplaceWith(slimmingTask, cms.Task(slimmingTask.copy(), packedCandidateMuonID))
+from PhysicsTools.PatAlgos.packedPFCandidateTrackChi2_cfi import packedPFCandidateTrackChi2
+lostTrackChi2 = packedPFCandidateTrackChi2.clone(candidates = "lostTracks", doLostTracks = True)
+(pp_on_AA_2018 | pp_on_PbPb_run3).toReplaceWith(slimmingTask, 
+	cms.Task(slimmingTask.copy(), packedCandidateMuonID, 
+	packedPFCandidateTrackChi2, lostTrackChi2))
+
+
 
 from Configuration.Eras.Modifier_phase2_timing_cff import phase2_timing
 _phase2_timing_slimmingTask = cms.Task(slimmingTask.copy(),