From 81db22c9d21778bf669fafbaea45aa1642fd957e Mon Sep 17 00:00:00 2001
From: Juraj Smiesko <juraj.smiesko@cern.ch>
Date: Wed, 15 Nov 2023 13:36:51 +0100
Subject: [PATCH] Adding selector by size and sorter by pT

---
 .../FCCAnalyses/ReconstructedParticle.h       | 17 +++++++
 .../dataframe/src/ReconstructedParticle.cc    | 44 +++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h b/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h
index 183f7d164db..9d2c1d9630a 100644
--- a/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h
+++ b/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h
@@ -121,6 +121,19 @@ namespace ReconstructedParticle{
   };
 
 
+  /**
+   * \brief Analyzer to select specified number of reconstructed particles
+   *
+   * \param size  Desired number of particles
+   */
+  struct selUpTo {
+    selUpTo(const size_t size);
+    const size_t m_size;
+    edm4hep::ReconstructedParticleCollection operator() (
+        const edm4hep::ReconstructedParticleCollection& inColl);
+  };
+
+
   /// return reconstructed particles
   ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> get(ROOT::VecOps::RVec<int> index, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
 
@@ -194,6 +207,10 @@ namespace ReconstructedParticle{
   /// get number of b-jets
   int getJet_ntags(ROOT::VecOps::RVec<bool> in);
 
+  /// reorder input ReconstructedParticles by pT
+  edm4hep::ReconstructedParticleCollection
+  sortByPt(const edm4hep::ReconstructedParticleCollection& inParticles);
+
 }//end NS ReconstructedParticle
 
 }//end NS FCCAnalyses
diff --git a/analyzers/dataframe/src/ReconstructedParticle.cc b/analyzers/dataframe/src/ReconstructedParticle.cc
index 838601a475c..8e774002952 100644
--- a/analyzers/dataframe/src/ReconstructedParticle.cc
+++ b/analyzers/dataframe/src/ReconstructedParticle.cc
@@ -131,6 +131,22 @@ edm4hep::ReconstructedParticleCollection selPDG::operator() (
 }
 
 
+selUpTo::selUpTo(const size_t size) : m_size(size) {};
+
+edm4hep::ReconstructedParticleCollection selUpTo::operator() (
+    const edm4hep::ReconstructedParticleCollection& inColl) {
+  edm4hep::ReconstructedParticleCollection result;
+  result.setSubsetCollection();
+
+  for (const auto& particle: inColl) {
+    if (result.size() >= m_size) {
+      break;
+    }
+    result.push_back(particle);
+  }
+
+  return result;
+}
 
 
 resonanceBuilder::resonanceBuilder(float arg_resonance_mass) {m_resonance_mass = arg_resonance_mass;}
@@ -501,6 +517,34 @@ int getJet_ntags(ROOT::VecOps::RVec<bool> in) {
   return result;
 }
 
+
+edm4hep::ReconstructedParticleCollection
+sortByPt(const edm4hep::ReconstructedParticleCollection& inColl) {
+  edm4hep::ReconstructedParticleCollection outColl;
+  outColl.setSubsetCollection();
+
+  std::vector<edm4hep::ReconstructedParticle> rpVec;
+  for (const auto& particle: inColl) {
+    rpVec.emplace_back(particle);
+  }
+
+  auto pt = [] (const auto& particle) {
+    return std::sqrt(std::pow(particle.getMomentum().x, 2) +
+                     std::pow(particle.getMomentum().y, 2));
+  };
+
+  std::sort(rpVec.begin(), rpVec.end(),
+            [&pt](const auto& lhs, const auto& rhs) {
+              return pt(lhs) > pt(rhs);
+            });
+
+  for (const auto& particle: rpVec) {
+    outColl.push_back(particle);
+  }
+
+  return outColl;
+}
+
 }//end NS ReconstructedParticle
 
 }//end NS FCCAnalyses