Added getters of related collections in ReconstructedParticle

HEP-FCC · Aug 23, 2024 · 0e352bf · 0e352bf
1 parent 2debc44
commit 0e352bf
Show file tree

Hide file tree

Showing 2 changed files with 109 additions and 0 deletions.
diff --git a/include/ral/ReconstructedParticle.h b/include/ral/ReconstructedParticle.h
@@ -4,7 +4,9 @@
 #include "Math/Vector3D.h"
 #include "Math/Vector4D.h"
 #include "ROOT/RVec.hxx"
+#include "edm4hep/ClusterData.h"
 #include "edm4hep/ReconstructedParticleData.h"
+#include "edm4hep/TrackData.h"
 #include "ral/LogicalOperators.h"
 #include <Math/GenVector/PxPyPzM4D.h>
 #include <Math/Vector4Dfwd.h>
@@ -225,6 +227,39 @@ get_m(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 ROOT::VecOps::RVec<float> get_goodnessOfPID(
     ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 
+/**
+ * Get all the daughter particles of one ReconstructedParticle
+ *
+ * @param main_particle Particle that is going to be look for daughters.
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
+get_daugthers(edm4hep::ReconstructedParticleData main_particle,
+              ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
+
+/**
+ * Get all the clusters related to one ReconstructedParticle
+ *
+ * @param main_particle Particle that is going to be look for clusters.
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+ROOT::VecOps::RVec<edm4hep::ClusterData>
+get_clusters(edm4hep::ReconstructedParticleData main_particle,
+             ROOT::VecOps::RVec<edm4hep::ClusterData> clusters);
+
+/**
+ * Get all the tracks related to one ReconstructedParticle
+ *
+ * @param main_particle Particle that is going to be look for tracks.
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+ROOT::VecOps::RVec<edm4hep::TrackData>
+get_tracks(edm4hep::ReconstructedParticleData main_particle,
+           ROOT::VecOps::RVec<edm4hep::TrackData> tracks);
+
 /**
  * Struct that can print the pdg of collection of ReconstructedParticles.
  */
@@ -517,6 +552,28 @@ ROOT::VecOps::RVec<bool>
 mask_abspdg(LogicalOperators::ComparisonOperator op, int value,
             ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 
+/**
+ * Select a subcollection of ReconstructedParticles
+ *
+ * @param n Number of elements that would be selected. Positive number selects
+ * from the begining and negative number select from the end
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> sel_n_elements(
+    int n, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
+
+/**
+ * Select one ReconstructedParticle
+ *
+ * @param n Index of the element selected
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+edm4hep::ReconstructedParticleData
+sel_element(int n,
+            ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
+
 /**
  * Select a subcollection of ReconstructedParticles based on the energy
  *

diff --git a/src/ReconstructedParticle.cc b/src/ReconstructedParticle.cc
@@ -5,7 +5,9 @@
 #include <Math/Vector4Dfwd.h>
 #include <ROOT/RVec.hxx>
 #include <cstdlib>
+#include <edm4hep/ClusterData.h>
 #include <edm4hep/ReconstructedParticleData.h>
+#include <edm4hep/TrackData.h>
 #include <iostream>
 #include <stdexcept>
 
@@ -255,6 +257,45 @@ ROOT::VecOps::RVec<float> get_goodnessOfPID(
   return result;
 }
 
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> get_daugthers(
+    edm4hep::ReconstructedParticleData main_particle,
+    ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
+  ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> result;
+  result.reserve(main_particle.particles_end - main_particle.particles_begin +
+                 1);
+  for (int i = main_particle.particles_begin; i <= main_particle.particles_end;
+       i++) {
+    result.emplace_back(particles[i]);
+  }
+  return result;
+}
+
+ROOT::VecOps::RVec<edm4hep::ClusterData>
+get_clusters(edm4hep::ReconstructedParticleData main_particle,
+             ROOT::VecOps::RVec<edm4hep::ClusterData> clusters) {
+  ROOT::VecOps::RVec<edm4hep::ClusterData> result;
+  result.reserve(main_particle.particles_end - main_particle.particles_begin +
+                 1);
+  for (int i = main_particle.particles_begin; i <= main_particle.particles_end;
+       i++) {
+    result.emplace_back(clusters[i]);
+  }
+  return result;
+}
+
+ROOT::VecOps::RVec<edm4hep::TrackData>
+get_tracks(edm4hep::ReconstructedParticleData main_particle,
+           ROOT::VecOps::RVec<edm4hep::TrackData> tracks) {
+  ROOT::VecOps::RVec<edm4hep::TrackData> result;
+  result.reserve(main_particle.particles_end - main_particle.particles_begin +
+                 1);
+  for (int i = main_particle.particles_begin; i <= main_particle.particles_end;
+       i++) {
+    result.emplace_back(tracks[i]);
+  }
+  return result;
+}
+
 print_pdg::print_pdg(int n_events) : m_n_events{n_events}, m_n_printed{0} {}
 
 int print_pdg::operator()(
@@ -571,6 +612,17 @@ mask_abspdg(LogicalOperators::ComparisonOperator op, int value,
   return result;
 }
 
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> sel_n_elements(
+    int n, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
+  return ROOT::VecOps::Take(particles, n);
+}
+
+edm4hep::ReconstructedParticleData
+sel_element(int n,
+            ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
+  return particles[n];
+}
+
 ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
 sel_e(LogicalOperators::ComparisonOperator op, float value,
       ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {