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

### Briefly, what does this PR introduce?
This PR adds the hooks for ONNX Runtime (CPU only) for fast inference,
and an example that doesn't actually do anything except... Enabled by
default but not actually doing anything.

### What kind of change does this PR introduce?
- [ ] Bug fix (issue #__)
- [x] New feature (issue: use ONNX for ML)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [ ] Tests for the changes have been added
- [ ] Documentation has been added / updated
- [x] Changes have been communicated to collaborators @rahmans1 

### Does this PR introduce breaking changes? What changes might users
need to make to their code?
No.

### Does this PR change default behavior?
No.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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

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