Merge pull request #118 from EmanuelPerez/RecoUpdates_from_GM

Some updates to ReconstructedParticle from Giovanni Marchiori: additi…
HEP-FCC · Feb 15, 2022 · bd4b0c4 · bd4b0c4
2 parents ab39a05 + a3d9ddf
commit bd4b0c4
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 3 deletions.
diff --git a/analyzers/dataframe/ReconstructedParticle.cc b/analyzers/dataframe/ReconstructedParticle.cc
@@ -14,13 +14,16 @@ ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>  ReconstructedParticle::s
   return result;
 }
 
-ReconstructedParticle::sel_p::sel_p(float arg_min_p) : m_min_p(arg_min_p) {};
+ReconstructedParticle::sel_p::sel_p(float arg_min_p, float arg_max_p) : m_min_p(arg_min_p), m_max_p(arg_max_p)  {};
 ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>  ReconstructedParticle::sel_p::operator() (ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in) {
   ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> result;
   result.reserve(in.size());
   for (size_t i = 0; i < in.size(); ++i) {
     auto & p = in[i];
-    if (std::sqrt(std::pow(p.momentum.x,2) + std::pow(p.momentum.y,2) + std::pow(p.momentum.z,2) ) > m_min_p) {
+    float momentum = std::sqrt(   std::pow(p.momentum.x,2)
+                                + std::pow(p.momentum.y,2)
+                                + std::pow(p.momentum.z,2) );
+    if ( momentum > m_min_p && momentum < m_max_p ) {
       result.emplace_back(p);
     }
   }
@@ -111,6 +114,22 @@ ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> ReconstructedParticle::se
 }
 
 
+ReconstructedParticle::sel_tag::sel_tag(bool arg_pass): m_pass(arg_pass) {};
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> ReconstructedParticle::sel_tag::operator()(ROOT::VecOps::RVec<bool> tags, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in){
+  ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> result;
+  for (size_t i = 0; i < in.size(); ++i) {
+    if (m_pass) {
+      if (tags.at(i)) result.push_back(in.at(i));
+    }
+    else {
+      if (!tags.at(i)) result.push_back(in.at(i));
+    }
+  }
+  return result;
+}
+
+
+
 // Angular separation between the particles of a collection:
 //   arg_delta = 0 / 1 / 2 :   return delta_max, delta_min, delta_average
 
@@ -165,6 +184,41 @@ ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> ReconstructedParticle::me
   return ROOT::VecOps::RVec(result);
 }
 
+
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> ReconstructedParticle::remove(
+  		ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> x,
+  		ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> y) {
+  //to be kept as ROOT::VecOps::RVec
+  std::vector<edm4hep::ReconstructedParticleData> result;
+  result.reserve( x.size() );
+  result.insert( result.end(), x.begin(), x.end() );
+  float epsilon = 1e-8;
+  for (size_t i = 0; i < y.size(); ++i) {
+    float mass1 = y.at(i).mass;
+    float px1 = y.at(i).momentum.x;
+    float py1 = y.at(i).momentum.y;
+    float pz1 = y.at(i).momentum.z;
+    for(std::vector<edm4hep::ReconstructedParticleData>::iterator
+          it = std::begin(result); it != std::end(result); ++it) {
+      float mass2 = it->mass;
+      float px2 = it->momentum.x;
+      float py2 = it->momentum.y;
+      float pz2 = it->momentum.z;
+      if ( abs(mass1-mass2) < epsilon && 
+	   abs(px1-px2) < epsilon &&
+	   abs(py1-py2) < epsilon &&
+	   abs(pz1-pz2) < epsilon ) {
+        result.erase(it);
+        //break;
+      }
+    }
+  }
+  return ROOT::VecOps::RVec(result);
+}
+
+
+
+
 ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> ReconstructedParticle::get(ROOT::VecOps::RVec<int> index, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in){
   ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> result;
   for (size_t i = 0; i < index.size(); ++i) {

diff --git a/analyzers/dataframe/ReconstructedParticle.h b/analyzers/dataframe/ReconstructedParticle.h
@@ -42,8 +42,9 @@ namespace ReconstructedParticle{
 
   /// select ReconstructedParticles with momentum greater than a minimum value [GeV]
   struct sel_p {
-    sel_p(float arg_min_p);
+    sel_p(float arg_min_p, float arg_max_p = 1e10);
     float m_min_p = 1.; //> momentum threshold [GeV]
+    float m_max_p = 1e10; //< momentum threshold [GeV]
     ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>  operator() (ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
   };
 
@@ -62,6 +63,13 @@ namespace ReconstructedParticle{
     ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> operator()(ROOT::VecOps::RVec<float> angle, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
   };
 
+  /// select a list of reconstructed particles depending on the status of a certain boolean flag
+  struct sel_tag {
+    bool m_pass; // if pass is true, select tagged jets. Otherwise select anti-tagged ones
+    sel_tag(bool arg_pass);
+    ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>  operator() (ROOT::VecOps::RVec<bool> tags, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
+  };
+
 
 
 
@@ -122,6 +130,9 @@ namespace ReconstructedParticle{
   /// concatenate both input vectors and return the resulting vector
   ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> merge(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> x, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> y);
 
+  /// remove elements of vector y from vector x
+  ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> remove( ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> x, ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> y);
+
   /// return the size of the input collection
   int get_n(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);