From c771513b6ca7058c310e8f8f1847909e8a714d46 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Tue, 20 Aug 2024 17:48:31 +0200
Subject: [PATCH 01/13] use dynamic columns from JSON file instead of
 hardcoding them

---
 Tools/PIDML/KaonPidTask.cxx                 |   8 +-
 Tools/PIDML/pidMLBatchEffAndPurProducer.cxx |   2 +-
 Tools/PIDML/pidMLEffAndPurProducer.cxx      |   5 +-
 Tools/PIDML/pidOnnxInterface.h              |  12 +-
 Tools/PIDML/pidOnnxModel.h                  | 173 +++++++++++++++-----
 Tools/PIDML/qaPidML.cxx                     |  12 +-
 Tools/PIDML/simpleApplyPidOnnxInterface.cxx |   4 +-
 Tools/PIDML/simpleApplyPidOnnxModel.cxx     |   4 +-
 8 files changed, 154 insertions(+), 66 deletions(-)

diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index acc29e47bee..baff6bf1178 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -92,26 +92,26 @@ struct KaonPidTask {
     histos.add("hdEdXvsMomentum", ";P_{K^{+}K^{-}}; dE/dx in TPC (keV/cm)", kTH2F, {{100, 0., 4.}, {200, 20., 400.}});
   }
 
-  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& /*tracks*/)
+  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& tracks)
   {
     auto groupPositive = positive->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
     auto groupNegative = negative->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
     for (auto track : groupPositive) {
       histos.fill(HIST("hChargePos"), track.sign());
-      if (pidModel.get()->applyModelBoolean(track)) {
+      if (pidModel.get()->applyModelBoolean(tracks, track)) {
         histos.fill(HIST("hdEdXvsMomentum"), track.p(), track.tpcSignal());
       }
     }
 
     for (auto track : groupNegative) {
       histos.fill(HIST("hChargeNeg"), track.sign());
-      if (pidModel.get()->applyModelBoolean(track)) {
+      if (pidModel.get()->applyModelBoolean(tracks, track)) {
         histos.fill(HIST("hdEdXvsMomentum"), track.p(), track.tpcSignal());
       }
     }
 
     for (auto& [pos, neg] : combinations(soa::CombinationsFullIndexPolicy(groupPositive, groupNegative))) {
-      if (!(pidModel.get()->applyModelBoolean(pos)) || !(pidModel.get()->applyModelBoolean(neg))) {
+      if (!(pidModel.get()->applyModelBoolean(tracks, pos)) || !(pidModel.get()->applyModelBoolean(tracks, neg))) {
         continue;
       }
 
diff --git a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
index fe4d591e590..62320bc961f 100644
--- a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
@@ -225,7 +225,7 @@ struct PidMlBatchEffAndPurProducer {
           fillTrackedHist(mcPart.pdgCode(), track.pt());
 
           for (size_t i = 0; i < cfgPids.value.size(); ++i) {
-            float mlCertainty = models[i].applyModel(track);
+            float mlCertainty = models[i].applyModel(tracks, track);
             nSigma_t nSigma = getNSigma(track, cfgPids.value[i]);
             bool isMCPid = mcPart.pdgCode() == cfgPids.value[i];
 
diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index 7e626a8d631..bc0ade8ef8d 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -15,6 +15,7 @@
 /// \author Michał Olędzki <m.oledzki@cern.ch>
 /// \author Marek Mytkowski <marek.mytkowski@cern.ch>
 
+#include <algorithm>
 #include <string>
 
 #include "Framework/AnalysisDataModel.h"
@@ -26,6 +27,7 @@
 #include "Tools/PIDML/pidOnnxModel.h"
 #include "pidOnnxModel.h"
 #include "Tools/PIDML/pidUtils.h"
+#include "Tools/PIDML/pidML.h"
 
 using namespace o2;
 using namespace o2::framework;
@@ -168,7 +170,8 @@ struct PidMlEffAndPurProducer {
       if (track.has_mcParticle()) {
         auto mcPart = track.mcParticle();
         if (mcPart.isPhysicalPrimary()) {
-          bool mlAccepted = pidModel.applyModelBoolean(track);
+
+          bool mlAccepted = pidModel.applyModelBoolean(tracks, track);
           nSigma_t nSigma = GetNSigma(track);
           bool nSigmaAccepted = IsNSigmaAccept(track, nSigma);
 
diff --git a/Tools/PIDML/pidOnnxInterface.h b/Tools/PIDML/pidOnnxInterface.h
index 03e7bc19146..35fdb8a461c 100644
--- a/Tools/PIDML/pidOnnxInterface.h
+++ b/Tools/PIDML/pidOnnxInterface.h
@@ -78,24 +78,24 @@ struct PidONNXInterface {
   PidONNXInterface& operator=(const PidONNXInterface&) = delete;
   ~PidONNXInterface() = default;
 
-  template <typename T>
-  float applyModel(const T& track, int pid)
+  template <typename Tb, typename T>
+  float applyModel(const Tb& table, const T& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
-        return mModels[i].applyModel(track);
+        return mModels[i].applyModel(table, track);
       }
     }
     LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid;
     return -1.0f;
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track, int pid)
+  template <typename Tb, typename T>
+  bool applyModelBoolean(const Tb& table, const T& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
-        return mModels[i].applyModelBoolean(track);
+        return mModels[i].applyModelBoolean(table, track);
       }
     }
     LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid;
diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 207bc18cd8e..db702359118 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -19,10 +19,12 @@
 
 #include <array>
 #include <cstdint>
+#include <cstring>
 #include <limits>
 #include <string>
 #include <algorithm>
 #include <map>
+#include <type_traits>
 #include <utility>
 #include <memory>
 #include <vector>
@@ -32,10 +34,15 @@
 #include <onnxruntime_cxx_api.h>
 #endif
 
+#include "Framework/TableBuilder.h"
+#include "Framework/Expressions.h"
+#include "arrow/table.h"
+
 #include "rapidjson/document.h"
 #include "rapidjson/filereadstream.h"
 #include "CCDB/CcdbApi.h"
 #include "Tools/PIDML/pidUtils.h"
+#include "Common/DataModel/PIDResponse.h"
 
 using namespace pidml::pidutils;
 
@@ -73,6 +80,7 @@ bool readJsonFile(const std::string& config, rapidjson::Document& d)
 }
 } // namespace
 
+
 struct PidONNXModel {
  public:
   PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp,
@@ -135,16 +143,42 @@ struct PidONNXModel {
   PidONNXModel& operator=(const PidONNXModel&) = delete;
   ~PidONNXModel() = default;
 
-  template <typename T>
-  float applyModel(const T& track)
+  template <typename Tb, typename T>
+  float applyModel(const Tb& table, const T& track)
+  {
+    return getModelOutput(table, track);
+  }
+
+  template <typename Tb, typename T>
+  bool applyModelBoolean(const Tb& table, const T& track)
+  {
+    return getModelOutput(table, track) >= mMinCertainty;
+  }
+
+  template <typename Tb>
+  std::vector<float> batchApplyModel(const Tb& table)
   {
-    return getModelOutput(track);
+    std::vector<float> outputs;
+    outputs.reserve(table.size());
+
+    for (const auto& track : table) {
+      outputs.push_back(applyModel(table, track));
+    }
+
+    return outputs;
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track)
+  template <typename Tb>
+  std::vector<bool> batchApplyModelBoolean(const Tb& table)
   {
-    return getModelOutput(track) >= mMinCertainty;
+    std::vector<bool> outputs;
+    outputs.reserve(table.size());
+
+    for (const auto& track : table) {
+      outputs.push_back(applyModelBoolean(table, track));
+    }
+
+    return outputs;
   }
 
   int mPid;
@@ -213,52 +247,103 @@ struct PidONNXModel {
     }
   }
 
-  template <typename T>
-  std::vector<float> createInputsSingle(const T& track)
-  {
-    // TODO: Hardcoded for now. Planning to implement RowView extension to get runtime access to selected columns
-    // sign is short, trackType and tpcNClsShared uint8_t
+  template <typename... P1, typename ...P2>
+  static constexpr bool is_equal_size(o2::framework::pack<P1...>, o2::framework::pack<P2...>) {
+    if constexpr (sizeof...(P1) == sizeof...(P2)) {
+      return true;
+    }
 
-    float scaledTPCSignal = (track.tpcSignal() - mScalingParams.at("fTPCSignal").first) / mScalingParams.at("fTPCSignal").second;
+    return false;
+  }
 
-    std::vector<float> inputValues{scaledTPCSignal};
+  static float scale(float value, const std::pair<float, float>& scalingParams) {
+    return (value - scalingParams.first) / scalingParams.second;
+  }
 
-    // When TRD Signal shouldn't be used we pass quiet_NaNs to the network
-    if (!inPLimit(track, mPLimits[kTPCTOFTRD]) || trdMissing(track)) {
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-    } else {
-      float scaledTRDSignal = (track.trdSignal() - mScalingParams.at("fTRDSignal").first) / mScalingParams.at("fTRDSignal").second;
-      inputValues.push_back(scaledTRDSignal);
-      inputValues.push_back(track.trdPattern());
-    }
+  template <typename T, typename C>
+  typename C::type getPersistentValue(arrow::Table* table, const T& rowIterator)
+  {
+    auto colIterator = static_cast<C>(rowIterator).getIterator();
+    uint64_t ci = colIterator.mCurrentChunk;
+    uint64_t ai = *(colIterator.mCurrentPos) - colIterator.mFirstIndex;
 
-    // When TOF Signal shouldn't be used we pass quiet_NaNs to the network
-    if (!inPLimit(track, mPLimits[kTPCTOF]) || tofMissing(track)) {
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-      inputValues.push_back(std::numeric_limits<float>::quiet_NaN());
-    } else {
-      float scaledTOFSignal = (track.tofSignal() - mScalingParams.at("fTOFSignal").first) / mScalingParams.at("fTOFSignal").second;
-      float scaledBeta = (track.beta() - mScalingParams.at("fBeta").first) / mScalingParams.at("fBeta").second;
-      inputValues.push_back(scaledTOFSignal);
-      inputValues.push_back(scaledBeta);
-    }
+    return std::static_pointer_cast<o2::soa::arrow_array_for_t<typename C::type>>(o2::soa::getIndexFromLabel(table, C::columnLabel())->chunk(ci))->raw_values()[ai];
+  }
 
-    float scaledX = (track.x() - mScalingParams.at("fX").first) / mScalingParams.at("fX").second;
-    float scaledY = (track.y() - mScalingParams.at("fY").first) / mScalingParams.at("fY").second;
-    float scaledZ = (track.z() - mScalingParams.at("fZ").first) / mScalingParams.at("fZ").second;
-    float scaledAlpha = (track.alpha() - mScalingParams.at("fAlpha").first) / mScalingParams.at("fAlpha").second;
-    float scaledTPCNClsShared = (static_cast<float>(track.tpcNClsShared()) - mScalingParams.at("fTPCNClsShared").first) / mScalingParams.at("fTPCNClsShared").second;
-    float scaledDcaXY = (track.dcaXY() - mScalingParams.at("fDcaXY").first) / mScalingParams.at("fDcaXY").second;
-    float scaledDcaZ = (track.dcaZ() - mScalingParams.at("fDcaZ").first) / mScalingParams.at("fDcaZ").second;
+  template <typename T, typename Tb,  typename... C>
+  std::vector<float> getValues(o2::framework::pack<C...>, const T& track, const Tb& table)
+  {
+    auto arrowTable = table.asArrowTable();
+    std::vector<float> output;
+    output.reserve(mTrainColumns.size());
+    for (const std::string& columnLabel : mTrainColumns) {
+      std::optional<std::pair<float, float>> scalingParams = std::nullopt;
+
+      auto scalingParamsEntry = mScalingParams.find(columnLabel);
+      if(scalingParamsEntry != mScalingParams.end()) {
+        scalingParams = scalingParamsEntry->second;
+      }
 
-    inputValues.insert(inputValues.end(), {track.p(), track.pt(), track.px(), track.py(), track.pz(), static_cast<float>(track.sign()), scaledX, scaledY, scaledZ, scaledAlpha, static_cast<float>(track.trackType()), scaledTPCNClsShared, scaledDcaXY, scaledDcaZ});
+      bool isInPLimitTrd = inPLimit(track, mPLimits[kTPCTOFTRD]);
+      bool isInPLimitTof = inPLimit(track, mPLimits[kTPCTOF]);
+      bool isTrdMissing = trdMissing(track);
+      bool isTofMissing = tofMissing(track);
+
+      ([&]() {
+        if constexpr (o2::soa::is_dynamic_v<C> && std::is_arithmetic_v<typename C::type>) {
+          // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
+          if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
+            std::string label = C::columnLabel();
+
+            // dynamic columns do not have "f" prefix in columnLabel() return string
+            if (std::strcmp(&columnLabel[1], label.data())) {
+              return;
+            }
+
+            float value = static_cast<float>(track.template getDynamicColumn<C>());
+
+            if(scalingParams) {
+              value = scale(value, scalingParams.value());
+            }
+
+            output.push_back(value);
+          }
+        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type> && !std::is_same_v<typename C::type, bool>) {
+          std::string label = C::columnLabel();
+
+          if (columnLabel != label) {
+            return;
+          }
+
+          if constexpr (std::is_same_v<C, o2::aod::track::TRDSignal> || std::is_same_v<C, o2::aod::track::TRDPattern>) {
+            if(isTrdMissing || !isInPLimitTrd) {
+              output.push_back(std::numeric_limits<float>::quiet_NaN());
+              return;
+            }
+          } else if constexpr (std::is_same_v<C, o2::aod::pidtofsignal::TOFSignal> || std::is_same_v<C, o2::aod::pidtofbeta::Beta>) {
+            if(isTofMissing || !isInPLimitTof) {
+              output.push_back(std::numeric_limits<float>::quiet_NaN());
+              return;
+            }
+          }
+
+          float value = static_cast<float>(getPersistentValue<T, C>(arrowTable.get(), track));
+
+          if(scalingParams) {
+            value = scale(value, scalingParams.value());
+          }
+
+          output.push_back(value);
+        }
+      }(),
+      ...);
+    }
 
-    return inputValues;
+    return output;
   }
 
-  template <typename T>
-  float getModelOutput(const T& track)
+  template <typename Tb, typename T>
+  float getModelOutput(const Tb& table, const T& track)
   {
     // First rank of the expected model input is -1 which means that it is dynamic axis.
     // Axis is exported as dynamic to make it possible to run model inference with the batch of
@@ -268,7 +353,7 @@ struct PidONNXModel {
     auto input_shape = mInputShapes[0];
     input_shape[0] = batch_size;
 
-    std::vector<float> inputTensorValues = createInputsSingle(track);
+    std::vector<float> inputTensorValues = getValues(typename Tb::table_t::columns{}, track, table);
     std::vector<Ort::Value> inputTensors;
 
 #if __has_include(<onnxruntime/core/session/onnxruntime_cxx_api.h>)
diff --git a/Tools/PIDML/qaPidML.cxx b/Tools/PIDML/qaPidML.cxx
index 54cafc437ab..b5eef17952d 100644
--- a/Tools/PIDML/qaPidML.cxx
+++ b/Tools/PIDML/qaPidML.cxx
@@ -324,13 +324,13 @@ struct QaPidML {
     }
   }
 
-  template <std::size_t i, typename T>
-  void pidML(const T& track, const int pdgCodeMC)
+  template <std::size_t i, typename Tb, typename T>
+  void pidML(const Tb& tracks, const T& track, const int pdgCodeMC)
   {
     float pidCertainties[3];
-    pidCertainties[0] = model211.applyModel(track);
-    pidCertainties[1] = model2212.applyModel(track);
-    pidCertainties[2] = model321.applyModel(track);
+    pidCertainties[0] = model211.applyModel(tracks, track);
+    pidCertainties[1] = model2212.applyModel(tracks, track);
+    pidCertainties[2] = model321.applyModel(tracks, track);
     int pid = getParticlePdg(pidCertainties);
     // condition for sign: we want to work only with pi, p and K, without antiparticles
     if (pid == particlesPdgCode[i] && track.sign() == 1) {
@@ -385,7 +385,7 @@ struct QaPidML {
 
       // only 3 particles can be predicted by model
       static_for<0, 2>([&](auto i) {
-        pidML<i>(track, pdgCodeMC);
+        pidML<i>(tracks, track, pdgCodeMC);
       });
     }
   }
diff --git a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
index 82f644103ed..4b75d8eb222 100644
--- a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
@@ -84,7 +84,7 @@ struct SimpleApplyOnnxInterface {
 
     for (auto& track : tracks) {
       for (int pid : cfgPids.value) {
-        bool accepted = pidInterface.applyModelBoolean(track, pid);
+        bool accepted = pidInterface.applyModelBoolean(tracks, track, pid);
         LOGF(info, "collision id: %d track id: %d pid: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
              track.collisionId(), track.index(), pid, accepted, track.p(), track.x(), track.y(), track.z());
         pidMLResults(track.index(), pid, accepted);
@@ -97,7 +97,7 @@ struct SimpleApplyOnnxInterface {
   {
     for (auto& track : tracks) {
       for (int pid : cfgPids.value) {
-        bool accepted = pidInterface.applyModelBoolean(track, pid);
+        bool accepted = pidInterface.applyModelBoolean(tracks, track, pid);
         LOGF(info, "collision id: %d track id: %d pid: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
              track.collisionId(), track.index(), pid, accepted, track.p(), track.x(), track.y(), track.z());
         pidMLResults(track.index(), pid, accepted);
diff --git a/Tools/PIDML/simpleApplyPidOnnxModel.cxx b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
index e261c1e10bc..bda913bedf4 100644
--- a/Tools/PIDML/simpleApplyPidOnnxModel.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
@@ -80,7 +80,7 @@ struct SimpleApplyOnnxModel {
     }
 
     for (auto& track : tracks) {
-      bool accepted = pidModel.applyModelBoolean(track);
+      bool accepted = pidModel.applyModelBoolean(tracks, track);
       LOGF(info, "collision id: %d track id: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
            track.collisionId(), track.index(), accepted, track.p(), track.x(), track.y(), track.z());
       pidMLResults(track.index(), cfgPid.value, accepted);
@@ -91,7 +91,7 @@ struct SimpleApplyOnnxModel {
   void processTracksOnly(BigTracks const& tracks)
   {
     for (auto& track : tracks) {
-      bool accepted = pidModel.applyModelBoolean(track);
+      bool accepted = pidModel.applyModelBoolean(tracks, track);
       LOGF(info, "collision id: %d track id: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
            track.collisionId(), track.index(), accepted, track.p(), track.x(), track.y(), track.z());
       pidMLResults(track.index(), cfgPid.value, accepted);

From 6ce8bb08d8a2df6c130181df08f2b29303c8930a Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Tue, 20 Aug 2024 17:49:47 +0200
Subject: [PATCH 02/13] formatting

---
 Tools/PIDML/pidOnnxModel.h | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index db702359118..4d692e23378 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -80,7 +80,6 @@ bool readJsonFile(const std::string& config, rapidjson::Document& d)
 }
 } // namespace
 
-
 struct PidONNXModel {
  public:
   PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp,
@@ -247,8 +246,9 @@ struct PidONNXModel {
     }
   }
 
-  template <typename... P1, typename ...P2>
-  static constexpr bool is_equal_size(o2::framework::pack<P1...>, o2::framework::pack<P2...>) {
+  template <typename... P1, typename... P2>
+  static constexpr bool is_equal_size(o2::framework::pack<P1...>, o2::framework::pack<P2...>)
+  {
     if constexpr (sizeof...(P1) == sizeof...(P2)) {
       return true;
     }
@@ -256,7 +256,8 @@ struct PidONNXModel {
     return false;
   }
 
-  static float scale(float value, const std::pair<float, float>& scalingParams) {
+  static float scale(float value, const std::pair<float, float>& scalingParams)
+  {
     return (value - scalingParams.first) / scalingParams.second;
   }
 
@@ -270,7 +271,7 @@ struct PidONNXModel {
     return std::static_pointer_cast<o2::soa::arrow_array_for_t<typename C::type>>(o2::soa::getIndexFromLabel(table, C::columnLabel())->chunk(ci))->raw_values()[ai];
   }
 
-  template <typename T, typename Tb,  typename... C>
+  template <typename T, typename Tb, typename... C>
   std::vector<float> getValues(o2::framework::pack<C...>, const T& track, const Tb& table)
   {
     auto arrowTable = table.asArrowTable();
@@ -280,7 +281,7 @@ struct PidONNXModel {
       std::optional<std::pair<float, float>> scalingParams = std::nullopt;
 
       auto scalingParamsEntry = mScalingParams.find(columnLabel);
-      if(scalingParamsEntry != mScalingParams.end()) {
+      if (scalingParamsEntry != mScalingParams.end()) {
         scalingParams = scalingParamsEntry->second;
       }
 
@@ -302,7 +303,7 @@ struct PidONNXModel {
 
             float value = static_cast<float>(track.template getDynamicColumn<C>());
 
-            if(scalingParams) {
+            if (scalingParams) {
               value = scale(value, scalingParams.value());
             }
 
@@ -316,12 +317,12 @@ struct PidONNXModel {
           }
 
           if constexpr (std::is_same_v<C, o2::aod::track::TRDSignal> || std::is_same_v<C, o2::aod::track::TRDPattern>) {
-            if(isTrdMissing || !isInPLimitTrd) {
+            if (isTrdMissing || !isInPLimitTrd) {
               output.push_back(std::numeric_limits<float>::quiet_NaN());
               return;
             }
           } else if constexpr (std::is_same_v<C, o2::aod::pidtofsignal::TOFSignal> || std::is_same_v<C, o2::aod::pidtofbeta::Beta>) {
-            if(isTofMissing || !isInPLimitTof) {
+            if (isTofMissing || !isInPLimitTof) {
               output.push_back(std::numeric_limits<float>::quiet_NaN());
               return;
             }
@@ -329,14 +330,14 @@ struct PidONNXModel {
 
           float value = static_cast<float>(getPersistentValue<T, C>(arrowTable.get(), track));
 
-          if(scalingParams) {
+          if (scalingParams) {
             value = scale(value, scalingParams.value());
           }
 
           output.push_back(value);
         }
       }(),
-      ...);
+       ...);
     }
 
     return output;

From 82ebb922cb9fd279988bb4d074938a3cd7c4c910 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 09:00:53 +0200
Subject: [PATCH 03/13] add missing includes

---
 Tools/PIDML/pidML.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Tools/PIDML/pidML.h b/Tools/PIDML/pidML.h
index 821da7efc42..6629b88475b 100644
--- a/Tools/PIDML/pidML.h
+++ b/Tools/PIDML/pidML.h
@@ -19,6 +19,8 @@
 
 #include "Framework/AnalysisDataModel.h"
 #include "Common/DataModel/PIDResponse.h"
+#include "Common/DataModel/Centrality.h"
+#include "Common/DataModel/Multiplicity.h"
 
 namespace o2::aod
 {

From 17e007e58d91d327dc1467c304c142c4308463bd Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 11:23:36 +0200
Subject: [PATCH 04/13] use getIterator - no more passing of table

---
 Tools/PIDML/KaonPidTask.cxx                 |  6 +--
 Tools/PIDML/pidMLBatchEffAndPurProducer.cxx |  2 +-
 Tools/PIDML/pidMLEffAndPurProducer.cxx      |  2 +-
 Tools/PIDML/pidOnnxInterface.h              | 12 +++---
 Tools/PIDML/pidOnnxModel.h                  | 41 +++++++++------------
 Tools/PIDML/qaPidML.cxx                     | 12 +++---
 Tools/PIDML/simpleApplyPidOnnxInterface.cxx |  4 +-
 Tools/PIDML/simpleApplyPidOnnxModel.cxx     |  4 +-
 8 files changed, 39 insertions(+), 44 deletions(-)

diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index baff6bf1178..bb1f0fe1551 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -98,20 +98,20 @@ struct KaonPidTask {
     auto groupNegative = negative->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
     for (auto track : groupPositive) {
       histos.fill(HIST("hChargePos"), track.sign());
-      if (pidModel.get()->applyModelBoolean(tracks, track)) {
+      if (pidModel.get()->applyModelBoolean(track)) {
         histos.fill(HIST("hdEdXvsMomentum"), track.p(), track.tpcSignal());
       }
     }
 
     for (auto track : groupNegative) {
       histos.fill(HIST("hChargeNeg"), track.sign());
-      if (pidModel.get()->applyModelBoolean(tracks, track)) {
+      if (pidModel.get()->applyModelBoolean(track)) {
         histos.fill(HIST("hdEdXvsMomentum"), track.p(), track.tpcSignal());
       }
     }
 
     for (auto& [pos, neg] : combinations(soa::CombinationsFullIndexPolicy(groupPositive, groupNegative))) {
-      if (!(pidModel.get()->applyModelBoolean(tracks, pos)) || !(pidModel.get()->applyModelBoolean(tracks, neg))) {
+      if (!(pidModel.get()->applyModelBoolean(pos)) || !(pidModel.get()->applyModelBoolean(neg))) {
         continue;
       }
 
diff --git a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
index 62320bc961f..fe4d591e590 100644
--- a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
@@ -225,7 +225,7 @@ struct PidMlBatchEffAndPurProducer {
           fillTrackedHist(mcPart.pdgCode(), track.pt());
 
           for (size_t i = 0; i < cfgPids.value.size(); ++i) {
-            float mlCertainty = models[i].applyModel(tracks, track);
+            float mlCertainty = models[i].applyModel(track);
             nSigma_t nSigma = getNSigma(track, cfgPids.value[i]);
             bool isMCPid = mcPart.pdgCode() == cfgPids.value[i];
 
diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index bc0ade8ef8d..74bbc380509 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -171,7 +171,7 @@ struct PidMlEffAndPurProducer {
         auto mcPart = track.mcParticle();
         if (mcPart.isPhysicalPrimary()) {
 
-          bool mlAccepted = pidModel.applyModelBoolean(tracks, track);
+          bool mlAccepted = pidModel.applyModelBoolean(track);
           nSigma_t nSigma = GetNSigma(track);
           bool nSigmaAccepted = IsNSigmaAccept(track, nSigma);
 
diff --git a/Tools/PIDML/pidOnnxInterface.h b/Tools/PIDML/pidOnnxInterface.h
index 35fdb8a461c..03e7bc19146 100644
--- a/Tools/PIDML/pidOnnxInterface.h
+++ b/Tools/PIDML/pidOnnxInterface.h
@@ -78,24 +78,24 @@ struct PidONNXInterface {
   PidONNXInterface& operator=(const PidONNXInterface&) = delete;
   ~PidONNXInterface() = default;
 
-  template <typename Tb, typename T>
-  float applyModel(const Tb& table, const T& track, int pid)
+  template <typename T>
+  float applyModel(const T& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
-        return mModels[i].applyModel(table, track);
+        return mModels[i].applyModel(track);
       }
     }
     LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid;
     return -1.0f;
   }
 
-  template <typename Tb, typename T>
-  bool applyModelBoolean(const Tb& table, const T& track, int pid)
+  template <typename T>
+  bool applyModelBoolean(const T& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
-        return mModels[i].applyModelBoolean(table, track);
+        return mModels[i].applyModelBoolean(track);
       }
     }
     LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid;
diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 4d692e23378..c0a2530b520 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -142,16 +142,16 @@ struct PidONNXModel {
   PidONNXModel& operator=(const PidONNXModel&) = delete;
   ~PidONNXModel() = default;
 
-  template <typename Tb, typename T>
-  float applyModel(const Tb& table, const T& track)
+  template <typename T>
+  float applyModel(const T& track)
   {
-    return getModelOutput(table, track);
+    return getModelOutput(track);
   }
 
-  template <typename Tb, typename T>
-  bool applyModelBoolean(const Tb& table, const T& track)
+  template <typename T>
+  bool applyModelBoolean(const T& track)
   {
-    return getModelOutput(table, track) >= mMinCertainty;
+    return getModelOutput(track) >= mMinCertainty;
   }
 
   template <typename Tb>
@@ -161,7 +161,7 @@ struct PidONNXModel {
     outputs.reserve(table.size());
 
     for (const auto& track : table) {
-      outputs.push_back(applyModel(table, track));
+      outputs.push_back(applyModel(track));
     }
 
     return outputs;
@@ -174,7 +174,7 @@ struct PidONNXModel {
     outputs.reserve(table.size());
 
     for (const auto& track : table) {
-      outputs.push_back(applyModelBoolean(table, track));
+      outputs.push_back(applyModelBoolean(track));
     }
 
     return outputs;
@@ -261,20 +261,15 @@ struct PidONNXModel {
     return (value - scalingParams.first) / scalingParams.second;
   }
 
-  template <typename T, typename C>
-  typename C::type getPersistentValue(arrow::Table* table, const T& rowIterator)
+  template <typename C, typename T>
+  typename C::type getPersistentValue(const T& rowIterator)
   {
-    auto colIterator = static_cast<C>(rowIterator).getIterator();
-    uint64_t ci = colIterator.mCurrentChunk;
-    uint64_t ai = *(colIterator.mCurrentPos) - colIterator.mFirstIndex;
-
-    return std::static_pointer_cast<o2::soa::arrow_array_for_t<typename C::type>>(o2::soa::getIndexFromLabel(table, C::columnLabel())->chunk(ci))->raw_values()[ai];
+    return *(static_cast<C>(rowIterator).getIterator());
   }
 
-  template <typename T, typename Tb, typename... C>
-  std::vector<float> getValues(o2::framework::pack<C...>, const T& track, const Tb& table)
+  template <typename T, typename... C>
+  std::vector<float> getValues(o2::framework::pack<C...>, const T& track)
   {
-    auto arrowTable = table.asArrowTable();
     std::vector<float> output;
     output.reserve(mTrainColumns.size());
     for (const std::string& columnLabel : mTrainColumns) {
@@ -309,7 +304,7 @@ struct PidONNXModel {
 
             output.push_back(value);
           }
-        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type> && !std::is_same_v<typename C::type, bool>) {
+        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type>) {
           std::string label = C::columnLabel();
 
           if (columnLabel != label) {
@@ -328,7 +323,7 @@ struct PidONNXModel {
             }
           }
 
-          float value = static_cast<float>(getPersistentValue<T, C>(arrowTable.get(), track));
+          float value = static_cast<float>(getPersistentValue<C>(track));
 
           if (scalingParams) {
             value = scale(value, scalingParams.value());
@@ -343,8 +338,8 @@ struct PidONNXModel {
     return output;
   }
 
-  template <typename Tb, typename T>
-  float getModelOutput(const Tb& table, const T& track)
+  template <typename T>
+  float getModelOutput(const T& track)
   {
     // First rank of the expected model input is -1 which means that it is dynamic axis.
     // Axis is exported as dynamic to make it possible to run model inference with the batch of
@@ -354,7 +349,7 @@ struct PidONNXModel {
     auto input_shape = mInputShapes[0];
     input_shape[0] = batch_size;
 
-    std::vector<float> inputTensorValues = getValues(typename Tb::table_t::columns{}, track, table);
+    std::vector<float> inputTensorValues = getValues(typename T::parent_t::table_t::columns{}, track);
     std::vector<Ort::Value> inputTensors;
 
 #if __has_include(<onnxruntime/core/session/onnxruntime_cxx_api.h>)
diff --git a/Tools/PIDML/qaPidML.cxx b/Tools/PIDML/qaPidML.cxx
index b5eef17952d..54cafc437ab 100644
--- a/Tools/PIDML/qaPidML.cxx
+++ b/Tools/PIDML/qaPidML.cxx
@@ -324,13 +324,13 @@ struct QaPidML {
     }
   }
 
-  template <std::size_t i, typename Tb, typename T>
-  void pidML(const Tb& tracks, const T& track, const int pdgCodeMC)
+  template <std::size_t i, typename T>
+  void pidML(const T& track, const int pdgCodeMC)
   {
     float pidCertainties[3];
-    pidCertainties[0] = model211.applyModel(tracks, track);
-    pidCertainties[1] = model2212.applyModel(tracks, track);
-    pidCertainties[2] = model321.applyModel(tracks, track);
+    pidCertainties[0] = model211.applyModel(track);
+    pidCertainties[1] = model2212.applyModel(track);
+    pidCertainties[2] = model321.applyModel(track);
     int pid = getParticlePdg(pidCertainties);
     // condition for sign: we want to work only with pi, p and K, without antiparticles
     if (pid == particlesPdgCode[i] && track.sign() == 1) {
@@ -385,7 +385,7 @@ struct QaPidML {
 
       // only 3 particles can be predicted by model
       static_for<0, 2>([&](auto i) {
-        pidML<i>(tracks, track, pdgCodeMC);
+        pidML<i>(track, pdgCodeMC);
       });
     }
   }
diff --git a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
index 4b75d8eb222..82f644103ed 100644
--- a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
@@ -84,7 +84,7 @@ struct SimpleApplyOnnxInterface {
 
     for (auto& track : tracks) {
       for (int pid : cfgPids.value) {
-        bool accepted = pidInterface.applyModelBoolean(tracks, track, pid);
+        bool accepted = pidInterface.applyModelBoolean(track, pid);
         LOGF(info, "collision id: %d track id: %d pid: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
              track.collisionId(), track.index(), pid, accepted, track.p(), track.x(), track.y(), track.z());
         pidMLResults(track.index(), pid, accepted);
@@ -97,7 +97,7 @@ struct SimpleApplyOnnxInterface {
   {
     for (auto& track : tracks) {
       for (int pid : cfgPids.value) {
-        bool accepted = pidInterface.applyModelBoolean(tracks, track, pid);
+        bool accepted = pidInterface.applyModelBoolean(track, pid);
         LOGF(info, "collision id: %d track id: %d pid: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
              track.collisionId(), track.index(), pid, accepted, track.p(), track.x(), track.y(), track.z());
         pidMLResults(track.index(), pid, accepted);
diff --git a/Tools/PIDML/simpleApplyPidOnnxModel.cxx b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
index bda913bedf4..e261c1e10bc 100644
--- a/Tools/PIDML/simpleApplyPidOnnxModel.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
@@ -80,7 +80,7 @@ struct SimpleApplyOnnxModel {
     }
 
     for (auto& track : tracks) {
-      bool accepted = pidModel.applyModelBoolean(tracks, track);
+      bool accepted = pidModel.applyModelBoolean(track);
       LOGF(info, "collision id: %d track id: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
            track.collisionId(), track.index(), accepted, track.p(), track.x(), track.y(), track.z());
       pidMLResults(track.index(), cfgPid.value, accepted);
@@ -91,7 +91,7 @@ struct SimpleApplyOnnxModel {
   void processTracksOnly(BigTracks const& tracks)
   {
     for (auto& track : tracks) {
-      bool accepted = pidModel.applyModelBoolean(tracks, track);
+      bool accepted = pidModel.applyModelBoolean(track);
       LOGF(info, "collision id: %d track id: %d accepted: %d p: %.3f; x: %.3f, y: %.3f, z: %.3f",
            track.collisionId(), track.index(), accepted, track.p(), track.x(), track.y(), track.z());
       pidMLResults(track.index(), cfgPid.value, accepted);

From 6524f6dffddc36ef9d29d7f42b8859db110ba077 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 11:24:35 +0200
Subject: [PATCH 05/13] comment out tracks table in KaonPidTask

---
 Tools/PIDML/KaonPidTask.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index bb1f0fe1551..4c28af7368f 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -92,7 +92,7 @@ struct KaonPidTask {
     histos.add("hdEdXvsMomentum", ";P_{K^{+}K^{-}}; dE/dx in TPC (keV/cm)", kTH2F, {{100, 0., 4.}, {200, 20., 400.}});
   }
 
-  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& tracks)
+  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& /* tracks */)
   {
     auto groupPositive = positive->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
     auto groupNegative = negative->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);

From cde7a0afa25b63920b42973d6d0c7ef724586314 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 12:46:19 +0200
Subject: [PATCH 06/13] refactor solution to be more generic

---
 Tools/PIDML/pidOnnxModel.h | 115 ++++++++++++++++++++-----------------
 1 file changed, 63 insertions(+), 52 deletions(-)

diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index c0a2530b520..672c2f89baf 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <cstring>
 #include <limits>
+#include <optional>
 #include <string>
 #include <algorithm>
 #include <map>
@@ -268,71 +269,81 @@ struct PidONNXModel {
   }
 
   template <typename T, typename... C>
-  std::vector<float> getValues(o2::framework::pack<C...>, const T& track)
+  std::optional<float> getColumnValueByLabel(o2::framework::pack<C...>, const T& rowIterator, const std::string& columnLabel)
   {
-    std::vector<float> output;
-    output.reserve(mTrainColumns.size());
-    for (const std::string& columnLabel : mTrainColumns) {
-      std::optional<std::pair<float, float>> scalingParams = std::nullopt;
+    std::optional<float> value = std::nullopt;
 
-      auto scalingParamsEntry = mScalingParams.find(columnLabel);
-      if (scalingParamsEntry != mScalingParams.end()) {
-        scalingParams = scalingParamsEntry->second;
+    ([&]() {
+      if (value) {
+        return;
       }
 
-      bool isInPLimitTrd = inPLimit(track, mPLimits[kTPCTOFTRD]);
-      bool isInPLimitTof = inPLimit(track, mPLimits[kTPCTOF]);
-      bool isTrdMissing = trdMissing(track);
-      bool isTofMissing = tofMissing(track);
+      if constexpr (o2::soa::is_dynamic_v<C> && std::is_arithmetic_v<typename C::type>) {
+        // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
+        if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
+          std::string label = C::columnLabel();
 
-      ([&]() {
-        if constexpr (o2::soa::is_dynamic_v<C> && std::is_arithmetic_v<typename C::type>) {
-          // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
-          if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
-            std::string label = C::columnLabel();
+          // dynamic columns do not have "f" prefix in columnLabel() return string
+          if (std::strcmp(&columnLabel[1], label.data()) != 0 && columnLabel != label) {
+            return;
+          }
 
-            // dynamic columns do not have "f" prefix in columnLabel() return string
-            if (std::strcmp(&columnLabel[1], label.data())) {
-              return;
-            }
+          value = static_cast<float>(rowIterator.template getDynamicColumn<C>());
+        }
+      } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type>) {
+        if (columnLabel != C::columnLabel()) {
+          return;
+        }
 
-            float value = static_cast<float>(track.template getDynamicColumn<C>());
+        value = static_cast<float>(getPersistentValue<C>(rowIterator));
+      }
+    }(),
+     ...);
 
-            if (scalingParams) {
-              value = scale(value, scalingParams.value());
-            }
+    return value;
+  }
 
-            output.push_back(value);
-          }
-        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type>) {
-          std::string label = C::columnLabel();
+  template <typename T>
+  std::optional<float> getColumnValueByLabel(const T& rowIterator, const std::string& columnLabel)
+  {
+    return getColumnValueByLabel(typename T::parent_t::table_t::columns{}, rowIterator, columnLabel);
+  }
 
-          if (columnLabel != label) {
-            return;
-          }
+  template <typename T>
+  std::vector<float> getValues(const T& track)
+  {
+    std::vector<float> output;
+    output.reserve(mTrainColumns.size());
 
-          if constexpr (std::is_same_v<C, o2::aod::track::TRDSignal> || std::is_same_v<C, o2::aod::track::TRDPattern>) {
-            if (isTrdMissing || !isInPLimitTrd) {
-              output.push_back(std::numeric_limits<float>::quiet_NaN());
-              return;
-            }
-          } else if constexpr (std::is_same_v<C, o2::aod::pidtofsignal::TOFSignal> || std::is_same_v<C, o2::aod::pidtofbeta::Beta>) {
-            if (isTofMissing || !isInPLimitTof) {
-              output.push_back(std::numeric_limits<float>::quiet_NaN());
-              return;
-            }
-          }
+    bool useTOF = !tofMissing(track) && inPLimit(track, mPLimits[kTPCTOF]);
+    bool useTRD = !trdMissing(track) && inPLimit(track, mPLimits[kTPCTOFTRD]);
 
-          float value = static_cast<float>(getPersistentValue<C>(track));
+    for (const std::string& columnLabel : mTrainColumns) {
+      if (
+        ((columnLabel == "fTRDSignal" || columnLabel == "fTRDPattern") && !useTRD) ||
+        ((columnLabel == "fTOFSignal" || columnLabel == "fBeta") && !useTOF)) {
+        output.push_back(std::numeric_limits<float>::quiet_NaN());
+        continue;
+      }
 
-          if (scalingParams) {
-            value = scale(value, scalingParams.value());
-          }
+      std::optional<std::pair<float, float>> scalingParams = std::nullopt;
 
-          output.push_back(value);
-        }
-      }(),
-       ...);
+      auto scalingParamsEntry = mScalingParams.find(columnLabel);
+      if (scalingParamsEntry != mScalingParams.end()) {
+        scalingParams = scalingParamsEntry->second;
+      }
+
+      std::optional<float> value = getColumnValueByLabel(track, columnLabel);
+
+      if (!value) {
+        LOG(fatal) << "Cannot obtain column of name: " << columnLabel;
+      }
+
+      if (scalingParams) {
+        value = scale(value.value(), scalingParams.value());
+      }
+
+      output.push_back(value.value());
     }
 
     return output;
@@ -349,7 +360,7 @@ struct PidONNXModel {
     auto input_shape = mInputShapes[0];
     input_shape[0] = batch_size;
 
-    std::vector<float> inputTensorValues = getValues(typename T::parent_t::table_t::columns{}, track);
+    std::vector<float> inputTensorValues = getValues(track);
     std::vector<Ort::Value> inputTensors;
 
 #if __has_include(<onnxruntime/core/session/onnxruntime_cxx_api.h>)

From 49738747933051b4cc397c106833a85cc9eb26bc Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:23:59 +0200
Subject: [PATCH 07/13] move out common condition

---
 Tools/PIDML/pidOnnxModel.h | 30 +++++++++++++++---------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 672c2f89baf..9f0431ee125 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -23,7 +23,6 @@
 #include <limits>
 #include <optional>
 #include <string>
-#include <algorithm>
 #include <map>
 #include <type_traits>
 #include <utility>
@@ -37,7 +36,6 @@
 
 #include "Framework/TableBuilder.h"
 #include "Framework/Expressions.h"
-#include "arrow/table.h"
 
 #include "rapidjson/document.h"
 #include "rapidjson/filereadstream.h"
@@ -278,24 +276,26 @@ struct PidONNXModel {
         return;
       }
 
-      if constexpr (o2::soa::is_dynamic_v<C> && std::is_arithmetic_v<typename C::type>) {
-        // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
-        if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
-          std::string label = C::columnLabel();
+      if constexpr (std::is_arithmetic_v<typename C::type>) {
+        if constexpr (o2::soa::is_dynamic_v<C>) {
+          // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
+          if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
+            std::string label = C::columnLabel();
 
-          // dynamic columns do not have "f" prefix in columnLabel() return string
-          if (std::strcmp(&columnLabel[1], label.data()) != 0 && columnLabel != label) {
+            // dynamic columns do not have "f" prefix in columnLabel() return string
+            if (std::strcmp(&columnLabel[1], label.data()) != 0 && columnLabel != label) {
+              return;
+            }
+
+            value = static_cast<float>(rowIterator.template getDynamicColumn<C>());
+          }
+        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C>) {
+          if (columnLabel != C::columnLabel()) {
             return;
           }
 
-          value = static_cast<float>(rowIterator.template getDynamicColumn<C>());
+          value = static_cast<float>(getPersistentValue<C>(rowIterator));
         }
-      } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C> && std::is_arithmetic_v<typename C::type>) {
-        if (columnLabel != C::columnLabel()) {
-          return;
-        }
-
-        value = static_cast<float>(getPersistentValue<C>(rowIterator));
       }
     }(),
      ...);

From 77ff0a79df41ccbebfd438773d9f8a3a5bc08f65 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:38:34 +0200
Subject: [PATCH 08/13] remove unused methods from ONNXModel and simplify
 is_equal_size function

---
 Tools/PIDML/KaonPidTask.cxx |  2 +-
 Tools/PIDML/pidOnnxModel.h  | 32 +-------------------------------
 2 files changed, 2 insertions(+), 32 deletions(-)

diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index 4c28af7368f..acc29e47bee 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -92,7 +92,7 @@ struct KaonPidTask {
     histos.add("hdEdXvsMomentum", ";P_{K^{+}K^{-}}; dE/dx in TPC (keV/cm)", kTH2F, {{100, 0., 4.}, {200, 20., 400.}});
   }
 
-  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& /* tracks */)
+  void process(MyFilteredCollision const& coll, o2::aod::MyTracks const& /*tracks*/)
   {
     auto groupPositive = positive->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
     auto groupNegative = negative->sliceByCached(aod::track::collisionId, coll.globalIndex(), cache);
diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 9f0431ee125..c470d9013d3 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -153,32 +153,6 @@ struct PidONNXModel {
     return getModelOutput(track) >= mMinCertainty;
   }
 
-  template <typename Tb>
-  std::vector<float> batchApplyModel(const Tb& table)
-  {
-    std::vector<float> outputs;
-    outputs.reserve(table.size());
-
-    for (const auto& track : table) {
-      outputs.push_back(applyModel(track));
-    }
-
-    return outputs;
-  }
-
-  template <typename Tb>
-  std::vector<bool> batchApplyModelBoolean(const Tb& table)
-  {
-    std::vector<bool> outputs;
-    outputs.reserve(table.size());
-
-    for (const auto& track : table) {
-      outputs.push_back(applyModelBoolean(track));
-    }
-
-    return outputs;
-  }
-
   int mPid;
   double mMinCertainty;
 
@@ -248,11 +222,7 @@ struct PidONNXModel {
   template <typename... P1, typename... P2>
   static constexpr bool is_equal_size(o2::framework::pack<P1...>, o2::framework::pack<P2...>)
   {
-    if constexpr (sizeof...(P1) == sizeof...(P2)) {
-      return true;
-    }
-
-    return false;
+    return sizeof...(P1) == sizeof...(P2);
   }
 
   static float scale(float value, const std::pair<float, float>& scalingParams)

From b8c119d573b3b70c78e18d6548bb8a2b9648b774 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:40:29 +0200
Subject: [PATCH 09/13] tidy up pidEffAndPurProducer

---
 Tools/PIDML/pidMLEffAndPurProducer.cxx | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index 74bbc380509..d7c1ecf54ba 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -25,9 +25,7 @@
 #include "Common/DataModel/TrackSelectionTables.h"
 #include "Common/DataModel/PIDResponse.h"
 #include "Tools/PIDML/pidOnnxModel.h"
-#include "pidOnnxModel.h"
 #include "Tools/PIDML/pidUtils.h"
-#include "Tools/PIDML/pidML.h"
 
 using namespace o2;
 using namespace o2::framework;
@@ -170,7 +168,6 @@ struct PidMlEffAndPurProducer {
       if (track.has_mcParticle()) {
         auto mcPart = track.mcParticle();
         if (mcPart.isPhysicalPrimary()) {
-
           bool mlAccepted = pidModel.applyModelBoolean(track);
           nSigma_t nSigma = GetNSigma(track);
           bool nSigmaAccepted = IsNSigmaAccept(track, nSigma);

From 50c6bd47651032aa24eecdec39fc1cfd5121db0f Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:41:18 +0200
Subject: [PATCH 10/13] remove unused include

---
 Tools/PIDML/pidMLEffAndPurProducer.cxx | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index d7c1ecf54ba..2fcaa7f8477 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -15,7 +15,6 @@
 /// \author Michał Olędzki <m.oledzki@cern.ch>
 /// \author Marek Mytkowski <marek.mytkowski@cern.ch>
 
-#include <algorithm>
 #include <string>
 
 #include "Framework/AnalysisDataModel.h"

From a3408c0b01c490b5979a91253d1431ad38d2a6c2 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:44:36 +0200
Subject: [PATCH 11/13] remove more unused includes

---
 Tools/PIDML/pidOnnxModel.h | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index c470d9013d3..113852ea251 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -34,14 +34,10 @@
 #include <onnxruntime_cxx_api.h>
 #endif
 
-#include "Framework/TableBuilder.h"
-#include "Framework/Expressions.h"
-
 #include "rapidjson/document.h"
 #include "rapidjson/filereadstream.h"
 #include "CCDB/CcdbApi.h"
 #include "Tools/PIDML/pidUtils.h"
-#include "Common/DataModel/PIDResponse.h"
 
 using namespace pidml::pidutils;
 

From 0ee085b8f96b62879d987ec6db99732a2497a56a Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Wed, 21 Aug 2024 14:51:57 +0200
Subject: [PATCH 12/13] add include to ONNXModel

---
 Tools/PIDML/pidOnnxModel.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 113852ea251..40813f55c0c 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -18,6 +18,7 @@
 #define TOOLS_PIDML_PIDONNXMODEL_H_
 
 #include <array>
+#include <algorithm>
 #include <cstdint>
 #include <cstring>
 #include <limits>

From 312a1c51ab9f87e694de39cf1e6e6b90cc5958f8 Mon Sep 17 00:00:00 2001
From: mytkom <marek.mytkowski.mm@gmail.com>
Date: Tue, 8 Oct 2024 09:01:54 +0200
Subject: [PATCH 13/13] use getGetterByLabel (not merged to O2Framework yet) as
 proof of concept

---
 Tools/PIDML/KaonPidTask.cxx                 |  4 +-
 Tools/PIDML/pidMLBatchEffAndPurProducer.cxx |  6 +-
 Tools/PIDML/pidMLEffAndPurProducer.cxx      |  6 +-
 Tools/PIDML/pidOnnxInterface.h              |  9 +--
 Tools/PIDML/pidOnnxModel.h                  | 88 ++++-----------------
 Tools/PIDML/qaPidML.cxx                     | 27 ++++---
 Tools/PIDML/simpleApplyPidOnnxInterface.cxx |  8 +-
 Tools/PIDML/simpleApplyPidOnnxModel.cxx     |  6 +-
 8 files changed, 49 insertions(+), 105 deletions(-)

diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx
index acc29e47bee..3105292d44c 100644
--- a/Tools/PIDML/KaonPidTask.cxx
+++ b/Tools/PIDML/KaonPidTask.cxx
@@ -46,7 +46,7 @@ struct KaonPidTask {
   SliceCache cache;
   Preslice<aod::Tracks> perCol = aod::track::collisionId;
 
-  std::shared_ptr<PidONNXModel> pidModel; // creates a shared pointer to a new instance 'pidmodel'.
+  std::shared_ptr<PidONNXModel<o2::aod::MyTracks>> pidModel; // creates a shared pointer to a new instance 'pidmodel'.
   HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject};
 
   Configurable<float> cfgZvtxCut{"cfgZvtxCut", 10, "Z vtx cut"};
@@ -84,7 +84,7 @@ struct KaonPidTask {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL); // Initializes ccdbApi when cfgUseCCDB is set to 'true'
     }
-    pidModel = std::make_shared<PidONNXModel>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, cfgCertainty.value);
+    pidModel = std::make_shared<PidONNXModel<o2::aod::MyTracks>>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, cfgCertainty.value);
 
     histos.add("hChargePos", ";z;", kTH1F, {{3, -1.5, 1.5}});
     histos.add("hChargeNeg", ";z;", kTH1F, {{3, -1.5, 1.5}});
diff --git a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
index fe4d591e590..8afede63e86 100644
--- a/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLBatchEffAndPurProducer.cxx
@@ -66,7 +66,6 @@ struct PidMlBatchEffAndPurProducer {
   std::array<std::shared_ptr<TH1>, kNPids> hMCPositive;
 
   o2::ccdb::CcdbApi ccdbApi;
-  std::vector<PidONNXModel> models;
 
   Configurable<std::vector<int32_t>> cfgPids{"pids", std::vector<int32_t>(kPids, kPids + kNPids), "PIDs to predict"};
   Configurable<std::array<double, kNDetectors>> cfgDetectorsPLimits{"detectors-p-limits", std::array<double, kNDetectors>(pidml_pt_cuts::defaultModelPLimits), "\"use {detector} when p >= y_{detector}\": array of 3 doubles [y_TPC, y_TOF, y_TRD]"};
@@ -82,6 +81,7 @@ struct PidMlBatchEffAndPurProducer {
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal, aod::McTrackLabels,
                                             aod::pidTPCFullPi, aod::pidTPCFullKa, aod::pidTPCFullPr, aod::pidTPCFullEl, aod::pidTPCFullMu,
                                             aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr, aod::pidTOFFullEl, aod::pidTOFFullMu>>;
+  std::vector<PidONNXModel<BigTracks>> models;
 
   void initHistos()
   {
@@ -203,11 +203,11 @@ struct PidMlBatchEffAndPurProducer {
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
       for (const int32_t& pid : cfgPids.value)
-        models.emplace_back(PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
+        models.emplace_back(PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
                                          ccdbApi, timestamp, pid, 1.1, &cfgDetectorsPLimits.value[0]));
     } else {
       for (int32_t& pid : cfgPids.value)
-        models.emplace_back(PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
+        models.emplace_back(PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value,
                                          ccdbApi, -1, pid, 1.1, &cfgDetectorsPLimits.value[0]));
     }
 
diff --git a/Tools/PIDML/pidMLEffAndPurProducer.cxx b/Tools/PIDML/pidMLEffAndPurProducer.cxx
index 2fcaa7f8477..620fa1b56b1 100644
--- a/Tools/PIDML/pidMLEffAndPurProducer.cxx
+++ b/Tools/PIDML/pidMLEffAndPurProducer.cxx
@@ -34,7 +34,6 @@ using namespace pidml::pidutils;
 struct PidMlEffAndPurProducer {
   HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject};
 
-  PidONNXModel pidModel;
   Configurable<int> cfgPid{"pid", 211, "PID to predict"};
   Configurable<double> cfgNSigmaCut{"n-sigma-cut", 3.0f, "TPC and TOF PID nSigma cut"};
   Configurable<std::array<double, kNDetectors>> cfgDetectorsPLimits{"detectors-p-limits", std::array<double, kNDetectors>(pidml_pt_cuts::defaultModelPLimits), "\"use {detector} when p >= y_{detector}\": array of 3 doubles [y_TPC, y_TOF, y_TRD]"};
@@ -56,6 +55,7 @@ struct PidMlEffAndPurProducer {
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal, aod::McTrackLabels,
                                             aod::pidTPCFullPi, aod::pidTPCFullKa, aod::pidTPCFullPr, aod::pidTPCFullEl, aod::pidTPCFullMu,
                                             aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr, aod::pidTOFFullEl, aod::pidTOFFullMu>>;
+  PidONNXModel<BigTracks> pidModel;
 
   typedef struct nSigma_t {
     double tpc, tof;
@@ -115,7 +115,7 @@ struct PidMlEffAndPurProducer {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1,
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1,
                               cfgPid.value, cfgCertainty.value, &cfgDetectorsPLimits.value[0]);
     }
 
@@ -152,7 +152,7 @@ struct PidMlEffAndPurProducer {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp,
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp,
                               cfgPid.value, cfgCertainty.value, &cfgDetectorsPLimits.value[0]);
     }
 
diff --git a/Tools/PIDML/pidOnnxInterface.h b/Tools/PIDML/pidOnnxInterface.h
index 03e7bc19146..9aa0a5c6fcb 100644
--- a/Tools/PIDML/pidOnnxInterface.h
+++ b/Tools/PIDML/pidOnnxInterface.h
@@ -45,6 +45,7 @@ static const std::vector<std::string> cutVarLabels = {
 
 } // namespace pidml_pt_cuts
 
+template <typename T>
 struct PidONNXInterface {
   PidONNXInterface(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, std::vector<int> const& pids, o2::framework::LabeledArray<double> const& pLimits, std::vector<double> const& minCertainties, bool autoMode) : mNPids{pids.size()}, mPLimits{pLimits}
   {
@@ -78,8 +79,7 @@ struct PidONNXInterface {
   PidONNXInterface& operator=(const PidONNXInterface&) = delete;
   ~PidONNXInterface() = default;
 
-  template <typename T>
-  float applyModel(const T& track, int pid)
+  float applyModel(const T::iterator& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
@@ -90,8 +90,7 @@ struct PidONNXInterface {
     return -1.0f;
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track, int pid)
+  bool applyModelBoolean(const T::iterator& track, int pid)
   {
     for (std::size_t i = 0; i < mNPids; i++) {
       if (mModels[i].mPid == pid) {
@@ -110,7 +109,7 @@ struct PidONNXInterface {
     minCertainties = std::vector<double>(mNPids, 0.5);
   }
 
-  std::vector<PidONNXModel> mModels;
+  std::vector<PidONNXModel<T>> mModels;
   std::size_t mNPids;
   o2::framework::LabeledArray<double> mPLimits;
 };
diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h
index 40813f55c0c..c1c9a59bfe0 100644
--- a/Tools/PIDML/pidOnnxModel.h
+++ b/Tools/PIDML/pidOnnxModel.h
@@ -17,6 +17,7 @@
 #ifndef TOOLS_PIDML_PIDONNXMODEL_H_
 #define TOOLS_PIDML_PIDONNXMODEL_H_
 
+#include <Framework/ASoA.h>
 #include <array>
 #include <algorithm>
 #include <cstdint>
@@ -76,6 +77,7 @@ bool readJsonFile(const std::string& config, rapidjson::Document& d)
 }
 } // namespace
 
+template <typename T>
 struct PidONNXModel {
  public:
   PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp,
@@ -138,14 +140,12 @@ struct PidONNXModel {
   PidONNXModel& operator=(const PidONNXModel&) = delete;
   ~PidONNXModel() = default;
 
-  template <typename T>
-  float applyModel(const T& track)
+  float applyModel(const typename T::iterator& track)
   {
     return getModelOutput(track);
   }
 
-  template <typename T>
-  bool applyModelBoolean(const T& track)
+  bool applyModelBoolean(const typename T::iterator& track)
   {
     return getModelOutput(track) >= mMinCertainty;
   }
@@ -206,7 +206,9 @@ struct PidONNXModel {
     LOG(info) << "Using configuration files: " << localTrainColumnsPath << ", " << localScalingParamsPath;
     if (readJsonFile(localTrainColumnsPath, trainColumnsDoc)) {
       for (auto& param : trainColumnsDoc["columns_for_training"].GetArray()) {
-        mTrainColumns.emplace_back(param.GetString());
+        auto columnLabel = param.GetString();
+        mTrainColumns.emplace_back(columnLabel);
+        mGetters.emplace_back(o2::soa::row_helpers::getColumnGetterByLabel<float, T>(columnLabel));
       }
     }
     if (readJsonFile(localScalingParamsPath, scalingParamsDoc)) {
@@ -216,68 +218,12 @@ struct PidONNXModel {
     }
   }
 
-  template <typename... P1, typename... P2>
-  static constexpr bool is_equal_size(o2::framework::pack<P1...>, o2::framework::pack<P2...>)
-  {
-    return sizeof...(P1) == sizeof...(P2);
-  }
-
   static float scale(float value, const std::pair<float, float>& scalingParams)
   {
     return (value - scalingParams.first) / scalingParams.second;
   }
 
-  template <typename C, typename T>
-  typename C::type getPersistentValue(const T& rowIterator)
-  {
-    return *(static_cast<C>(rowIterator).getIterator());
-  }
-
-  template <typename T, typename... C>
-  std::optional<float> getColumnValueByLabel(o2::framework::pack<C...>, const T& rowIterator, const std::string& columnLabel)
-  {
-    std::optional<float> value = std::nullopt;
-
-    ([&]() {
-      if (value) {
-        return;
-      }
-
-      if constexpr (std::is_arithmetic_v<typename C::type>) {
-        if constexpr (o2::soa::is_dynamic_v<C>) {
-          // check if bindings have the same size as lambda parameters (getter do not have additional parameters)
-          if constexpr (is_equal_size(typename C::bindings_t{}, typename C::callable_t::args{})) {
-            std::string label = C::columnLabel();
-
-            // dynamic columns do not have "f" prefix in columnLabel() return string
-            if (std::strcmp(&columnLabel[1], label.data()) != 0 && columnLabel != label) {
-              return;
-            }
-
-            value = static_cast<float>(rowIterator.template getDynamicColumn<C>());
-          }
-        } else if constexpr (o2::soa::is_persistent_v<C> && !o2::soa::is_index_column_v<C>) {
-          if (columnLabel != C::columnLabel()) {
-            return;
-          }
-
-          value = static_cast<float>(getPersistentValue<C>(rowIterator));
-        }
-      }
-    }(),
-     ...);
-
-    return value;
-  }
-
-  template <typename T>
-  std::optional<float> getColumnValueByLabel(const T& rowIterator, const std::string& columnLabel)
-  {
-    return getColumnValueByLabel(typename T::parent_t::table_t::columns{}, rowIterator, columnLabel);
-  }
-
-  template <typename T>
-  std::vector<float> getValues(const T& track)
+  std::vector<float> getValues(const typename T::iterator& track)
   {
     std::vector<float> output;
     output.reserve(mTrainColumns.size());
@@ -285,7 +231,9 @@ struct PidONNXModel {
     bool useTOF = !tofMissing(track) && inPLimit(track, mPLimits[kTPCTOF]);
     bool useTRD = !trdMissing(track) && inPLimit(track, mPLimits[kTPCTOFTRD]);
 
-    for (const std::string& columnLabel : mTrainColumns) {
+    for (uint32_t i = 0; i < mTrainColumns.size(); ++i) {
+      auto& columnLabel = mTrainColumns[i];
+
       if (
         ((columnLabel == "fTRDSignal" || columnLabel == "fTRDPattern") && !useTRD) ||
         ((columnLabel == "fTOFSignal" || columnLabel == "fBeta") && !useTOF)) {
@@ -300,24 +248,19 @@ struct PidONNXModel {
         scalingParams = scalingParamsEntry->second;
       }
 
-      std::optional<float> value = getColumnValueByLabel(track, columnLabel);
-
-      if (!value) {
-        LOG(fatal) << "Cannot obtain column of name: " << columnLabel;
-      }
+      float value = mGetters[i](track);
 
       if (scalingParams) {
-        value = scale(value.value(), scalingParams.value());
+        value = scale(value, scalingParams.value());
       }
 
-      output.push_back(value.value());
+      output.push_back(value);
     }
 
     return output;
   }
 
-  template <typename T>
-  float getModelOutput(const T& track)
+  float getModelOutput(const typename T::iterator& track)
   {
     // First rank of the expected model input is -1 which means that it is dynamic axis.
     // Axis is exported as dynamic to make it possible to run model inference with the batch of
@@ -382,6 +325,7 @@ struct PidONNXModel {
   }
 
   std::vector<std::string> mTrainColumns;
+  std::vector<float(*)(const typename T::iterator&)> mGetters;
   std::map<std::string, std::pair<float, float>> mScalingParams;
 
   std::shared_ptr<Ort::Env> mEnv = nullptr;
diff --git a/Tools/PIDML/qaPidML.cxx b/Tools/PIDML/qaPidML.cxx
index 54cafc437ab..a34963e8f72 100644
--- a/Tools/PIDML/qaPidML.cxx
+++ b/Tools/PIDML/qaPidML.cxx
@@ -342,11 +342,6 @@ struct QaPidML {
     }
   }
 
-  // one model for one particle
-  PidONNXModel model211;
-  PidONNXModel model2212;
-  PidONNXModel model321;
-
   Configurable<std::string> cfgPathCCDB{"ccdb-path", "Users/m/mkabus/PIDML", "base path to the CCDB directory with ONNX models"};
   Configurable<std::string> cfgCCDBURL{"ccdb-url", "http://alice-ccdb.cern.ch", "URL of the CCDB repository"};
   Configurable<bool> cfgUseCCDB{"useCCDB", true, "Whether to autofetch ML model from CCDB. If false, local file will be used."};
@@ -355,26 +350,32 @@ struct QaPidML {
   o2::ccdb::CcdbApi ccdbApi;
   int currentRunNumber = -1;
 
+  Filter trackFilter = requireGlobalTrackInFilter();
+  using pidTracks = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels, aod::TracksDCA, aod::TrackSelection, aod::pidTOFbeta, aod::TOFSignal>>;
+
+  // one model for one particle
+  PidONNXModel<pidTracks> model211;
+  PidONNXModel<pidTracks> model2212;
+  PidONNXModel<pidTracks> model321;
+
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, 0.5f, pSwitchValue[0]);
-      model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, 0.5f, pSwitchValue[1]);
-      model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, 0.5f, pSwitchValue[2]);
+      model211 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, 0.5f, pSwitchValue[0]);
+      model2212 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, 0.5f, pSwitchValue[1]);
+      model321 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, 0.5f, pSwitchValue[2]);
     }
   }
 
-  Filter trackFilter = requireGlobalTrackInFilter();
-  using pidTracks = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels, aod::TracksDCA, aod::TrackSelection, aod::pidTOFbeta, aod::TOFSignal>>;
   void process(aod::Collisions const& collisions, pidTracks const& tracks, aod::McParticles const& /*mcParticles*/, aod::BCsWithTimestamps const&)
   {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
-      model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, 0.5f, pSwitchValue[0]);
-      model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, 0.5f, pSwitchValue[1]);
-      model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, 0.5f, pSwitchValue[2]);
+      model211 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, 0.5f, pSwitchValue[0]);
+      model2212 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, 0.5f, pSwitchValue[1]);
+      model321 = PidONNXModel<pidTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, 0.5f, pSwitchValue[2]);
     }
 
     for (auto& track : tracks) {
diff --git a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
index 82f644103ed..d023f7390ee 100644
--- a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx
@@ -39,8 +39,6 @@ DECLARE_SOA_TABLE(MlPidResults, "AOD", "MLPIDRESULTS", o2::soa::Index<>, mlpidre
 } // namespace o2::aod
 
 struct SimpleApplyOnnxInterface {
-  PidONNXInterface pidInterface; // One instance to manage all needed ONNX models
-
   Configurable<LabeledArray<double>> cfgPTCuts{"pT_cuts", {pidml_pt_cuts::cuts[0], pidml_pt_cuts::nPids, pidml_pt_cuts::nCutVars, pidml_pt_cuts::pidLabels, pidml_pt_cuts::cutVarLabels}, "pT cuts for each output pid and each detector configuration"};
   Configurable<std::vector<int>> cfgPids{"pids", std::vector<int>{pidml_pt_cuts::pids_v}, "PIDs to predict"};
   Configurable<std::vector<double>> cfgCertainties{"certainties", std::vector<double>{pidml_pt_cuts::certainties_v}, "Min certainties of the models to accept given particle to be of given kind"};
@@ -65,12 +63,14 @@ struct SimpleApplyOnnxInterface {
   // Filter on isGlobalTrack (TracksSelection)
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal>>;
 
+  PidONNXInterface<BigTracks> pidInterface; // One instance to manage all needed ONNX models
+
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidInterface = PidONNXInterface(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
+      pidInterface = PidONNXInterface<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
     }
   }
 
@@ -79,7 +79,7 @@ struct SimpleApplyOnnxInterface {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidInterface = PidONNXInterface(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
+      pidInterface = PidONNXInterface<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPids.value, cfgPTCuts.value, cfgCertainties.value, cfgAutoMode.value);
     }
 
     for (auto& track : tracks) {
diff --git a/Tools/PIDML/simpleApplyPidOnnxModel.cxx b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
index e261c1e10bc..c2ba88e9b82 100644
--- a/Tools/PIDML/simpleApplyPidOnnxModel.cxx
+++ b/Tools/PIDML/simpleApplyPidOnnxModel.cxx
@@ -39,7 +39,6 @@ DECLARE_SOA_TABLE(MlPidResults, "AOD", "MLPIDRESULTS", o2::soa::Index<>, mlpidre
 } // namespace o2::aod
 
 struct SimpleApplyOnnxModel {
-  PidONNXModel pidModel; // One instance per model, e.g., one per each pid to predict
   Configurable<int> cfgPid{"pid", 211, "PID to predict"};
   Configurable<double> cfgCertainty{"certainty", 0.5, "Min certainty of the model to accept given particle to be of given kind"};
 
@@ -61,13 +60,14 @@ struct SimpleApplyOnnxModel {
   // TPC signal (FullTracks), TOF signal (TOFSignal), TOF beta (pidTOFbeta), dcaXY and dcaZ (TracksDCA)
   // Filter on isGlobalTrack (TracksSelection)
   using BigTracks = soa::Filtered<soa::Join<aod::FullTracks, aod::TracksDCA, aod::pidTOFbeta, aod::TrackSelection, aod::TOFSignal>>;
+  PidONNXModel<BigTracks> pidModel; // One instance per model, e.g., one per each pid to predict
 
   void init(InitContext const&)
   {
     if (cfgUseCCDB) {
       ccdbApi.init(cfgCCDBURL);
     } else {
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, cfgCertainty.value);
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, cfgCertainty.value);
     }
   }
 
@@ -76,7 +76,7 @@ struct SimpleApplyOnnxModel {
     auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
     if (cfgUseCCDB && bc.runNumber() != currentRunNumber) {
       uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp();
-      pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, cfgCertainty.value);
+      pidModel = PidONNXModel<BigTracks>(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, cfgCertainty.value);
     }
 
     for (auto& track : tracks) {