feat: example of how to use ONNX Runtime (Ort) in algorithms

CMakeLists.txt

@@ -219,6 +219,12 @@ if(NOT "cxx_std_${CMAKE_CXX_STANDARD}" IN_LIST ROOT_COMPILE_FEATURES)
   )
 endif()
 
+# ONNX Runtime
+option(USE_ONNX "Compile with ONNX support" ON)
+if(${USE_ONNX})
+  find_package(onnxruntime)
+endif()
+
 # Add CMake additional functionality:
 include(cmake/jana_plugin.cmake) # Add common settings for plugins
 list(APPEND CMAKE_MODULE_PATH ${EICRECON_SOURCE_DIR}/cmake

cmake/jana_plugin.cmake

@@ -376,3 +376,15 @@ macro(plugin_add_fastjet _name)
   plugin_link_libraries(${PLUGIN_NAME} ${FASTJET_LIBRARIES})
 
 endmacro()
+
+# Adds ONNX Runtime for a plugin
+macro(plugin_add_onnxruntime _name)
+
+  if(NOT onnxruntime_FOUND)
+    find_package(onnxruntime)
+  endif()
+
+  # Add libraries
+  plugin_link_libraries(${PLUGIN_NAME} onnxruntime::onnxruntime)
+
+endmacro()

src/algorithms/CMakeLists.txt

@@ -7,3 +7,7 @@ add_subdirectory(pid)
 add_subdirectory(digi)
 add_subdirectory(reco)
 add_subdirectory(fardetectors)
+
+if(USE_ONNX)
+  add_subdirectory(onnx)
+endif()

src/algorithms/onnx/CMakeLists.txt

@@ -0,0 +1,21 @@
+cmake_minimum_required(VERSION 3.16)
+
+set(PLUGIN_NAME "algorithms_onnx")
+
+# Function creates ${PLUGIN_NAME}_plugin and ${PLUGIN_NAME}_library targets
+# Setting default includes, libraries and installation paths
+plugin_add(${PLUGIN_NAME} WITH_SHARED_LIBRARY WITHOUT_PLUGIN)
+
+# The macro grabs sources as *.cc *.cpp *.c and headers as *.h *.hh *.hpp Then
+# correctly sets sources for ${_name}_plugin and ${_name}_library targets Adds
+# headers to the correct installation directory
+plugin_glob_all(${PLUGIN_NAME})
+
+# Find dependencies
+plugin_add_event_model(${PLUGIN_NAME})
+plugin_add_onnxruntime(${PLUGIN_NAME})
+
+# The macro grabs sources as *.cc *.cpp *.c and headers as *.h *.hh *.hpp Then
+# correctly sets sources for ${_name}_plugin and ${_name}_library targets Adds
+# headers to the correct installation directory
+plugin_glob_all(${PLUGIN_NAME})

src/algorithms/onnx/InclusiveKinematicsML.cc

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2022, 2023 Wouter Deconinck, Tooba Ali
+
+#include <fmt/core.h>
+#include <onnxruntime_c_api.h>
+#include <onnxruntime_cxx_api.h>
+#include <algorithm>
+#include <cstddef>
+#include <exception>
+#include <gsl/pointers>
+#include <iterator>
+#include <ostream>
+
+#include "InclusiveKinematicsML.h"
+
+namespace eicrecon {
+
+  static std::string print_shape(const std::vector<std::int64_t>& v) {
+    std::stringstream ss("");
+    for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x";
+    ss << v[v.size() - 1];
+    return ss.str();
+  }
+
+  template <typename T>
+  Ort::Value vec_to_tensor(std::vector<T>& data, const std::vector<std::int64_t>& shape) {
+    Ort::MemoryInfo mem_info =
+        Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
+    auto tensor = Ort::Value::CreateTensor<T>(mem_info, data.data(), data.size(), shape.data(), shape.size());
+    return tensor;
+  }
+
+  void InclusiveKinematicsML::init(std::shared_ptr<spdlog::logger>& logger) {
+    m_log = logger;
+
+    // onnxruntime setup
+    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "inclusive-kinematics-ml");
+    Ort::SessionOptions session_options;
+    try {
+      m_session = Ort::Session(env, m_cfg.modelPath.c_str(), session_options);
+
+      // print name/shape of inputs
+      Ort::AllocatorWithDefaultOptions allocator;
+      m_log->debug("Input Node Name/Shape:");
+      for (std::size_t i = 0; i < m_session.GetInputCount(); i++) {
+        m_input_names.emplace_back(m_session.GetInputNameAllocated(i, allocator).get());
+        m_input_shapes.emplace_back(m_session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
+        m_log->debug("\t{} : {}", m_input_names.at(i), print_shape(m_input_shapes.at(i)));
+      }
+
+      // print name/shape of outputs
+      m_log->debug("Output Node Name/Shape:");
+      for (std::size_t i = 0; i < m_session.GetOutputCount(); i++) {
+        m_output_names.emplace_back(m_session.GetOutputNameAllocated(i, allocator).get());
+        m_output_shapes.emplace_back(m_session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
+        m_log->debug("\t{} : {}", m_output_names.at(i), print_shape(m_output_shapes.at(i)));
+      }
+
+      // convert names to char*
+      m_input_names_char.resize(m_input_names.size(), nullptr);
+      std::transform(std::begin(m_input_names), std::end(m_input_names), std::begin(m_input_names_char),
+                     [&](const std::string& str) { return str.c_str(); });
+      m_output_names_char.resize(m_output_names.size(), nullptr);
+      std::transform(std::begin(m_output_names), std::end(m_output_names), std::begin(m_output_names_char),
+                     [&](const std::string& str) { return str.c_str(); });
+
+    } catch(std::exception& e) {
+      m_log->error(e.what());
+    }
+  }
+
+  void InclusiveKinematicsML::process(
+      const InclusiveKinematicsML::Input& input,
+      const InclusiveKinematicsML::Output& output) const {
+
+    const auto [electron, da] = input;
+    auto [ml] = output;
+
+    // Require valid inputs
+    if (electron->size() == 0 || da->size() == 0) {
+      m_log->debug("skipping because input collections have no entries");
+      return;
+    }
+
+    // Assume model has 1 input nodes and 1 output node.
+    if (m_input_names.size() != 1 || m_output_names.size() != 1) {
+      m_log->debug("skipping because model has incorrect input and output size");
+      return;
+    }
+
+    // Prepare input tensor
+    std::vector<float> input_tensor_values;
+    std::vector<Ort::Value> input_tensors;
+    for (std::size_t i = 0; i < electron->size(); i++) {
+      input_tensor_values.push_back(electron->at(i).getX());
+    }
+    input_tensors.emplace_back(vec_to_tensor<float>(input_tensor_values, m_input_shapes.front()));
+
+    // Double-check the dimensions of the input tensor
+    if (! input_tensors[0].IsTensor() || input_tensors[0].GetTensorTypeAndShapeInfo().GetShape() != m_input_shapes.front()) {
+      m_log->debug("skipping because input tensor shape incorrect");
+      return;
+    }
+
+    // Attempt inference
+    try {
+      auto output_tensors = m_session.Run(Ort::RunOptions{nullptr}, m_input_names_char.data(), input_tensors.data(),
+                                          m_input_names_char.size(), m_output_names_char.data(), m_output_names_char.size());
+
+      // Double-check the dimensions of the output tensors
+      if (!output_tensors[0].IsTensor() || output_tensors.size() != m_output_names.size()) {
+        m_log->debug("skipping because output tensor shape incorrect");
+        return;
+      }
+
+      // Convert output tensor
+      float* output_tensor_data = output_tensors[0].GetTensorMutableData<float>();
+      auto x  = output_tensor_data[0];
+      auto kin = ml->create();
+      kin.setX(x);
+
+    } catch (const Ort::Exception& exception) {
+      m_log->error("error running model inference: {}", exception.what());
+    }
+  }
+
+} // namespace eicrecon

src/algorithms/onnx/InclusiveKinematicsML.h

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2022, 2023 Sylvester Joosten, Dmitry Romanov, Wouter Deconinck
+
+#pragma once
+
+#include <algorithms/algorithm.h>
+#include <edm4eic/InclusiveKinematicsCollection.h>
+#include <onnxruntime_cxx_api.h>
+#include <spdlog/logger.h>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "algorithms/interfaces/WithPodConfig.h"
+#include "algorithms/onnx/InclusiveKinematicsMLConfig.h"
+
+namespace eicrecon {
+
+  using InclusiveKinematicsMLAlgorithm = algorithms::Algorithm<
+    algorithms::Input<
+      edm4eic::InclusiveKinematicsCollection,
+      edm4eic::InclusiveKinematicsCollection
+    >,
+    algorithms::Output<
+      edm4eic::InclusiveKinematicsCollection
+    >
+  >;
+
+  class InclusiveKinematicsML
+  : public InclusiveKinematicsMLAlgorithm,
+    public WithPodConfig<InclusiveKinematicsMLConfig> {
+
+  public:
+    InclusiveKinematicsML(std::string_view name)
+      : InclusiveKinematicsMLAlgorithm{name,
+                            {"inclusiveKinematicsElectron", "inclusiveKinematicsDA"},
+                            {"inclusiveKinematicsML"},
+                            "Determine inclusive kinematics using combined ML method."} {}
+
+    void init(std::shared_ptr<spdlog::logger>& logger);
+    void process(const Input&, const Output&) const final;
+
+  private:
+    std::shared_ptr<spdlog::logger> m_log;
+
+    mutable Ort::Session m_session{nullptr};
+
+    std::vector<std::string> m_input_names;
+    std::vector<const char*> m_input_names_char;
+    std::vector<std::vector<std::int64_t>> m_input_shapes;
+
+    std::vector<std::string> m_output_names;
+    std::vector<const char*> m_output_names_char;
+    std::vector<std::vector<std::int64_t>> m_output_shapes;
+
+  };
+
+} // namespace eicrecon

src/algorithms/onnx/InclusiveKinematicsMLConfig.h

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 Wouter Deconinck
+
+#pragma once
+
+#include <string>
+
+namespace eicrecon {
+
+  struct InclusiveKinematicsMLConfig {
+
+    std::string modelPath{"calibrations/onnx/identity_gemm_w1x1_b1.onnx"};
+
+  };
+
+} // eicrecon

src/factories/reco/InclusiveKinematicsML_factory.h

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2022 Wouter Deconinck
+
+#pragma once
+
+#include <JANA/JEvent.h>
+#include <edm4eic/InclusiveKinematicsCollection.h>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "algorithms/onnx/InclusiveKinematicsML.h"
+#include "extensions/jana/JOmniFactory.h"
+
+namespace eicrecon {
+
+class InclusiveKinematicsML_factory :
+        public JOmniFactory<InclusiveKinematicsML_factory, InclusiveKinematicsMLConfig> {
+
+public:
+    using AlgoT = eicrecon::InclusiveKinematicsML;
+private:
+    std::unique_ptr<AlgoT> m_algo;
+
+    PodioInput<edm4eic::InclusiveKinematics> m_inclusive_kinematics_electron_input {this};
+    PodioInput<edm4eic::InclusiveKinematics> m_inclusive_kinematics_da_input {this};
+    PodioOutput<edm4eic::InclusiveKinematics> m_inclusive_kinematics_output {this};
+
+    ParameterRef<std::string> m_modelPath {this, "modelPath", config().modelPath};
+
+public:
+    void Configure() {
+        m_algo = std::make_unique<AlgoT>(GetPrefix());
+        m_algo->applyConfig(config());
+        m_algo->init(logger());
+    }
+
+    void ChangeRun(int64_t run_number) {
+    }
+
+    void Process(int64_t run_number, uint64_t event_number) {
+        m_algo->process({m_inclusive_kinematics_electron_input(), m_inclusive_kinematics_da_input()}, {m_inclusive_kinematics_output().get()});
+    }
+};
+
+} // eicrecon

src/global/reco/CMakeLists.txt

@@ -26,5 +26,6 @@ plugin_glob_all(${PLUGIN_NAME})
 
 # Add libraries (same as target_include_directories but for both plugin and
 # library)
-plugin_link_libraries(${PLUGIN_NAME} algorithms_digi_library
-                      algorithms_tracking_library algorithms_reco_library)
+plugin_link_libraries(
+  ${PLUGIN_NAME} algorithms_digi_library algorithms_tracking_library
+  algorithms_onnx_library algorithms_reco_library)

src/global/reco/reco.cc

@@ -11,13 +11,15 @@
 #include <map>
 #include <memory>
 
+#include "algorithms/interfaces/WithPodConfig.h"
 #include "algorithms/reco/InclusiveKinematicsDA.h"
 #include "algorithms/reco/InclusiveKinematicsElectron.h"
 #include "algorithms/reco/InclusiveKinematicsJB.h"
 #include "algorithms/reco/InclusiveKinematicsSigma.h"
 #include "algorithms/reco/InclusiveKinematicseSigma.h"
 #include "extensions/jana/JOmniFactoryGeneratorT.h"
 #include "factories/meta/CollectionCollector_factory.h"
+#include "factories/reco/InclusiveKinematicsML_factory.h"
 #include "factories/reco/InclusiveKinematicsReconstructed_factory.h"
 #include "factories/reco/InclusiveKinematicsTruth_factory.h"
 #include "factories/reco/JetReconstruction_factory.h"
@@ -155,6 +157,18 @@ void InitPlugin(JApplication *app) {
         app
     ));
 
+    app->Add(new JOmniFactoryGeneratorT<InclusiveKinematicsML_factory>(
+        "InclusiveKinematicsML",
+        {
+          "InclusiveKinematicsElectron",
+          "InclusiveKinematicsDA"
+        },
+        {
+          "InclusiveKinematicsML"
+        },
+        app
+    ));
+
     app->Add(new JOmniFactoryGeneratorT<ReconstructedElectrons_factory>(
         "ReconstructedElectrons",
         {"ReconstructedParticles"},

src/services/io/podio/JEventProcessorPODIO.cc

@@ -136,6 +136,7 @@ JEventProcessorPODIO::JEventProcessorPODIO() {
             "CentralCKFSeededTrackParameters",
             "InclusiveKinematicsDA",
             "InclusiveKinematicsJB",
+            "InclusiveKinematicsML",
             "InclusiveKinematicsSigma",
             "InclusiveKinematicseSigma",
             "InclusiveKinematicsElectron",