Skip to content

Commit 9be047e

Browse files
goblirscgoblirsc
committed
support index-based selection
1 parent 8a2684c commit 9be047e

7 files changed

Lines changed: 255 additions & 72 deletions

File tree

‎analyzers/dataframe/FCCAnalyses/MCParticle.h‎

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
#include <cmath>
66
#include <vector>
7+
#include <functional>
78

89
#include "ROOT/RVec.hxx"
910
#include "TLorentzVector.h"
@@ -30,26 +31,41 @@ namespace MCParticle{
3031
bool operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in);
3132
};
3233

34+
/// @brief Helper struct to select entries matching a certain predicate.
35+
/// Supports two signatures - either a list of candidates is passed and a list of accepted candidates returned,
36+
/// Or a list of indices in a vector of candidates is passed and a list of accepted indices returned.
37+
/// The latter is more compatible with index-based selection logic.
38+
struct selByPredicate{
39+
selByPredicate(std::function<bool(const edm4hep::MCParticleData &)> thePredicate):m_predicate(thePredicate){}
40+
std::function<bool(const edm4hep::MCParticleData &)> m_predicate;
41+
ROOT::VecOps::RVec<edm4hep::MCParticleData> operator() (const ROOT::VecOps::RVec<edm4hep::MCParticleData> & in);
42+
ROOT::VecOps::RVec<int> operator() (const ROOT::VecOps::RVec<int> & indices, const ROOT::VecOps::RVec<edm4hep::MCParticleData> & in);
43+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> operator() (const ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> & indices, const ROOT::VecOps::RVec<edm4hep::MCParticleData> & in);
44+
};
45+
3346
/// select MCParticles with transverse momentum greater than a minimum value [GeV]
34-
struct sel_pt {
47+
struct sel_pt : selByPredicate{
3548
sel_pt(float arg_min_pt);
36-
float m_min_pt = 20; //> transverse momentum threshold [GeV]
37-
ROOT::VecOps::RVec<edm4hep::MCParticleData> operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in);
49+
};
50+
51+
/// select MCParticles with absolute pseudorapidity less than a max value
52+
struct sel_eta : selByPredicate {
53+
sel_eta(float arg_max_eta);
3854
};
3955

4056
/// select MCParticles with their status
41-
struct sel_genStatus {
57+
struct sel_genStatus : selByPredicate {
4258
sel_genStatus(int arg_status);
43-
int m_status = 1; //> Generator status
44-
ROOT::VecOps::RVec<edm4hep::MCParticleData> operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in);
4559
};
4660

4761
/// select MCParticles with their PDG id
48-
struct sel_pdgID {
62+
struct sel_pdgID : selByPredicate {
4963
sel_pdgID(int arg_pdg, bool arg_chargeconjugate);
50-
int m_pdg = 13;
51-
bool m_chargeconjugate = true;
52-
ROOT::VecOps::RVec<edm4hep::MCParticleData> operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in);
64+
};
65+
66+
/// select MCParticles with a non-zero charge
67+
struct sel_charged : selByPredicate {
68+
sel_charged();
5369
};
5470

5571
/// get MC history tree for a given MCParticle index
@@ -116,7 +132,6 @@ namespace MCParticle{
116132
ROOT::VecOps::RVec<edm4hep::MCParticleData> in ,
117133
ROOT::VecOps::RVec<int> ind);
118134

119-
120135
/// return the parent index of a given list of MC particles
121136
ROOT::VecOps::RVec<int> get_parentid(ROOT::VecOps::RVec<int> mcind, ROOT::VecOps::RVec<edm4hep::MCParticleData> mc, ROOT::VecOps::RVec<int> parents);
122137

@@ -210,6 +225,11 @@ namespace MCParticle{
210225
/// return the list of stable particles from the decay of a mother particle, looking at the full decay chain recursively. i is the mother index in the Particle block
211226
std::vector<int> get_list_of_stable_particles_from_decay( int i, ROOT::VecOps::RVec<edm4hep::MCParticleData> in, ROOT::VecOps::RVec<int> ind) ;
212227

228+
/// return the list of stable particles from the decays of a mother particle, looking at the full decay chain recursively.
229+
/// i is the list of mother indices to process in the Particle block.
230+
/// Will return a vector of vectors - each vector is the set of children for one of the mothers in the input vector
231+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> get_lists_of_stable_particles_from_decays( ROOT::VecOps::RVec<int> i, ROOT::VecOps::RVec<edm4hep::MCParticleData> in, ROOT::VecOps::RVec<int> ind) ;
232+
213233
/// return the list of particles from the decay of a mother particle. i is the mother index in the Particle block.
214234
std::vector<int> get_list_of_particles_from_decay( int i,
215235
ROOT::VecOps::RVec<edm4hep::MCParticleData> in,

‎analyzers/dataframe/FCCAnalyses/ReconstructedParticle2MC.h‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,12 +125,35 @@ namespace ReconstructedParticle2MC{
125125
ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> reco,
126126
ROOT::VecOps::RVec<edm4hep::MCParticleData> mc) ;
127127

128-
/// select ReconstructedParticles matched to the (stable) MC particles whose indices are passed in a list
129-
ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> selRP_matched_to_list( ROOT::VecOps::RVec<int> mcParticles_indices,
130-
ROOT::VecOps::RVec<int> recind,
131-
ROOT::VecOps::RVec<int> mcind,
132-
ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> reco,
133-
ROOT::VecOps::RVec<edm4hep::MCParticleData> mc) ;
128+
/// select ReconstructedParticles matched to the MC particles whose indices are passed in a list
129+
/// @param mcParticles_indices indices of the MC particles to look up
130+
/// @param recind: reco index component of the MCRecoAssociations
131+
/// @param mcind: mc index component of the MCRecoAssociations
132+
/// @param reco: full reco particle list (ReconstructedParticles)
133+
/// @param mc: full mc particle list (Particles)
134+
/// @param require_stable: if set to true, will only match stable particles.
135+
/// @return List of ReconstructedParticle candidates with length corresponding to the number of *stable* MC particles in the mcParticles_indices vector. In presence of unstable particles, no 1:1 correspondence. For non-reconstructed stable MC particles, a dummy particle will be inserted. If 1:1 length correspondence is required, set require_stable to false.
136+
ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> selRP_matched_to_list(
137+
const ROOT::VecOps::RVec<int> & mcParticles_indices,
138+
const ROOT::VecOps::RVec<int> & recind,
139+
const ROOT::VecOps::RVec<int> & mcind,
140+
const ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> & reco,
141+
const ROOT::VecOps::RVec<edm4hep::MCParticleData> & mc,
142+
bool require_stable = true) ;
143+
/// select indices of ReconstructedParticles matched to the MC particles whose indices are passed in a list
144+
/// @param mcParticles_indices indices of the MC particles to look up
145+
/// @param recind: reco index component of the MCRecoAssociations
146+
/// @param mcind: mc index component of the MCRecoAssociations
147+
/// @param mc: full mc particle list (Particles)
148+
/// @param require_stable: if set to true, will only match stable particles.
149+
/// @return List of ReconstructedParticle candidates with length corresponding to the number of *stable* MC particles in the mcParticles_indices vector. In presence of unstable particles, no 1:1 correspondence. For non-reconstructed stable MC particles, a "-1" entry will be inserted. If 1:1 length correspondence is required, set require_stable to false.
150+
ROOT::VecOps::RVec<int> selRP_indices_matched_to_list(
151+
const ROOT::VecOps::RVec<int> & mcParticles_indices,
152+
const ROOT::VecOps::RVec<int> & recind,
153+
const ROOT::VecOps::RVec<int> & mcind,
154+
const ROOT::VecOps::RVec<edm4hep::MCParticleData> & mc,
155+
bool require_stable = true
156+
) ;
134157

135158
/// return the index of the MC particle that is associated to a given track (via the track-reco association)
136159
int getTrack2MC_index ( int track_index,

‎analyzers/dataframe/FCCAnalyses/Utils.h‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,40 @@
22
#define UTILS_ANALYZERS_H
33

44
#include <cmath>
5+
#include <algorithm>
56

67
namespace FCCAnalyses {
78
namespace Utils {
89

910
template<typename T> inline auto getsize( T& vec){ return vec.size();};
10-
template<typename T> inline ROOT::VecOps::RVec<ROOT::VecOps::RVec<T> > as_vector(const ROOT::VecOps::RVec<T>& in){return ROOT::VecOps::RVec<ROOT::VecOps::RVec<T> >(1, in);};
11+
template<typename T> inline ROOT::VecOps::RVec<ROOT::VecOps::RVec<T> > as_vector(const ROOT::VecOps::RVec<T>& in){return ROOT::VecOps::RVec<ROOT::VecOps::RVec<T> >(1, in);};
12+
template <typename T> inline ROOT::VecOps::RVec<int> index_range(const ROOT::VecOps::RVec<T> & in){
13+
ROOT::VecOps::RVec<int> indices(in.size());
14+
std::iota(indices.begin(),indices.end(), 0);
15+
return indices;
16+
};
17+
/// @brief count the number of valid (>=0, < size of collection) indices in an index list
18+
/// @param in: index list
19+
/// @param ref: particle vector to which the indices refer.
20+
/// @return integer count of valid indices
21+
template <typename T> inline int count_valid_indices(const ROOT::VecOps::RVec<int> & in,
22+
const ROOT::VecOps::RVec<T> & ref){
23+
int maxSize = ref.size();
24+
return std::count_if(in.begin(),in.end(),[maxSize](const int & i){return (i >=0 && i < maxSize);});
25+
};
26+
27+
// @brief for a given list of indices, returns a list of particle (copies)
28+
/// @param idx : Indices of desired particles within the full set ("in")
29+
/// @param in : Full set of particles
30+
/// @return A vector of particles with the desired indices.
31+
template <typename T> inline ROOT::VecOps::RVec<T> sel_byIndex( const ROOT::VecOps::RVec<int> & idx, const ROOT::VecOps::RVec<T> & in){
32+
ROOT::VecOps::RVec<T> found;
33+
for (int index : idx){
34+
if (index < 0 || index >= in.size()) continue;
35+
found.push_back(in.at(index));
36+
}
37+
return found;
38+
}
1139
}
1240
}
1341

‎analyzers/dataframe/FCCAnalyses/VertexingUtils.h‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,22 @@ namespace VertexingUtils{
107107
/// Retrieve the indices of the tracks fitted to that vertex, but now in the collection of RecoParticles
108108
ROOT::VecOps::RVec<int> get_VertexRecoParticlesInd( FCCAnalysesVertex TheVertex,
109109
const ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>& reco );
110-
110+
111+
/// Retrieve the indices of the tracks fitted to a vector of vertices, but now in the collection of RecoParticles
112+
ROOT::VecOps::RVec<int> get_VerticesRecoParticlesInd( ROOT::VecOps::RVec<FCCAnalysesVertex > vertices,
113+
const ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>& reco );
114+
115+
/// @brief Find (by index) a vertex that is made of reco particles from a passed list.
116+
/// @param vertices: List of possible vertices to match
117+
/// @param recoParticleIndices: Indices of the reco particles to find in the vertex
118+
/// @param reco: list of all reco particles (ReconstructedParticles)
119+
/// @param require_all: If set, require one vertex to contain all recoParticleIndices (no 'missed' tracks). Else only require that all tracks in the vertex come from the vector (but allow for unused tracks)
120+
/// @return: Index within the list of the (first) vertex fulfilling the criteria. If none are found, return -1
121+
int getVertex_matching_recoParticles(const ROOT::VecOps::RVec<FCCAnalysesVertex > & vertices,
122+
const ROOT::VecOps::RVec<int> & recoParticleIndices,
123+
const ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> &reco,
124+
bool require_all = false);
125+
111126
/// Return the number of tracks in a given track collection
112127
int get_nTracks(ROOT::VecOps::RVec<edm4hep::TrackState> tracks);
113128

‎analyzers/dataframe/src/MCParticle.cc‎

Lines changed: 63 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -8,35 +8,66 @@ namespace FCCAnalyses{
88

99
namespace MCParticle{
1010

11-
sel_genStatus::sel_genStatus(int arg_status) : m_status(arg_status) {};
12-
ROOT::VecOps::RVec<edm4hep::MCParticleData> sel_genStatus::operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in) {
11+
12+
ROOT::VecOps::RVec<edm4hep::MCParticleData> selByPredicate::operator() (const ROOT::VecOps::RVec<edm4hep::MCParticleData> & in){
1313
ROOT::VecOps::RVec<edm4hep::MCParticleData> result;
1414
result.reserve(in.size());
15-
for (size_t i = 0; i < in.size(); ++i) {
16-
auto & p = in[i];
17-
if (p.generatorStatus == m_status) {
18-
result.emplace_back(p);
19-
}
20-
}
21-
return result;
15+
for (auto & p : in) {
16+
if (m_predicate(p)) result.emplace_back(p);
17+
}
18+
return result;
2219
}
2320

24-
sel_pdgID::sel_pdgID(int arg_pdg, bool arg_chargeconjugate) : m_pdg(arg_pdg), m_chargeconjugate( arg_chargeconjugate ) {};
25-
ROOT::VecOps::RVec<edm4hep::MCParticleData> sel_pdgID::operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in) {
26-
ROOT::VecOps::RVec<edm4hep::MCParticleData> result;
21+
ROOT::VecOps::RVec<int> selByPredicate::operator() (const ROOT::VecOps::RVec<int>& indices, const ROOT::VecOps::RVec<edm4hep::MCParticleData> &in){
22+
ROOT::VecOps::RVec<int> result;
2723
result.reserve(in.size());
28-
for (size_t i = 0; i < in.size(); ++i) {
29-
auto & p = in[i];
30-
if ( m_chargeconjugate ) {
31-
if ( std::abs( p.PDG ) == std::abs( m_pdg) ) result.emplace_back(p);
32-
}
33-
else {
34-
if ( p.PDG == m_pdg ) result.emplace_back(p);
35-
}
24+
for (int index : indices) {
25+
if (index < 0 || index >= in.size()) continue;
26+
if (m_predicate(in[index])) result.emplace_back(index);
27+
}
28+
return result;
29+
}
30+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> selByPredicate::operator() (const ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>>& setsOfIndices, const ROOT::VecOps::RVec<edm4hep::MCParticleData> &in){
31+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> result(setsOfIndices.size());
32+
for (int elem = 0; elem < setsOfIndices.size(); ++elem){
33+
result[elem] = this->operator()(setsOfIndices[elem], in);
3634
}
37-
return result;
35+
return result;
36+
}
37+
38+
sel_pt::sel_pt(float arg_min_pt) :
39+
selByPredicate([arg_min_pt](const edm4hep::MCParticleData & p)->bool{
40+
return (p.momentum.x*p.momentum.x + p.momentum.y*p.momentum.y > arg_min_pt*arg_min_pt);
41+
}){
42+
}
43+
44+
sel_eta::sel_eta(float arg_max_eta) :
45+
selByPredicate([arg_max_eta](const edm4hep::MCParticleData & p)->bool{
46+
ROOT::Math::PxPyPzM4D vec(p.momentum.x, p.momentum.y, p.momentum.z, p.mass);
47+
return std::abs(vec.Eta()) < std::abs(arg_max_eta);
48+
}){
3849
}
3950

51+
sel_genStatus::sel_genStatus(int arg_status):
52+
selByPredicate([arg_status](const edm4hep::MCParticleData & p)->bool{
53+
return p.generatorStatus == arg_status;
54+
}){
55+
}
56+
57+
sel_pdgID::sel_pdgID(int arg_pdg, bool arg_chargeconjugate):
58+
selByPredicate([arg_pdg, arg_chargeconjugate](const edm4hep::MCParticleData & p)->bool{
59+
if (arg_chargeconjugate) return std::abs( p.PDG ) == std::abs( arg_pdg);
60+
else return p.PDG == arg_pdg;
61+
}){
62+
}
63+
64+
sel_charged::sel_charged():
65+
selByPredicate([](const edm4hep::MCParticleData & p)->bool{
66+
// try to avoid floating comp to zero
67+
// and thank you for the person who gave some neutral particles a -999 charge!
68+
return std::abs(p.charge) > 1e-6 && p.charge != -999;
69+
}){
70+
}
4071

4172

4273
get_decay::get_decay(int arg_mother, int arg_daughters, bool arg_inf){m_mother=arg_mother; m_daughters=arg_daughters; m_inf=arg_inf;};
@@ -59,19 +90,6 @@ bool get_decay::operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in, ROO
5990
return result;
6091
}
6192

62-
sel_pt::sel_pt(float arg_min_pt) : m_min_pt(arg_min_pt) {};
63-
ROOT::VecOps::RVec<edm4hep::MCParticleData> sel_pt::operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in) {
64-
ROOT::VecOps::RVec<edm4hep::MCParticleData> result;
65-
result.reserve(in.size());
66-
for (size_t i = 0; i < in.size(); ++i) {
67-
auto & p = in[i];
68-
if (std::sqrt(std::pow(p.momentum.x,2) + std::pow(p.momentum.y,2)) > m_min_pt) {
69-
result.emplace_back(p);
70-
}
71-
}
72-
return result;
73-
}
74-
7593

7694
filter_pdgID::filter_pdgID(int arg_pdgid, bool arg_abs){m_pdgid = arg_pdgid; m_abs = arg_abs;};
7795
bool filter_pdgID::operator() (ROOT::VecOps::RVec<edm4hep::MCParticleData> in) {
@@ -540,6 +558,16 @@ std::vector<int> get_list_of_particles_from_decay(int i, ROOT::VecOps::RVec<edm4
540558
return res;
541559
}
542560

561+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> get_lists_of_stable_particles_from_decays(ROOT::VecOps::RVec<int> i, ROOT::VecOps::RVec<edm4hep::MCParticleData> in, ROOT::VecOps::RVec<int> ind) {
562+
563+
ROOT::VecOps::RVec<ROOT::VecOps::RVec<int>> ret;
564+
ret.reserve(i.size());
565+
for (int ix : i){
566+
ret.push_back(get_list_of_stable_particles_from_decay(ix, in, ind));
567+
}
568+
return ret;
569+
}
570+
543571

544572
// ----------------------------------------------------------------------------------------------------------------------------------
545573

@@ -560,6 +588,7 @@ std::vector<int> list_of_particles_from_decay(int i, ROOT::VecOps::RVec<edm4hep:
560588

561589

562590

591+
563592
// ----------------------------------------------------------------------------------------------------------------------------------
564593

565594
ROOT::VecOps::RVec<int> get_indices_MotherByIndex ( int imother,

0 commit comments

Comments
 (0)