diff --git a/include/ral/ReconstructedParticle.h b/include/ral/ReconstructedParticle.h index ec302e7..c3bfec1 100644 --- 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 #include @@ -225,6 +227,39 @@ get_m(ROOT::VecOps::RVec particles); ROOT::VecOps::RVec get_goodnessOfPID( ROOT::VecOps::RVec 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 +get_daugthers(edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec 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 +get_clusters(edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec 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 +get_tracks(edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec tracks); + /** * Struct that can print the pdg of collection of ReconstructedParticles. */ @@ -517,6 +552,28 @@ ROOT::VecOps::RVec mask_abspdg(LogicalOperators::ComparisonOperator op, int value, ROOT::VecOps::RVec 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 sel_n_elements( + int n, ROOT::VecOps::RVec 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 particles); + /** * Select a subcollection of ReconstructedParticles based on the energy * diff --git a/src/ReconstructedParticle.cc b/src/ReconstructedParticle.cc index 3359d24..4b14195 100644 --- a/src/ReconstructedParticle.cc +++ b/src/ReconstructedParticle.cc @@ -5,7 +5,9 @@ #include #include #include +#include #include +#include #include #include @@ -255,6 +257,45 @@ ROOT::VecOps::RVec get_goodnessOfPID( return result; } +ROOT::VecOps::RVec get_daugthers( + edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec particles) { + ROOT::VecOps::RVec 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 +get_clusters(edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec clusters) { + ROOT::VecOps::RVec 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 +get_tracks(edm4hep::ReconstructedParticleData main_particle, + ROOT::VecOps::RVec tracks) { + ROOT::VecOps::RVec 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 sel_n_elements( + int n, ROOT::VecOps::RVec particles) { + return ROOT::VecOps::Take(particles, n); +} + +edm4hep::ReconstructedParticleData +sel_element(int n, + ROOT::VecOps::RVec particles) { + return particles[n]; +} + ROOT::VecOps::RVec sel_e(LogicalOperators::ComparisonOperator op, float value, ROOT::VecOps::RVec particles) {