Adding support for PODIO DataSource

include/ral/LogicalOperators.h

@@ -1,17 +1,59 @@
 #ifndef LOGICAL_OPERATORS_RAL_H
 #define LOGICAL_OPERATORS_RAL_H
 
+// std
+#include <type_traits>
+
+// ROOT
 #include "ROOT/RVec.hxx"
 
+// PODIO
+#include "podio/CollectionBase.h"
+
 namespace k4::ral {
 
 namespace LogicalOperators {
 
 enum class ComparisonOperator { LESS, LESSEQ, EQ, GREATEREQ, GREATER };
 
 template <typename T>
-ROOT::VecOps::RVec<T> filter(ROOT::VecOps::RVec<bool> mask,
-                             ROOT::VecOps::RVec<T> collection);
+ROOT::VecOps::RVec<T> filter(const ROOT::VecOps::RVec<bool> &mask,
+                             const ROOT::VecOps::RVec<T> &collection) {
+  if (mask.size() != collection.size()) {
+    auto msg = "Different vector lengths: " + std::to_string(mask.size()) +
+               " vs. " + std::to_string(collection.size()) + "!";
+    throw std::length_error(msg);
+  }
+
+  ROOT::VecOps::RVec<T> result;
+  for (int i = 0; i < collection.size(); i++) {
+    if (mask[i]) {
+      result.emplace_back(collection[i]);
+    }
+  }
+
+  return result;
+}
+
+template <typename C, typename = std::enable_if<
+                          std::is_base_of<podio::CollectionBase, C>::value>>
+C filter(const ROOT::VecOps::RVec<bool> &mask, const C &collection) {
+  if (mask.size() != collection.size()) {
+    auto msg = "Different vector lengths: " + std::to_string(mask.size()) +
+               " vs. " + std::to_string(collection.size()) + "!";
+    throw std::length_error(msg);
+  }
+
+  C result;
+  result.setSubsetCollection();
+  for (int i = 0; i < collection.size(); i++) {
+    if (mask[i]) {
+      result.push_back(collection[i]);
+    }
+  }
+
+  return result;
+}
 
 ROOT::VecOps::RVec<bool> operator&&(const ROOT::VecOps::RVec<bool> &vec1,
                                     const ROOT::VecOps::RVec<bool> &vec2);

include/ral/ReconstructedParticle.h

@@ -5,7 +5,7 @@
 #include "Math/Vector4D.h"
 #include "ROOT/RVec.hxx"
 #include "edm4hep/ClusterData.h"
-#include "edm4hep/ReconstructedParticleData.h"
+#include "edm4hep/ReconstructedParticleCollection.h"
 #include "edm4hep/TrackData.h"
 #include "ral/LogicalOperators.h"
 #include <Math/GenVector/PxPyPzM4D.h>
@@ -218,6 +218,15 @@ get_absq(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 ROOT::VecOps::RVec<float>
 get_m(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 
+/**
+ * Get mass member from ReconstructedParticles
+ *
+ * @param particles List of reconstructed particles in an event
+ *
+ */
+ROOT::VecOps::RVec<float>
+get_m(const edm4hep::ReconstructedParticleCollection &particles);
+
 /**
  * Get goodnessOfPID member from ReconstructedParticles
  *
@@ -586,6 +595,18 @@ ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
 sel_e(LogicalOperators::ComparisonOperator op, float value,
       ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles);
 
+/**
+ * Select a subcollection of ReconstructedParticles based on the energy
+ *
+ * @param op Comparison operator to apply with the value
+ * @param value Value of the energy that is used in the comparison
+ * @param particles Collection of reconstructed particles in an event
+ *
+ */
+edm4hep::ReconstructedParticleCollection
+sel_e(LogicalOperators::ComparisonOperator op, float value,
+      const edm4hep::ReconstructedParticleCollection &particles);
+
 /**
  * Select a subcollection of ReconstructedParticles based on the 3D momentum
  * modulus

src/LogicalOperators.cc

@@ -1,37 +1,12 @@
 #include "ral/LogicalOperators.h"
 #include "ROOT/RVec.hxx"
-#include <edm4hep/ReconstructedParticleData.h>
 #include <stdexcept>
 #include <string>
 
 namespace k4::ral {
 
 namespace LogicalOperators {
 
-template <typename T>
-ROOT::VecOps::RVec<T> filter(ROOT::VecOps::RVec<bool> mask,
-                             ROOT::VecOps::RVec<T> collection) {
-  if (mask.size() != collection.size()) {
-    auto msg = "Different vector lenghts: " + std::to_string(mask.size()) +
-               " and " + std::to_string(collection.size()) + ".";
-    throw std::length_error(msg);
-  }
-  ROOT::VecOps::RVec<T> result;
-  for (int i = 0; i < collection.size(); i++) {
-    if (mask[i]) {
-      result.emplace_back(collection[i]);
-    }
-  }
-  return result;
-}
-
-/// It is important to instanciate the methods that we want to use in order
-/// them to be available
-template ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
-filter<edm4hep::ReconstructedParticleData>(
-    ROOT::VecOps::RVec<bool> mask,
-    ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> collection);
-
 ROOT::VecOps::RVec<bool> operator&&(const ROOT::VecOps::RVec<bool> &vec1,
                                     const ROOT::VecOps::RVec<bool> &vec2) {
   if (vec1.size() != vec2.size()) {

src/ReconstructedParticle.cc

@@ -49,6 +49,16 @@ get_e(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
   return result;
 }
 
+ROOT::VecOps::RVec<float>
+get_e(const edm4hep::ReconstructedParticleCollection &particles) {
+  ROOT::VecOps::RVec<float> result;
+  result.reserve(particles.size());
+  for (edm4hep::ReconstructedParticle p : particles) {
+    result.emplace_back(p.getEnergy());
+  }
+  return result;
+}
+
 ROOT::VecOps::RVec<ROOT::Math::PxPyPzMVector>
 get_p(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
   ROOT::VecOps::RVec<ROOT::Math::PxPyPzMVector> result;
@@ -247,6 +257,16 @@ get_m(ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
   return result;
 }
 
+ROOT::VecOps::RVec<float>
+get_m(const edm4hep::ReconstructedParticleCollection &particles) {
+  ROOT::VecOps::RVec<float> result;
+  result.reserve(particles.size());
+  for (const auto &p : particles) {
+    result.emplace_back(p.getMass());
+  }
+  return result;
+}
+
 ROOT::VecOps::RVec<float> get_goodnessOfPID(
     ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
   ROOT::VecOps::RVec<float> result;
@@ -420,7 +440,7 @@ int print_goodnessOfPID::operator()(
 
 #define MASKING(T, getter, op, value, collection, result)                      \
   ROOT::VecOps::RVec<T> vec = getter(collection);                              \
-  for (T & item : vec) {                                                       \
+  for (const T &item : vec) {                                                  \
     switch (op) {                                                              \
     case LogicalOperators::ComparisonOperator::LESS:                           \
       result.emplace_back(item < value);                                       \
@@ -449,6 +469,15 @@ mask_e(LogicalOperators::ComparisonOperator op, float value,
   return result;
 }
 
+ROOT::VecOps::RVec<bool>
+mask_e(LogicalOperators::ComparisonOperator op, float value,
+       const edm4hep::ReconstructedParticleCollection &particles) {
+  ROOT::VecOps::RVec<bool> result;
+  result.reserve(particles.size());
+  MASKING(float, get_e, op, value, particles, result);
+  return result;
+}
+
 ROOT::VecOps::RVec<bool>
 mask_pmod(LogicalOperators::ComparisonOperator op, float value,
           ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {
@@ -631,6 +660,13 @@ sel_e(LogicalOperators::ComparisonOperator op, float value,
       mask, particles);
 }
 
+edm4hep::ReconstructedParticleCollection
+sel_e(LogicalOperators::ComparisonOperator op, float value,
+      const edm4hep::ReconstructedParticleCollection &particles) {
+  auto mask = mask_e(op, value, particles);
+  return LogicalOperators::filter(mask, particles);
+}
+
 ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData>
 sel_pmod(LogicalOperators::ComparisonOperator op, float value,
          ROOT::VecOps::RVec<edm4hep::ReconstructedParticleData> particles) {

test/integration/CMakeLists.txt

@@ -1,2 +1,5 @@
 add_test(NAME integration_ReconstructedParticle
          COMMAND python "${CMAKE_CURRENT_LIST_DIR}/ReconstructedParticle.py")
+
+       add_test(NAME integration_ReconstructedParticleSource
+         COMMAND python "${CMAKE_CURRENT_LIST_DIR}/ReconstructedParticleSource.py")

test/integration/ReconstructedParticle.py

@@ -66,7 +66,7 @@ def get_file(url: str, path: str) -> None:
     if not ROOT.gSystem.Load("libral") and ROOT.gSystem.Load("libedm4hep"):
         print('RAL Found!')
         print('EDM4HEP Found!')
-    if ROOT.loadRal() and ROOT.edm4hep.ReconstructedParticleData():
+    if ROOT.loadRal and ROOT.edm4hep.ReconstructedParticleData:
         print('RAL Loaded!')
         print('EDM4HEP Loaded!')
 

test/integration/ReconstructedParticleSource.py

@@ -0,0 +1,96 @@
+import sys
+import subprocess
+import urllib.request
+from pathlib import Path
+import ROOT
+
+def get_file(url: str, path: str) -> None:
+    size = 0
+    with urllib.request.urlopen(url) as stream, open(path, "wb") as file:
+        buf_size = 1024 * 8
+        while block := stream.read(buf_size):
+            file.write(block)
+            size += buf_size
+    if size == 0:
+        raise RuntimeError(f"Cannot download from url: {url}")
+
+
+PYTHIA_CARD = {
+    "url": "https://raw.githubusercontent.com/HEP-FCC/FCC-config/winter2023/FCCee/Generator/Pythia8/p8_ee_WW_ecm240.cmd",
+    "path": "/tmp/test_pythia_card.cmd"
+}
+DETECTOR_CARD = {
+    "url": "https://raw.githubusercontent.com/HEP-FCC/FCC-config/winter2023/FCCee/Delphes/card_IDEA.tcl",
+    "path": "/tmp/test_detector_card.tcl"
+}
+OUTPUT_CARD = {
+    "url": "https://raw.githubusercontent.com/HEP-FCC/FCC-config/winter2023/FCCee/Delphes/card_IDEA.tcl",
+    "path": "/tmp/test_output_card.tcl"
+}
+ROOT_FILE = "/tmp/test.root"
+FINAL_ROOT_FILE = "/tmp/test_final.root"
+
+if __name__ == "__main__":
+
+    print("ReconstructedParticle Integration Test: PRE-TEST".center(80, "-"))
+
+    if not Path(ROOT_FILE).exists():
+        print("Downloading Pythia files")
+        get_file(**PYTHIA_CARD)
+        get_file(**DETECTOR_CARD)
+        get_file(**OUTPUT_CARD)
+        print("Running Pythia to generate root file")
+        subprocess.run(
+            [
+                "DelphesPythia8_EDM4HEP",
+                DETECTOR_CARD["path"],
+                OUTPUT_CARD["path"],
+                PYTHIA_CARD["path"],
+                ROOT_FILE
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL
+        ).check_returncode()
+    # else:
+    #     subprocess.run(
+    #         [
+    #             "podio-dump",
+    #             ROOT_FILE
+    #         ],
+    #         stdout=subprocess.DEVNULL,
+    #         stderr=subprocess.DEVNULL
+    #     ).check_returncode()
+
+    print("ReconstructedParticle Integration Test: START".center(80, "-"))
+
+    if ROOT.podio.DataSource:
+        print('PODIO::DataSource Loaded!')
+    if ROOT.edm4hep.ReconstructedParticle:
+        print('EDM4HEP Loaded!')
+    if ROOT.gSystem.Load("libral") < 0:
+        print('Loading of the Key4hep RAL library failed!\nAborting...')
+        sys.exit(1)
+    if ROOT.loadRal:
+        print('Key4hep RAL Loaded!')
+
+    print("Loading ROOT dataframe and test ral functions")
+    df = ROOT.podio.CreateDataFrame(ROOT_FILE)
+    ROOT.gInterpreter.ProcessLine("using namespace k4::ral;")
+    ROOT.gInterpreter.ProcessLine("using namespace LogicalOperators;")
+    df = (
+        df
+        .Define("e1",
+                "ReconstructedParticle::sel_e(ComparisonOperator::LESS, 1., ReconstructedParticles)")
+        .Define("e2",
+                "ReconstructedParticle::sel_e(ComparisonOperator::GREATER, 5., e1)")
+        .Define("mass",
+                "ReconstructedParticle::get_m(e2)")
+        .Define("total_mass",
+                "ROOT::VecOps::Sum(mass)")
+         )
+    print("Output test result in a new dataframe")
+    df.Snapshot("events", FINAL_ROOT_FILE,
+                [
+                    "mass",
+                    "total_mass"
+                ])

test/unittest/ReconstructedParticle.cpp

@@ -11,7 +11,6 @@
 #include "ral/LogicalOperators.h"
 #include "ral/ReconstructedParticle.h"
 #include <Math/Vector3Dfwd.h>
-#include <edm4hep/ReconstructedParticleData.h>
 
 using namespace k4::ral;