TPC PID Skim Tree Creation

Common/Core/PID/TPCPIDResponse.h

@@ -23,6 +23,7 @@
 #include "Framework/Logger.h"
 // O2 includes
 #include "ReconstructionDataFormats/PID.h"
+#include "Framework/DataTypes.h"
 #include "DataFormatsTPC/BetheBlochAleph.h"
 
 namespace o2::pid::tpc
@@ -44,13 +45,15 @@ class Response
   void SetMIP(const float mip) { mMIP = mip; }
   void SetChargeFactor(const float chargeFactor) { mChargeFactor = chargeFactor; }
   void SetMultiplicityNormalization(const float multNormalization) { mMultNormalization = multNormalization; }
+  void SetNClNormalization(const float nclnorm) { nClNorm = nclnorm; }
   void SetUseDefaultResolutionParam(const bool useDefault) { mUseDefaultResolutionParam = useDefault; }
   void SetParameters(const Response* response)
   {
     mBetheBlochParams = response->GetBetheBlochParams();
     mResolutionParamsDefault = response->GetResolutionParamsDefault();
     mResolutionParams = response->GetResolutionParams();
     mMIP = response->GetMIP();
+    nClNorm = response->GetNClNormalization();
     mChargeFactor = response->GetChargeFactor();
     mMultNormalization = response->GetMultiplicityNormalization();
     mUseDefaultResolutionParam = response->GetUseDefaultResolutionParam();
@@ -60,6 +63,7 @@ class Response
   const std::array<float, 2> GetResolutionParamsDefault() const { return mResolutionParamsDefault; }
   const std::vector<double> GetResolutionParams() const { return mResolutionParams; }
   const float GetMIP() const { return mMIP; }
+  const float GetNClNormalization() const { return nClNorm; }
   const float GetChargeFactor() const { return mChargeFactor; }
   const float GetMultiplicityNormalization() const { return mMultNormalization; }
   const bool GetUseDefaultResolutionParam() const { return mUseDefaultResolutionParam; }
@@ -89,30 +93,37 @@ class Response
   float mChargeFactor = 2.299999952316284f;
   float mMultNormalization = 11000.;
   bool mUseDefaultResolutionParam = true;
+  float nClNorm = 152.f;
 
-  ClassDefNV(Response, 2);
+  ClassDefNV(Response, 3);
 
 }; // class Response
 
 /// Get expected Signal of the measurement
 template <typename TrackType>
 inline float Response::GetExpectedSignal(const TrackType& track, const o2::track::PID::ID id) const
 {
+  if (!track.hasTPC()) {
+    return -999.f;
+  }
   const float bethe = mMIP * o2::tpc::BetheBlochAleph(track.tpcInnerParam() / o2::track::pid_constants::sMasses[id], mBetheBlochParams[0], mBetheBlochParams[1], mBetheBlochParams[2], mBetheBlochParams[3], mBetheBlochParams[4]) * std::pow((float)o2::track::pid_constants::sCharges[id], mChargeFactor);
-  return bethe >= 0.f ? bethe : 0.f;
+  return bethe >= 0.f ? bethe : -999.f;
 }
 
 /// Gets the expected resolution of the measurement
 template <typename CollisionType, typename TrackType>
 inline float Response::GetExpectedSigma(const CollisionType& collision, const TrackType& track, const o2::track::PID::ID id) const
 {
+  if (!track.hasTPC()) {
+    return -999.f;
+  }
   float resolution = 0.;
   if (mUseDefaultResolutionParam) {
     const float reso = track.tpcSignal() * mResolutionParamsDefault[0] * ((float)track.tpcNClsFound() > 0 ? std::sqrt(1. + mResolutionParamsDefault[1] / (float)track.tpcNClsFound()) : 1.f);
-    reso >= 0.f ? resolution = reso : resolution = 0.f;
+    reso >= 0.f ? resolution = reso : resolution = -999.f;
   } else {
 
-    const double ncl = 159. / track.tpcNClsFound(); //
+    const double ncl = nClNorm / track.tpcNClsFound(); //
     const double p = track.tpcInnerParam();
     const double mass = o2::track::pid_constants::sMasses[id];
     const double bg = p / mass;
@@ -122,7 +133,7 @@ inline float Response::GetExpectedSigma(const CollisionType& collision, const Tr
     const std::vector<double> values{1.f / dEdx, track.tgl(), std::sqrt(ncl), relReso, track.signed1Pt(), collision.multTPC() / mMultNormalization};
 
     const float reso = sqrt(pow(mResolutionParams[0], 2) * values[0] + pow(mResolutionParams[1], 2) * (values[2] * mResolutionParams[5]) * pow(values[0] / sqrt(1 + pow(values[1], 2)), mResolutionParams[2]) + values[2] * pow(values[3], 2) + pow(mResolutionParams[4] * values[4], 2) + pow(values[5] * mResolutionParams[6], 2) + pow(values[5] * (values[0] / sqrt(1 + pow(values[1], 2))) * mResolutionParams[7], 2)) * dEdx * mMIP;
-    reso >= 0.f ? resolution = reso : resolution = 0.f;
+    reso >= 0.f ? resolution = reso : resolution = -999.f;
   }
   return resolution;
 }
@@ -131,13 +142,28 @@ inline float Response::GetExpectedSigma(const CollisionType& collision, const Tr
 template <typename CollisionType, typename TrackType>
 inline float Response::GetNumberOfSigma(const CollisionType& collision, const TrackType& trk, const o2::track::PID::ID id) const
 {
+  if (GetExpectedSigma(collision, trk, id) < 0.) {
+    return -999.f;
+  }
+  if (GetExpectedSignal(trk, id) < 0.) {
+    return -999.f;
+  }
+  if (!trk.hasTPC()) {
+    return -999.f;
+  }
   return ((trk.tpcSignal() - GetExpectedSignal(trk, id)) / GetExpectedSigma(collision, trk, id));
 }
 
 /// Gets the deviation between the actual signal and the expected signal
 template <typename TrackType>
 inline float Response::GetSignalDelta(const TrackType& trk, const o2::track::PID::ID id) const
 {
+  if (GetExpectedSignal(trk, id) < 0.) {
+    return -999.f;
+  }
+  if (!trk.hasTPC()) {
+    return -999.f;
+  }
   return (trk.tpcSignal() - GetExpectedSignal(trk, id));
 }
 
@@ -173,6 +199,7 @@ inline void Response::PrintAll() const
   LOGP(info, "mMIP = {}", mMIP);
   LOGP(info, "mChargeFactor = {}", mChargeFactor);
   LOGP(info, "mMultNormalization = {}", mMultNormalization);
+  LOGP(info, "nClNorm = {}", nClNorm);
 }
 
 } // namespace o2::pid::tpc

Common/DataModel/PIDResponse.h

@@ -888,6 +888,25 @@ DECLARE_SOA_TABLE(pidTOFAl, "AOD", "pidTOFAl", //! Table of the TOF response wit
 namespace pidtpc
 {
 // Expected signals
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalEl, tpcExpSignalEl, //! Expected signal with the TPC detector for electron
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalMu, tpcExpSignalMu, //! Expected signal with the TPC detector for muon
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalPi, tpcExpSignalPi, //! Expected signal with the TPC detector for pion
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalKa, tpcExpSignalKa, //! Expected signal with the TPC detector for kaon
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalPr, tpcExpSignalPr, //! Expected signal with the TPC detector for proton
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalDe, tpcExpSignalDe, //! Expected signal with the TPC detector for deuteron
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalTr, tpcExpSignalTr, //! Expected signal with the TPC detector for triton
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalHe, tpcExpSignalHe, //! Expected signal with the TPC detector for helium3
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalAl, tpcExpSignalAl, //! Expected signal with the TPC detector for alpha
+                           [](float nsigma, float sigma, float tpcsignal) -> float { return tpcsignal - nsigma * sigma; });
+// Expected signals difference
 DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalDiffEl, tpcExpSignalDiffEl, //! Difference between signal and expected for electron
                            [](float nsigma, float sigma) -> float { return nsigma * sigma; });
 DECLARE_SOA_DYNAMIC_COLUMN(TPCExpSignalDiffMu, tpcExpSignalDiffMu, //! Difference between signal and expected for muon
@@ -968,23 +987,23 @@ DEFINE_UNWRAP_NSIGMA_COLUMN(TPCNSigmaAl, tpcNSigmaAl); //! Unwrapped (float) nsi
 
 // Per particle tables
 DECLARE_SOA_TABLE(pidTPCFullEl, "AOD", "pidTPCFullEl", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for electron
-                  pidtpc::TPCExpSignalDiffEl<pidtpc::TPCNSigmaEl, pidtpc::TPCExpSigmaEl>, pidtpc::TPCExpSigmaEl, pidtpc::TPCNSigmaEl);
+                  pidtpc::TPCExpSignalEl<pidtpc::TPCNSigmaEl, pidtpc::TPCExpSigmaEl>, pidtpc::TPCExpSignalDiffEl<pidtpc::TPCNSigmaEl, pidtpc::TPCExpSigmaEl>, pidtpc::TPCExpSigmaEl, pidtpc::TPCNSigmaEl);
 DECLARE_SOA_TABLE(pidTPCFullMu, "AOD", "pidTPCFullMu", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for muon
-                  pidtpc::TPCExpSignalDiffMu<pidtpc::TPCNSigmaMu, pidtpc::TPCExpSigmaMu>, pidtpc::TPCExpSigmaMu, pidtpc::TPCNSigmaMu);
+                  pidtpc::TPCExpSignalMu<pidtpc::TPCNSigmaMu, pidtpc::TPCExpSigmaMu>, pidtpc::TPCExpSignalDiffMu<pidtpc::TPCNSigmaMu, pidtpc::TPCExpSigmaMu>, pidtpc::TPCExpSigmaMu, pidtpc::TPCNSigmaMu);
 DECLARE_SOA_TABLE(pidTPCFullPi, "AOD", "pidTPCFullPi", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for pion
-                  pidtpc::TPCExpSignalDiffPi<pidtpc::TPCNSigmaPi, pidtpc::TPCExpSigmaPi>, pidtpc::TPCExpSigmaPi, pidtpc::TPCNSigmaPi);
+                  pidtpc::TPCExpSignalPi<pidtpc::TPCNSigmaPi, pidtpc::TPCExpSigmaPi>, pidtpc::TPCExpSignalDiffPi<pidtpc::TPCNSigmaPi, pidtpc::TPCExpSigmaPi>, pidtpc::TPCExpSigmaPi, pidtpc::TPCNSigmaPi);
 DECLARE_SOA_TABLE(pidTPCFullKa, "AOD", "pidTPCFullKa", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for kaon
-                  pidtpc::TPCExpSignalDiffKa<pidtpc::TPCNSigmaKa, pidtpc::TPCExpSigmaKa>, pidtpc::TPCExpSigmaKa, pidtpc::TPCNSigmaKa);
+                  pidtpc::TPCExpSignalKa<pidtpc::TPCNSigmaKa, pidtpc::TPCExpSigmaKa>, pidtpc::TPCExpSignalDiffKa<pidtpc::TPCNSigmaKa, pidtpc::TPCExpSigmaKa>, pidtpc::TPCExpSigmaKa, pidtpc::TPCNSigmaKa);
 DECLARE_SOA_TABLE(pidTPCFullPr, "AOD", "pidTPCFullPr", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for proton
-                  pidtpc::TPCExpSignalDiffPr<pidtpc::TPCNSigmaPr, pidtpc::TPCExpSigmaPr>, pidtpc::TPCExpSigmaPr, pidtpc::TPCNSigmaPr);
+                  pidtpc::TPCExpSignalPr<pidtpc::TPCNSigmaPr, pidtpc::TPCExpSigmaPr>, pidtpc::TPCExpSignalDiffPr<pidtpc::TPCNSigmaPr, pidtpc::TPCExpSigmaPr>, pidtpc::TPCExpSigmaPr, pidtpc::TPCNSigmaPr);
 DECLARE_SOA_TABLE(pidTPCFullDe, "AOD", "pidTPCFullDe", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for deuteron
-                  pidtpc::TPCExpSignalDiffDe<pidtpc::TPCNSigmaDe, pidtpc::TPCExpSigmaDe>, pidtpc::TPCExpSigmaDe, pidtpc::TPCNSigmaDe);
+                  pidtpc::TPCExpSignalDe<pidtpc::TPCNSigmaDe, pidtpc::TPCExpSigmaDe>, pidtpc::TPCExpSignalDiffDe<pidtpc::TPCNSigmaDe, pidtpc::TPCExpSigmaDe>, pidtpc::TPCExpSigmaDe, pidtpc::TPCNSigmaDe);
 DECLARE_SOA_TABLE(pidTPCFullTr, "AOD", "pidTPCFullTr", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for triton
-                  pidtpc::TPCExpSignalDiffTr<pidtpc::TPCNSigmaTr, pidtpc::TPCExpSigmaTr>, pidtpc::TPCExpSigmaTr, pidtpc::TPCNSigmaTr);
+                  pidtpc::TPCExpSignalTr<pidtpc::TPCNSigmaTr, pidtpc::TPCExpSigmaTr>, pidtpc::TPCExpSignalDiffTr<pidtpc::TPCNSigmaTr, pidtpc::TPCExpSigmaTr>, pidtpc::TPCExpSigmaTr, pidtpc::TPCNSigmaTr);
 DECLARE_SOA_TABLE(pidTPCFullHe, "AOD", "pidTPCFullHe", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for helium3
-                  pidtpc::TPCExpSignalDiffHe<pidtpc::TPCNSigmaHe, pidtpc::TPCExpSigmaHe>, pidtpc::TPCExpSigmaHe, pidtpc::TPCNSigmaHe);
+                  pidtpc::TPCExpSignalHe<pidtpc::TPCNSigmaHe, pidtpc::TPCExpSigmaHe>, pidtpc::TPCExpSignalDiffHe<pidtpc::TPCNSigmaHe, pidtpc::TPCExpSigmaHe>, pidtpc::TPCExpSigmaHe, pidtpc::TPCNSigmaHe);
 DECLARE_SOA_TABLE(pidTPCFullAl, "AOD", "pidTPCFullAl", //! Table of the TPC (full) response with expected signal, expected resolution and Nsigma for alpha
-                  pidtpc::TPCExpSignalDiffAl<pidtpc::TPCNSigmaAl, pidtpc::TPCExpSigmaAl>, pidtpc::TPCExpSigmaAl, pidtpc::TPCNSigmaAl);
+                  pidtpc::TPCExpSignalAl<pidtpc::TPCNSigmaAl, pidtpc::TPCExpSigmaAl>, pidtpc::TPCExpSignalDiffAl<pidtpc::TPCNSigmaAl, pidtpc::TPCExpSigmaAl>, pidtpc::TPCExpSigmaAl, pidtpc::TPCNSigmaAl);
 
 // Tiny size tables
 DECLARE_SOA_TABLE(pidTPCEl, "AOD", "pidTPCEl", //! Table of the TPC response with binned Nsigma for electron

Common/TableProducer/PID/pidTPC.cxx

@@ -158,6 +158,7 @@ struct tpcPid {
                          enableNetworkOptimizations.value);
         network = temp_net;
         network.evalNetwork(std::vector<float>(network.getInputDimensions(), 1.)); // This is an initialisation and might reduce the overhead of the model
+        network.SetNClNormalization(response.GetNClNormalization());
       }
     }
   }
@@ -186,6 +187,7 @@ struct tpcPid {
     reserveTable(pidAl, tablePIDAl);
 
     std::vector<float> network_prediction;
+    const float nNclNormalization = response.GetNClNormalization();
 
     if (useNetworkCorrection) {
 
@@ -230,7 +232,7 @@ struct tpcPid {
           track_properties[counter_track_props + 2] = trk.signed1Pt();
           track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[i];
           track_properties[counter_track_props + 4] = collisions.iteratorAt(trk.collisionId()).multTPC() / 11000.;
-          track_properties[counter_track_props + 5] = std::sqrt(159. / trk.tpcNClsFound());
+          track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
           counter_track_props += input_dimensions;
         }
 

Common/TableProducer/PID/pidTPCFull.cxx

@@ -158,6 +158,7 @@ struct tpcPidFull {
                          enableNetworkOptimizations.value);
         network = temp_net;
         network.evalNetwork(std::vector<float>(network.getInputDimensions(), 1.)); // This is an initialisation and might reduce the overhead of the model
+        network.SetNClNormalization(response.GetNClNormalization());
       }
     }
   }
@@ -214,7 +215,7 @@ struct tpcPidFull {
       const unsigned long prediction_size = output_dimensions * tracks_size;
 
       network_prediction = std::vector<float>(prediction_size * 9); // For each mass hypotheses
-
+      const float nNclNormalization = response.GetNClNormalization();
       float duration_network = 0;
 
       std::vector<float> track_properties(track_prop_size);
@@ -230,7 +231,7 @@ struct tpcPidFull {
           track_properties[counter_track_props + 2] = trk.signed1Pt();
           track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[i];
           track_properties[counter_track_props + 4] = collisions.iteratorAt(trk.collisionId()).multTPC() / 11000.;
-          track_properties[counter_track_props + 5] = std::sqrt(159. / trk.tpcNClsFound());
+          track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
           counter_track_props += input_dimensions;
         }
 

Common/TableProducer/PID/pidTPCML.cxx

@@ -124,7 +124,7 @@ std::array<float, 6> Network::createInputFromTrack(const C& collision_it, const
   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());
+  const float ncl = std::sqrt(nClNorm / track.tpcNClsFound());
 
   return {p, tgl, signed1Pt, mass, multTPC, ncl};
 }

Common/TableProducer/PID/pidTPCML.h

@@ -46,12 +46,14 @@ class Network
   template <typename C, typename T>
   std::array<float, 6> createInputFromTrack(const C&, const T&, const uint8_t) const; // create a std::vector<float> with all the inputs for the network
   std::vector<Ort::Value> createTensor(std::array<float, 6>) const;                   // create a std::vector<Ort::Value> (= ONNX tensor) for model input
-  float* evalNetwork(std::vector<Ort::Value>);                                      // evaluate the network on a std::vector<Ort::Value> (= ONNX tensor)
-  float* evalNetwork(std::vector<float>);                                           // evaluate the network on a std::vector<float>
+  float* evalNetwork(std::vector<Ort::Value>);                                        // evaluate the network on a std::vector<Ort::Value> (= ONNX tensor)
+  float* evalNetwork(std::vector<float>);                                             // evaluate the network on a std::vector<float>
 
   // Getters & Setters
   int getInputDimensions() const { return mInputShapes[0][1]; };
   int getOutputDimensions() const { return mOutputShapes[0][1]; };
+  void SetNClNormalization(const float nclnorm) { nClNorm = nclnorm; }
+  const float GetNClNormalization() const { return nClNorm; }
 
  private:
   // Environment variables for the ONNX runtime
@@ -65,6 +67,8 @@ class Network
   std::vector<std::string> mOutputNames;
   std::vector<std::vector<int64_t>> mOutputShapes;
 
+  float nClNorm = 152.f;
+
   // Internal function for printing the shape of tensors: See https://github.com/saganatt/PID_ML_in_O2 or O2Physics/Tools/PIDML/simpleApplyPidOnnxModel.cxx
   std::string printShape(const std::vector<int64_t>& v);
 
@@ -75,4 +79,4 @@ class Network
 
 } // namespace o2::pid::tpc
 
-#endif // O2_PID_TPC_ML_H_
+#endif // O2_PID_TPC_ML_H_

Common/Tools/CMakeLists.txt

@@ -24,4 +24,5 @@ o2physics_add_executable(pidparam-tof-reso
 o2physics_add_executable(pidparam-tpc-response
                   SOURCES handleParamTPCResponse.cxx
                   PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore
-                 )
+                 )
+

Common/Tools/handleParamTPCResponse.cxx

@@ -41,6 +41,7 @@ bool initOptionsAndParse(bpo::options_description& options, int argc, char* argv
     "paramMIP", bpo::value<float>()->default_value(50.f), "MIP parameter value")(
     "paramChargeFactor", bpo::value<float>()->default_value(2.299999952316284f), "Charge factor value")(
     "paramMultNormalization", bpo::value<float>()->default_value(11000.), "Multiplicity Normalization")(
+    "paramnClNormalization", bpo::value<float>()->default_value(152.), "Maximum nClusters for normalisation (159 for run 2, 152 for run 3)")(
     "useDefaultParam", bpo::value<bool>()->default_value(true), "Use default sigma parametrisation")(
     "mode", bpo::value<string>()->default_value(""), "Running mode ('read' from file, 'write' to file, 'pull' from CCDB, 'push' to CCDB)")(
     "help,h", "Produce help message.");
@@ -93,6 +94,7 @@ int main(int argc, char* argv[])
   const float mipval = arguments["paramMIP"].as<float>();
   const float chargefacval = arguments["paramChargeFactor"].as<float>();
   const float multNormval = arguments["paramMultNormalization"].as<float>();
+  const float nClNormval = arguments["paramnClNormalization"].as<float>();
   const bool useDefaultParam = arguments["useDefaultParam"].as<bool>();
   const std::string optMode = arguments["mode"].as<std::string>();
   if (optMode.empty()) {
@@ -172,6 +174,7 @@ int main(int argc, char* argv[])
       tpc->SetMIP(mipval);
       tpc->SetChargeFactor(chargefacval);
       tpc->SetMultiplicityNormalization(multNormval);
+      tpc->SetNClNormalization(nClNormval);
       tpc->SetUseDefaultResolutionParam(useDefaultParam);
       tpc->PrintAll();
     }

DPG/Tasks/CMakeLists.txt

@@ -10,5 +10,5 @@
 # or submit itself to any jurisdiction.
 
 add_subdirectory(AOTTrack)
+add_subdirectory(TPC)
 add_subdirectory(FDD)
-