From 6b0858f01fd6586353d74be9588ebe16bfaeef7b Mon Sep 17 00:00:00 2001
From: Andre Govinda Stahl Leiton <andre.govinda.stahl.leiton@cern.ch>
Date: Fri, 30 Aug 2024 22:52:47 +0200
Subject: [PATCH] Add option to sort UParT input candidates by pT

---
 ...edParticleTransformerAK4TagInfoProducer.cc | 36 ++++++++++++++-----
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc b/RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc
index bcd6fb22cf543..29bbfe5794c12 100644
--- a/RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc
+++ b/RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc
@@ -79,6 +79,7 @@ class UnifiedParticleTransformerAK4TagInfoProducer : public edm::stream::EDProdu
   const double jet_radius_;
   const double min_candidate_pt_;
   const bool flip_;
+  const bool sort_cand_by_pt_;
 
   const edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
   const edm::EDGetTokenT<VertexCollection> vtx_token_;
@@ -108,6 +109,7 @@ UnifiedParticleTransformerAK4TagInfoProducer::UnifiedParticleTransformerAK4TagIn
     : jet_radius_(iConfig.getParameter<double>("jet_radius")),
       min_candidate_pt_(iConfig.getParameter<double>("min_candidate_pt")),
       flip_(iConfig.getParameter<bool>("flip")),
+      sort_cand_by_pt_(iConfig.getParameter<bool>("sort_cand_by_pt")),
       jet_token_(consumes<edm::View<reco::Jet>>(iConfig.getParameter<edm::InputTag>("jets"))),
       vtx_token_(consumes<VertexCollection>(iConfig.getParameter<edm::InputTag>("vertices"))),
       lt_token_(consumes<edm::View<reco::Candidate>>(iConfig.getParameter<edm::InputTag>("losttracks"))),
@@ -154,6 +156,7 @@ void UnifiedParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::Configu
   desc.add<double>("jet_radius", 0.4);
   desc.add<double>("min_candidate_pt", 0.10);
   desc.add<bool>("flip", false);
+  desc.add<bool>("sort_cand_by_pt", false);
   desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
   desc.add<edm::InputTag>("losttracks", edm::InputTag("lostTracks"));
   desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
@@ -275,10 +278,16 @@ void UnifiedParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, c
             auto& trackinfo = lt_trackinfos.emplace(i, track_builder).first->second;
             trackinfo.buildTrackInfo(PackedCandidate_, jet_dir, jet_ref_track_dir, pv);
 
-            lt_sorted.emplace_back(i,
-                                   trackinfo.getTrackSip2dSig(),
-                                   -btagbtvdeep::mindrsvpfcand(svs_unsorted, PackedCandidate_),
-                                   PackedCandidate_->pt() / jet.pt());
+            if (sort_cand_by_pt_)
+              lt_sorted.emplace_back(i,
+                                     PackedCandidate_->pt() / jet.pt(),
+                                     trackinfo.getTrackSip2dSig(),
+                                     -btagbtvdeep::mindrsvpfcand(svs_unsorted, PackedCandidate_));
+            else
+              lt_sorted.emplace_back(i,
+                                     trackinfo.getTrackSip2dSig(),
+                                     -btagbtvdeep::mindrsvpfcand(svs_unsorted, PackedCandidate_),
+                                     PackedCandidate_->pt() / jet.pt());
 
             ltPtrs.push_back(cand);
           }
@@ -353,12 +362,21 @@ void UnifiedParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, c
             auto& trackinfo = trackinfos.emplace(i, track_builder).first->second;
             trackinfo.buildTrackInfo(cand, jet_dir, jet_ref_track_dir, pv);
 
-            c_sorted.emplace_back(i,
-                                  trackinfo.getTrackSip2dSig(),
-                                  -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand),
-                                  cand->pt() / jet.pt());
+            if (sort_cand_by_pt_)
+              c_sorted.emplace_back(i,
+                                    cand->pt() / jet.pt(),
+                                    trackinfo.getTrackSip2dSig(),
+                                    -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand));
+            else
+              c_sorted.emplace_back(i,
+                                    trackinfo.getTrackSip2dSig(),
+                                    -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand),
+                                    cand->pt() / jet.pt());
           } else {
-            n_sorted.emplace_back(i, -1, -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand), cand->pt() / jet.pt());
+            if (sort_cand_by_pt_)
+              n_sorted.emplace_back(i, cand->pt() / jet.pt(), -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand), -1);
+            else
+              n_sorted.emplace_back(i, -1, -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand), cand->pt() / jet.pt());
           }
         }
       }