Add utility function to load parameters in RDataFrame

AIDASoft · Jun 12, 2024 · b95f710 · b95f710
1 parent b38c1d6
commit b95f710
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 4 deletions.
diff --git a/include/podio/GenericParameters.h b/include/podio/GenericParameters.h
@@ -105,8 +105,8 @@ class GenericParameters {
   }
 
   /// Load multiple key value pairs simultaneously
-  template <typename T>
-  void loadFrom(std::vector<std::string> keys, std::vector<std::vector<T>> values);
+  template <typename T, template <typename...> typename VecLike>
+  void loadFrom(VecLike<std::string> keys, VecLike<std::vector<T>> values);
 
   /// Get the number of elements stored under the given key for a type
   template <typename T, typename = EnableIfValidGenericDataType<T>>
@@ -277,8 +277,8 @@ std::vector<std::vector<T>> GenericParameters::getValues() const {
   return values;
 }
 
-template <typename T>
-void GenericParameters::loadFrom(std::vector<std::string> keys, std::vector<std::vector<T>> values) {
+template <typename T, template <typename...> typename VecLike>
+void GenericParameters::loadFrom(VecLike<std::string> keys, VecLike<std::vector<T>> values) {
   auto& map = getMap<T>();
   auto& mtx = getMutex<T>();
 

diff --git a/include/podio/utilities/RootHelpers.h b/include/podio/utilities/RootHelpers.h
@@ -1,6 +1,9 @@
 #ifndef PODIO_UTILITIES_ROOTHELPERS_H
 #define PODIO_UTILITIES_ROOTHELPERS_H
 
+#include "podio/GenericParameters.h"
+
+#include "ROOT/RVec.hxx"
 #include "TBranch.h"
 
 #include <string>
@@ -77,6 +80,12 @@ namespace root_utils {
     std::vector<std::vector<T>>* m_valuesPtr{nullptr};
   };
 
+  GenericParameters
+  loadParamsFrom(ROOT::VecOps::RVec<std::string> intKeys, ROOT::VecOps::RVec<std::vector<int>> intValues,
+                 ROOT::VecOps::RVec<std::string> floatKeys, ROOT::VecOps::RVec<std::vector<float>> floatValues,
+                 ROOT::VecOps::RVec<std::string> doubleKeys, ROOT::VecOps::RVec<std::vector<double>> doubleValues,
+                 ROOT::VecOps::RVec<std::string> stringKeys, ROOT::VecOps::RVec<std::vector<std::string>> stringValues);
+
 } // namespace root_utils
 } // namespace podio
 

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -81,6 +81,7 @@ SET(root_sources
   ROOTReader.cc
   ROOTLegacyReader.cc
   ROOTFrameData.cc
+  RootHelpers.cc
 )
 if(ENABLE_RNTUPLE)
   list(APPEND root_sources

diff --git a/src/RootHelpers.cc b/src/RootHelpers.cc
@@ -0,0 +1,17 @@
+#include "podio/utilities/RootHelpers.h"
+
+namespace podio::root_utils {
+GenericParameters
+loadParamsFrom(ROOT::VecOps::RVec<std::string> intKeys, ROOT::VecOps::RVec<std::vector<int>> intValues,
+               ROOT::VecOps::RVec<std::string> floatKeys, ROOT::VecOps::RVec<std::vector<float>> floatValues,
+               ROOT::VecOps::RVec<std::string> doubleKeys, ROOT::VecOps::RVec<std::vector<double>> doubleValues,
+               ROOT::VecOps::RVec<std::string> stringKeys, ROOT::VecOps::RVec<std::vector<std::string>> stringValues) {
+  GenericParameters params{};
+  params.loadFrom(std::move(intKeys), std::move(intValues));
+  params.loadFrom(std::move(floatKeys), std::move(floatValues));
+  params.loadFrom(std::move(doubleKeys), std::move(doubleValues));
+  params.loadFrom(std::move(stringKeys), std::move(stringValues));
+  return params;
+}
+
+} // namespace podio::root_utils
diff --git a/tests/root_io/CMakeLists.txt b/tests/root_io/CMakeLists.txt
@@ -62,3 +62,7 @@ if (ENABLE_RNTUPLE)
 
   set_property(TEST read_python_frame_rntuple PROPERTY DEPENDS write_python_frame_rntuple)
 endif()
+
+add_test(NAME param_reading_rdataframe COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/param_reading_rdataframe.py example_frame.root)
+PODIO_SET_TEST_ENV(param_reading_rdataframe)
+set_property(TEST param_reading_rdataframe PROPERTY DEPENDS write_frame_root)
diff --git a/tests/root_io/param_reading_rdataframe.py b/tests/root_io/param_reading_rdataframe.py
@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+import sys
+import ROOT
+
+
+ROOT.gInterpreter.Declare("#include <podio/utilities/RootHelpers.h>")
+df = ROOT.RDataFrame("events", sys.argv[1])
+
+evtWeights = (
+    df.Define(
+        "params",
+        "podio::root_utils::loadParamsFrom(GPIntKeys, GPIntValues, GPFloatKeys, GPFloatValues, GPDoubleKeys, GPDoubleValues, GPStringKeys, GPStringValues)",
+    )
+    .Define("eventweight", 'params.get<float>("UserEventWeight").value()')
+    .Histo1D("eventweight")
+)
+
+if evtWeights.GetMean() != 450.0:
+    print("Parameter values not as expected")
+    sys.exit(1)