Improve pidTPCML.cxx

Common/TableProducer/PID/CMakeLists.txt

@@ -11,6 +11,11 @@
 
 # TOF
 
+
+o2physics_add_library(PIDTPCCore
+             SOURCES pidTPCML.cxx
+             PUBLIC_LINK_LIBRARIES O2::Framework ONNXRuntime::ONNXRuntime
+)
 o2physics_add_dpl_workflow(pid-tof-base
                     SOURCES pidTOFBase.cxx
                     PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2::TOFBase
@@ -35,12 +40,12 @@ o2physics_add_dpl_workflow(pid-tof-full
 
 o2physics_add_dpl_workflow(pid-tpc
                     SOURCES pidTPC.cxx
-                    PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime
+                    PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(pid-tpc-full
                     SOURCES pidTPCFull.cxx
-                    PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime
+                    PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime
                     COMPONENT_NAME Analysis)
 
 # HMPID
@@ -55,4 +60,4 @@ o2physics_add_dpl_workflow(pid-hmpid-qa
 o2physics_add_dpl_workflow(pid-bayes
                     SOURCES pidBayes.cxx
                     PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore
-                    COMPONENT_NAME Analysis)
+                    COMPONENT_NAME Analysis)

Common/TableProducer/PID/pidTPCML.cxx

@@ -0,0 +1,217 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+///
+/// \file     pidTPCML.h
+///
+/// \author   Christian Sonnabend <christian.sonnabend@cern.ch>
+/// \author   Nicolò Jacazio <nicolo.jacazio@cern.ch>
+///
+/// \brief    A class for loading an ONNX neural network and evaluating it for the TPC PID response
+///
+#include <vector>
+
+#include "TSystem.h"
+
+// O2 includes
+#include "Framework/Logger.h"
+#include "Common/TableProducer/PID/pidTPCML.h"
+#include "ReconstructionDataFormats/PID.h"
+#include <onnxruntime/core/session/experimental_onnxruntime_cxx_api.h>
+
+namespace o2::pid::tpc
+{
+
+std::string Network::printShape(const std::vector<int64_t>& v)
+{
+  std::stringstream ss("");
+  for (size_t i = 0; i < v.size() - 1; i++)
+    ss << v[i] << "x";
+  ss << v[v.size() - 1];
+  return ss.str();
+}
+
+Network::Network(std::string path,
+                 bool enableOptimization = true)
+{
+
+  /*
+  Constructor: Creating a class instance from a file and enabling optimizations with the boolean option.
+  - Input:
+    -- pathLocally:         std::string   ; Path to the model file;
+    -- loadFromAlien:       bool          ; Download network from AliEn directory (true) or use local file (false)
+    -- pathAlien:           std::string   ; if loadFromAlien is true, then the network will be downloaded from pathAlien
+    -- enableOptimization:  bool          ; enabling optimizations for the loaded model in the session options;
+  */
+
+  LOG(info) << "--- Neural Network for the TPC PID response correction ---";
+
+  mEnv = std::make_shared<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "pid-neural-network");
+  if (enableOptimization) {
+    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
+  }
+
+  mSession.reset(new Ort::Experimental::Session{*mEnv, path, sessionOptions});
+
+  mInputNames = mSession->GetInputNames();
+  mInputShapes = mSession->GetInputShapes();
+  mOutputNames = mSession->GetOutputNames();
+  mOutputShapes = mSession->GetOutputShapes();
+
+  LOG(info) << "Input Nodes:";
+  for (size_t i = 0; i < mInputNames.size(); i++) {
+    LOG(info) << "\t" << mInputNames[i] << " : " << printShape(mInputShapes[i]);
+  }
+
+  LOG(info) << "Output Nodes:";
+  for (size_t i = 0; i < mOutputNames.size(); i++) {
+    LOG(info) << "\t" << mOutputNames[i] << " : " << printShape(mOutputShapes[i]);
+  }
+
+  LOG(info) << "--- Network initialized! ---";
+
+} // Network::Network(std::string, bool)
+
+Network& Network::operator=(Network& inst)
+{
+
+  /*
+  Operator: Setting instances of one class to the instances of an input class
+  - Input:
+    -- inst:    const Network&  ; An instance of the network class;
+  - Output:
+    -- *this:   Network&        ; An instance with the private properties of the "inst" (input) instance;
+  */
+
+  mEnv = inst.mEnv;
+  mSession = inst.mSession;
+  // sessionOptions = inst.sessionOptions; // Comment (Christian): Somehow the session options throw an error when trying to copy
+  mInputNames = inst.mInputNames;
+  mInputShapes = inst.mInputShapes;
+  mOutputNames = inst.mOutputNames;
+  mOutputShapes = inst.mOutputShapes;
+
+  LOG(debug) << "Network copied!";
+
+  return *this;
+
+} // Network& Network::operator=(const Network &)
+
+template <typename C, typename T>
+std::array<float, 6> Network::createInputFromTrack(const C& collision_it, const T& track, const uint8_t id) const
+{
+
+  /*
+  Function: Creating a std::vector<float> from a track with the variables that the network has been trained on
+  - Input:
+    -- collisions_it    const C&            ;   An iterator of a collisions table of the form: soa::Join<aod::Collisions, aod::Mults>::iterator const& collision
+    -- track:           const T&            ;   A track, typically from soa::Join<...> tables or of their iterators;
+    -- id:              uint8_t             ;   The id of a particle used for the mass assignment with o2::track::pid_constants::sMasses[id];
+  - Output:
+    -- inputValues:     std::vector<float>  ;   A std::vector<float> with the input variables for the network;
+  */
+
+  const float p = track.tpcInnerParam();
+  const float tgl = track.tgl();
+  const float signed1Pt = track.signed1Pt();
+  const float mass = o2::track::pid_constants::sMasses[id];
+  const float multTPC = collision_it.multTPC() / 11000.;
+  const float ncl = std::sqrt(159. / track.tpcNClsFound());
+
+  return {p, tgl, signed1Pt, mass, multTPC, ncl};
+}
+
+std::vector<Ort::Value> Network::createTensor(std::array<float, 6> input) const
+{
+
+  /*
+  Function: Creating a std::vector<Ort::Value> from a std::vector<float>
+  - Input:
+    -- input:        std::vector<float>       ; The vector from which the tensor should be created;
+  - Output:
+    -- inputValues:  std::vector<Ort::Value>  ; An ONNX tensor;
+  */
+
+  int64_t size = input.size();
+  std::vector<int64_t> input_shape{size / mInputShapes[0][1], mInputShapes[0][1]};
+  std::vector<Ort::Value> inputTensors;
+  inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor<float>(input.data(), size, input_shape));
+
+  return inputTensors;
+
+} // function std::vector<Ort::Value> Network::createTensor(std::vector<float>)
+
+float* Network::evalNetwork(std::vector<Ort::Value> input)
+{
+
+  /*
+  Function: Evaluating the network for a std::vector<Ort::Value>
+  - Input:
+    -- input:         std::vector<Ort::Value>     ; The tensor which should be evaluated by the network;
+  - Output:
+    -- output_values: const float*                ; A float array which can be indexed;
+  */
+
+  try {
+    LOG(debug) << "Shape of input (tensor): " << printShape(input[0].GetTensorTypeAndShapeInfo().GetShape());
+
+    auto outputTensors = mSession->Run(mInputNames, input, mOutputNames);
+    float* output_values = outputTensors[0].GetTensorMutableData<float>();
+    LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape());
+
+    // std::vector<float> output_vals(std::begin(output_values), std::end(output_values));
+
+    return output_values;
+
+  } catch (const Ort::Exception& exception) {
+    LOG(error) << "Error running model inference: " << exception.what();
+
+    return 0;
+  }
+
+} // function Network::evalNetwork(std::vector<Ort::Value>)
+
+float* Network::evalNetwork(std::vector<float> input)
+{
+
+  /*
+ Function: Evaluating the network for a std::vector<float>
+ - Input:
+   -- input:         std::vector<float>    ; The vector which should be evaluated by the network.
+                                             The network will evaluate n inputs where n = input.size()/input_nodes with input_nodes = #(input neurons);
+ - Output:
+   -- output_values: const float*          ; A float array which can be indexed;
+ */
+
+  int64_t size = input.size();
+  std::vector<int64_t> input_shape{size / mInputShapes[0][1], mInputShapes[0][1]};
+  std::vector<Ort::Value> inputTensors;
+  inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor<float>(input.data(), size, input_shape));
+
+  try {
+
+    LOG(debug) << "Shape of input (vector): " << printShape(input_shape);
+    auto outputTensors = mSession->Run(mInputNames, inputTensors, mOutputNames);
+    LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape());
+    float* output_values = outputTensors[0].GetTensorMutableData<float>();
+
+    return output_values;
+
+  } catch (const Ort::Exception& exception) {
+
+    LOG(error) << "Error running model inference: " << exception.what();
+
+    return 0;
+  }
+
+} // function Network::evalNetwork(std::vector<float>)
+
+} // namespace o2::pid::tpc