Adding selector by size and sorter by pT

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

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