Add some more RDF utility functionality and clean up implementation

test/test_rdf.cc

@@ -22,14 +22,17 @@ int main(int argc, char* argv[]) {
                  .Define("MCParticles_cosTheta", edm4hep::utils::cos_theta<edm4hep::MCParticleData>, {"MCParticles"})
                  .Define("SimTrackerHits_r", edm4hep::utils::r<edm4hep::SimTrackerHitData>, {"SimTrackerHits"})
                  .Define("SimTrackerHit_pt", edm4hep::utils::pt<edm4hep::SimTrackerHitData>, {"SimTrackerHits"})
-                 .Define("TrackerHits_r", edm4hep::utils::r<edm4hep::TrackerHitPlaneData>, {"TrackerHitPlanes"});
+                 .Define("TrackerHits_r", edm4hep::utils::r<edm4hep::TrackerHitPlaneData>, {"TrackerHitPlanes"})
+                 .Define("MCParticle_p4", edm4hep::utils::p4M<edm4hep::MCParticleData>, {"MCParticles"})
+
+                 .Define("MCParticle_energy", edm4hep::utils::E<edm4hep::LorentzVectorM>, {"MCParticle_p4"});
 
   std::string outfilename = "edm4hep_events_rdf.root";
   std::cout << "Writing snapshot to disk ... \t" << outfilename << std::endl;
 
   df2.Snapshot("events", outfilename,
                {"MCParticles_pt", "MCParticles_eta", "MCParticles_cosTheta", "SimTrackerHits_r", "SimTrackerHit_pt",
-                "TrackerHits_r"});
+                "TrackerHits_r", "MCParticle_energy"});
 
   return 0;
 }

test/test_rdf.py

@@ -20,6 +20,8 @@
        .Define('SimTrackerHits_r', 'edm4hep::utils::r(SimTrackerHits)')
        .Define('SimTrackerHits_pt', 'edm4hep::utils::pt(SimTrackerHits)')
        .Define('TrackerHits_r', 'edm4hep::utils::r(TrackerHitPlanes)')
+       .Define('MCParticle_p4', 'edm4hep::utils::p4M(MCParticles)')
+       .Define('MCParticle_energy', 'edm4hep::utils::E(MCParticle_p4)')
        )
 
 filename = 'edm4hep_events_py_rdf.root'
@@ -31,4 +33,5 @@
               'MCParticles_cosTheta',
               'SimTrackerHits_r',
               'SimTrackerHits_pt',
-              'TrackerHits_r'])
+              'TrackerHits_r',
+              'MCParticle_energy'])

test/utils/CMakeLists.txt

@@ -24,4 +24,5 @@ add_test(NAME pyunittests COMMAND python -m unittest discover -s ${CMAKE_CURRENT
 set_property(TEST pyunittests
   PROPERTY ENVIRONMENT
   LD_LIBRARY_PATH=$<TARGET_FILE_DIR:edm4hep>:$ENV{LD_LIBRARY_PATH}
+  ROOT_INCLUDE_PATH=${PROJECT_SOURCE_DIR}/utils/include
   )

test/utils/test_kinematics.py

@@ -11,7 +11,7 @@
   print('Cannot load edm4hep dictionary for tests')
   sys.exit(1)
 
-ROOT.gInterpreter.LoadFile(os.path.dirname(__file__) + '/../../utils/include/edm4hep/utils/kinematics.h')
+ROOT.gInterpreter.LoadFile('edm4hep/utils/kinematics.h')
 
 from ROOT import edm4hep
 

utils/include/edm4hep/utils/common.h

@@ -0,0 +1,19 @@
+#ifndef EDM4HEP_UTILS_COMMON_H
+#define EDM4HEP_UTILS_COMMON_H
+
+#include "Math/Vector4D.h"
+
+namespace edm4hep {
+/**
+ * A LorentzVector with (px, py, pz) and M
+ */
+using LorentzVectorM = ROOT::Math::PxPyPzMVector;
+
+/**
+ * A LorentzVector with (px, py, pz) and E
+ */
+using LorentzVectorE = ROOT::Math::PxPyPzEVector;
+
+} // namespace edm4hep
+
+#endif

utils/include/edm4hep/utils/dataframe.h

@@ -1,6 +1,8 @@
 #ifndef EDM4HEP_UTILS_DATAFRAME_H
 #define EDM4HEP_UTILS_DATAFRAME_H
 
+#include "edm4hep/utils/common.h"
+
 #include "ROOT/RVec.hxx"
 
 namespace edm4hep::utils {
@@ -21,6 +23,22 @@ ROOT::VecOps::RVec<float> cos_theta(ROOT::VecOps::RVec<T> const& in);
 template <typename T>
 ROOT::VecOps::RVec<float> r(ROOT::VecOps::RVec<T> const& in);
 
+/// Get the 4 momentum of the passed particle using the momentum and the mass
+template <typename T>
+ROOT::VecOps::RVec<edm4hep::LorentzVectorM> p4M(ROOT::VecOps::RVec<T> const& in);
+
+/// Get the 4 momentum of the passed particle using the momentum and the energy
+template <typename T>
+ROOT::VecOps::RVec<edm4hep::LorentzVectorE> p4E(ROOT::VecOps::RVec<T> const& in);
+
+/// Get the energy from a four momentum vector
+template <typename T>
+ROOT::VecOps::RVec<float> E(ROOT::VecOps::RVec<T> const& fourMom);
+
+/// Get the mass from a four momentum vector
+template <typename T>
+ROOT::VecOps::RVec<float> M(ROOT::VecOps::RVec<T> const& fourMom);
+
 } // namespace edm4hep::utils
 
 #endif // EDM4HEP_UTILS_DATAFRAME_H

utils/include/edm4hep/utils/kinematics.h

@@ -1,21 +1,11 @@
 #ifndef EDM4HEP_UTILS_KINEMATICS_H
 #define EDM4HEP_UTILS_KINEMATICS_H
 
-#include "Math/Vector4D.h"
+#include "edm4hep/utils/common.h"
 
 #include <cmath>
 
 namespace edm4hep {
-/**
- * A LorentzVector with (px, py, pz) and M
- */
-using LorentzVectorM = ROOT::Math::PxPyPzMVector;
-
-/**
- * A LorentzVector with (px, py, pz) and E
- */
-using LorentzVectorE = ROOT::Math::PxPyPzEVector;
-
 namespace utils {
 
   /**

utils/src/dataframe.cc

@@ -1,4 +1,5 @@
 #include "edm4hep/utils/dataframe.h"
+#include "edm4hep/utils/common.h"
 
 #include "edm4hep/CalorimeterHitData.h"
 #include "edm4hep/ClusterData.h"
@@ -16,44 +17,55 @@ namespace edm4hep::utils {
 
 template <typename T>
 ROOT::VecOps::RVec<float> pt(ROOT::VecOps::RVec<T> const& in) {
-  ROOT::VecOps::RVec<float> result;
-  result.reserve(in.size());
-  for (size_t i = 0; i < in.size(); ++i) {
-    result.push_back(std::sqrt(in[i].momentum.x * in[i].momentum.x + in[i].momentum.y * in[i].momentum.y));
-  }
-  return result;
+  return ROOT::VecOps::Map(
+      in, [](const auto& p) { return std::sqrt(p.momentum.x * p.momentum.x + p.momentum.y + p.momentum.y); });
 }
 
 template <typename T>
 ROOT::VecOps::RVec<float> eta(ROOT::VecOps::RVec<T> const& in) {
-  ROOT::VecOps::RVec<float> result;
-  result.reserve(in.size());
-  for (size_t i = 0; i < in.size(); ++i) {
-    ROOT::Math::XYZVector lv{in[i].momentum.x, in[i].momentum.y, in[i].momentum.z};
-    result.push_back(lv.Eta());
-  }
-  return result;
+  return ROOT::VecOps::Map(in, [](const auto& p) {
+    ROOT::Math::XYZVector lv{p.momentum.x, p.momentum.y, p.momentum.z};
+    return lv.Eta();
+  });
 }
 
 template <typename T>
 ROOT::VecOps::RVec<float> cos_theta(ROOT::VecOps::RVec<T> const& in) {
-  ROOT::VecOps::RVec<float> result;
-  result.reserve(in.size());
-  for (size_t i = 0; i < in.size(); ++i) {
-    ROOT::Math::XYZVector lv{in[i].momentum.x, in[i].momentum.y, in[i].momentum.z};
-    result.push_back(cos(lv.Theta()));
-  }
-  return result;
+  return ROOT::VecOps::Map(in, [](const auto& p) {
+    ROOT::Math::XYZVector lv{p.momentum.x, p.momentum.y, p.momentum.z};
+    return std::cos(lv.Theta());
+  });
 }
 
 template <typename T>
 ROOT::VecOps::RVec<float> r(ROOT::VecOps::RVec<T> const& in) {
-  ROOT::VecOps::RVec<double> result;
-  result.reserve(in.size());
-  for (size_t i = 0; i < in.size(); ++i) {
-    result.push_back(std::sqrt(in[i].position.x * in[i].position.x + in[i].position.y * in[i].position.y));
-  }
-  return result;
+  return ROOT::VecOps::Map(in, [](const auto& p) {
+    return std::sqrt(p.position.x * p.position.x + p.position.y * p.position.y + p.position.z + p.position.z);
+  });
+}
+
+template <typename T>
+ROOT::VecOps::RVec<edm4hep::LorentzVectorM> p4M(ROOT::VecOps::RVec<T> const& in) {
+  return ROOT::VecOps::Map(in, [](const auto& p) {
+    return edm4hep::LorentzVectorM{p.momentum.x, p.momentum.y, p.momentum.z, p.mass};
+  });
+}
+
+template <typename T>
+ROOT::VecOps::RVec<edm4hep::LorentzVectorE> p4E(ROOT::VecOps::RVec<T> const& in) {
+  return ROOT::VecOps::Map(in, [](const auto& p) {
+    return edm4hep::LorentzVectorE{p.momentum.x, p.momentum.y, p.momentum.z, p.energy};
+  });
+}
+
+template <typename T>
+ROOT::VecOps::RVec<float> E(ROOT::VecOps::RVec<T> const& fourMom) {
+  return ROOT::VecOps::Map(fourMom, [](const auto& p) { return p.E(); });
+}
+
+template <typename T>
+ROOT::VecOps::RVec<float> M(ROOT::VecOps::RVec<T> const& fourMom) {
+  return ROOT::VecOps::Map(fourMom, [](const auto& p) { return p.M(); });
 }
 
 // Explicitly instantiate the template functions here to have them available in
@@ -69,13 +81,15 @@ ROOT::VecOps::RVec<float> r(ROOT::VecOps::RVec<T> const& in) {
   INST_DATA_TO_FLOAT_VEC_FUNC(eta, DATATYPE);                                                                          \
   INST_DATA_TO_FLOAT_VEC_FUNC(cos_theta, DATATYPE)
 
-// Macro to instantiate all position related functions for a datatype
-#define INST_POSITION_FUNCS(DATATYPE) INST_DATA_TO_FLOAT_VEC_FUNC(r, DATATYPE)
-
 INST_MOMENTUM_FUNCS(edm4hep::MCParticleData);
 INST_MOMENTUM_FUNCS(edm4hep::ReconstructedParticleData);
 INST_MOMENTUM_FUNCS(edm4hep::SimTrackerHitData);
 
+#undef INST_MOMENTUM_FUNCS
+
+// Macro to instantiate all position related functions for a datatype
+#define INST_POSITION_FUNCS(DATATYPE) INST_DATA_TO_FLOAT_VEC_FUNC(r, DATATYPE)
+
 INST_POSITION_FUNCS(edm4hep::SimTrackerHitData);
 INST_POSITION_FUNCS(edm4hep::TrackerHitData);
 INST_POSITION_FUNCS(edm4hep::TrackerHitPlaneData);
@@ -84,8 +98,26 @@ INST_POSITION_FUNCS(edm4hep::CalorimeterHitData);
 INST_POSITION_FUNCS(edm4hep::ClusterData);
 INST_POSITION_FUNCS(edm4hep::VertexData);
 
-#undef INST_DATA_TO_FLOAT_VEC
 #undef INST_POSITION_FUNCS
-#undef INST_MOMENTUM_FUNCS
+
+INST_DATA_TO_FLOAT_VEC_FUNC(E, edm4hep::LorentzVectorE);
+INST_DATA_TO_FLOAT_VEC_FUNC(E, edm4hep::LorentzVectorM);
+INST_DATA_TO_FLOAT_VEC_FUNC(M, edm4hep::LorentzVectorE);
+INST_DATA_TO_FLOAT_VEC_FUNC(M, edm4hep::LorentzVectorM);
+
+#undef INST_DATA_TO_FLOAT_VEC
+
+#define INST_4MOM_MASS_FUNCS(DATATYPE)                                                                                 \
+  template ROOT::VecOps::RVec<edm4hep::LorentzVectorM> p4M(ROOT::VecOps::RVec<DATATYPE> const&)
+
+#define INST_4MOM_ENERGY_FUNCS(DATATYPE)                                                                               \
+  template ROOT::VecOps::RVec<edm4hep::LorentzVectorE> p4E(ROOT::VecOps::RVec<DATATYPE> const&)
+
+INST_4MOM_ENERGY_FUNCS(edm4hep::ReconstructedParticleData);
+INST_4MOM_MASS_FUNCS(edm4hep::ReconstructedParticleData);
+INST_4MOM_MASS_FUNCS(edm4hep::MCParticleData);
+
+#undef INST_4MOM_ENERGY_FUNCS
+#undef INST_4MOM_MASS_FUNCS
 
 } // namespace edm4hep::utils