Merge pull request #312 from kjvbrt/sel_type

Implementing selection by type for Reconstructed particles
HEP-FCC · Sep 19, 2023 · b969668 · b969668
2 parents f1078cd + 813303d
commit b969668
Show file tree

Hide file tree

Showing 4 changed files with 114 additions and 1 deletion.
diff --git a/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h b/analyzers/dataframe/FCCAnalyses/ReconstructedParticle.h
@@ -35,6 +35,24 @@ namespace ReconstructedParticle{
     float operator() (ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in) ;
   };
 
+  /// select ReconstructedParticles by type
+  /// Note: type might not correspond to PDG ID
+  struct sel_type {
+    sel_type(const int type);
+    const int m_type;
+    ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
+    operator()(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
+  };
+
+  /// select ReconstructedParticles by type absolute value
+  /// Note: type might not correspond to PDG ID
+  struct sel_absType {
+    sel_absType(const int type);
+    const int m_type;
+    ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
+    operator()(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in);
+  };
+
   /// select ReconstructedParticles with transverse momentum greater than a minimum value [GeV]
   struct sel_pt {
     sel_pt(float arg_min_pt);

diff --git a/analyzers/dataframe/src/ReconstructedParticle.cc b/analyzers/dataframe/src/ReconstructedParticle.cc
@@ -1,10 +1,50 @@
 #include "FCCAnalyses/ReconstructedParticle.h"
+
+// std
+#include <cstdlib>
 #include <iostream>
+#include <stdexcept>
 
 namespace FCCAnalyses{
 
 namespace ReconstructedParticle{
 
+/// sel_type
+sel_type::sel_type(const int type) : m_type(type) {}
+
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> sel_type::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) {
+    if (in[i].type == m_type) {
+      result.emplace_back(in[i]);
+    }
+  }
+  return result;
+}
+
+/// sel_absType
+sel_absType::sel_absType(const int type) : m_type(type) {
+  if (m_type < 0) {
+    throw std::invalid_argument(
+        "ReconstructedParticle::sel_absType: Received negative value!");
+  }
+}
+
+ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> sel_absType::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) {
+    if (std::abs(in[i].type) == m_type) {
+      result.emplace_back(in[i]);
+    }
+  }
+  return result;
+}
+
+/// sel_pt
 sel_pt::sel_pt(float arg_min_pt) : m_min_pt(arg_min_pt) {};
 ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>  sel_pt::operator() (ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> in) {
   ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> result;

diff --git a/tests/unittest/CMakeLists.txt b/tests/unittest/CMakeLists.txt
@@ -4,7 +4,11 @@ find_catch_instance()
 # list of labels that we want to ignore
 set(filter_tests "")
 
-add_executable(unittest unittest.cpp myutils.cpp algorithms.cpp)
+add_executable(unittest unittest.cpp
+                        myutils.cpp
+                        algorithms.cpp
+                        ReconstructedParticle.cpp
+)
 target_link_libraries(unittest PUBLIC FCCAnalyses gfortran PRIVATE Catch2::Catch2WithMain)
 target_include_directories(unittest PUBLIC ${VDT_INCLUDE_DIR})
 target_compile_definitions(unittest PUBLIC "-DTEST_INPUT_DATA_DIR=${TEST_INPUT_DATA_DIR}")

diff --git a/tests/unittest/ReconstructedParticle.cpp b/tests/unittest/ReconstructedParticle.cpp
@@ -0,0 +1,51 @@
+#include "FCCAnalyses/ReconstructedParticle.h"
+
+// Catch2
+#include "catch2/catch_test_macros.hpp"
+#include <catch2/catch_approx.hpp>
+
+TEST_CASE("sel_type", "[ReconstructedParticle]") {
+  ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> pVec;
+  edm4hep::ReconstructedParticleData p1;
+  p1.type = 11;
+  pVec.push_back(p1);
+  edm4hep::ReconstructedParticleData p2;
+  p2.type = 13;
+  pVec.push_back(p2);
+  edm4hep::ReconstructedParticleData p3;
+  p3.type = -11;
+  pVec.push_back(p3);
+  edm4hep::ReconstructedParticleData p4;
+  p4.type = -13;
+  pVec.push_back(p4);
+  FCCAnalyses::ReconstructedParticle::sel_type selType{11};
+  auto res = selType(pVec);
+  REQUIRE(res.size() == 1);
+  REQUIRE(res[0].type == 11);
+}
+
+TEST_CASE("sel_absType", "[ReconstructedParticle]") {
+  ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> pVec;
+  edm4hep::ReconstructedParticleData p1;
+  p1.type = 11;
+  pVec.push_back(p1);
+  edm4hep::ReconstructedParticleData p2;
+  p2.type = 13;
+  pVec.push_back(p2);
+  edm4hep::ReconstructedParticleData p3;
+  p3.type = -11;
+  pVec.push_back(p3);
+  edm4hep::ReconstructedParticleData p4;
+  p4.type = -13;
+  pVec.push_back(p4);
+  FCCAnalyses::ReconstructedParticle::sel_absType selAbsType{11};
+  auto res = selAbsType(pVec);
+  REQUIRE(res.size() == 2);
+  REQUIRE(res[0].type == 11);
+  REQUIRE(res[1].type == -11);
+}
+
+TEST_CASE("sel_absType__neg_type", "[ReconstructedParticle]") {
+  REQUIRE_THROWS_AS(FCCAnalyses::ReconstructedParticle::sel_absType(-17),
+                    std::invalid_argument);
+}