[Common] [TOF] introduce generic resolution parametrization (#8616)

Author: njacazio

Common/Core/PID/PIDTOF.h

@@ -28,17 +28,19 @@
 #include "TMath.h"
 #include "TGraph.h"
 #include "TFile.h"
+#include "TF2.h"
 
 // O2 includes
 #include "DataFormatsTOF/ParameterContainers.h"
 #include "Framework/Logger.h"
 #include "ReconstructionDataFormats/PID.h"
 #include "Framework/DataTypes.h"
+#include "CommonConstants/PhysicsConstants.h"
 
 namespace o2::pid::tof
 {
 
-// Utility values
+// Utility values (to remove!)
 static constexpr float kCSPEED = TMath::C() * 1.0e2f * 1.0e-12f; /// Speed of light in TOF units (cm/ps)
 static constexpr float kCSPEDDInv = 1.f / kCSPEED;               /// Inverse of the Speed of light in TOF units (ps/cm)
 static constexpr float defaultReturnValue = -999.f;              /// Default return value in case TOF measurement is not available
@@ -55,7 +57,7 @@ class Beta
   /// \param length Length in cm of the track
   /// \param tofSignal TOF signal in ps for the track
   /// \param collisionTime collision time in ps for the event of the track
-  static float GetBeta(const float length, const float tofSignal, const float collisionTime) { return length / (tofSignal - collisionTime) * kCSPEDDInv; }
+  static float GetBeta(const float length, const float tofSignal, const float collisionTime) { return length / (tofSignal - collisionTime) * o2::constants::physics::invLightSpeedCm2PS; }
 
   /// Gets the beta for the track of interest
   /// \param track Track of interest
@@ -139,6 +141,11 @@ class TOFResoParamsV2 : public o2::tof::Parameters<13>
 
   ~TOFResoParamsV2() = default;
 
+  template <o2::track::PID::ID pid>
+  float getResolution(const float, const float) const
+  {
+    return -1.f;
+  }
   // Momentum shift for charge calibration
   void setMomentumChargeShiftParameters(std::unordered_map<std::string, float> const& pars)
   {
@@ -254,14 +261,12 @@ class TOFResoParamsV2 : public o2::tof::Parameters<13>
 class TOFResoParamsV3 : public o2::tof::Parameters<13>
 {
  public:
-  TOFResoParamsV3() : Parameters(std::array<std::string, 13>{"TrkRes.Pi.P0", "TrkRes.Pi.P1", "TrkRes.Pi.P2", "TrkRes.Pi.P3", "time_resolution",
-                                                             "TrkRes.Ka.P0", "TrkRes.Ka.P1", "TrkRes.Ka.P2", "TrkRes.Ka.P3",
-                                                             "TrkRes.Pr.P0", "TrkRes.Pr.P1", "TrkRes.Pr.P2", "TrkRes.Pr.P3"},
+  TOFResoParamsV3() : Parameters(std::array<std::string, 13>{"time_resolution", "time_resolution", "time_resolution", "time_resolution", "time_resolution",
+                                                             "time_resolution", "time_resolution", "time_resolution", "time_resolution",
+                                                             "time_resolution", "time_resolution", "time_resolution", "time_resolution"},
                                  "TOFResoParamsV3")
   {
-    setParameters(std::array<float, 13>{0.008, 0.008, 0.002, 40.0, 60.0,
-                                        0.008, 0.008, 0.002, 40.0,
-                                        0.008, 0.008, 0.002, 40.0});
+    setParameters(std::array<float, 13>{60.0});
   } // Default constructor with default parameters
 
   ~TOFResoParamsV3() = default;
@@ -313,7 +318,7 @@ class TOFResoParamsV3 : public o2::tof::Parameters<13>
     }
   }
 
-  // Time shift for post calibration
+  // Time shift for post calibration to realign as a function of eta
   void setTimeShiftParameters(std::unordered_map<std::string, float> const& pars, bool positive)
   {
     std::string baseOpt = positive ? "TimeShift.Pos." : "TimeShift.Neg.";
@@ -367,13 +372,55 @@ class TOFResoParamsV3 : public o2::tof::Parameters<13>
     return gNegEtaTimeCorr->Eval(eta);
   }
 
+  void setResolutionParametrization(std::unordered_map<std::string, float> const& pars)
+  {
+    static constexpr std::array<const char*, 9> particleNames = {"El", "Mu", "Pi", "Ka", "Pr", "De", "Tr", "He", "Al"};
+    for (int i = 0; i < 9; ++i) {
+      const std::string baseOpt = Form("tofResTrack.%s_", particleNames[i]);
+      // Check if a key begins with a string
+      for (const auto& [key, value] : pars) {
+        if (key.find(baseOpt) == 0) {
+          // Remove from the key the baseOpt
+          const std::string fun = key.substr(baseOpt.size());
+          mResolution[i] = new TF2(baseOpt.c_str(), fun.c_str(), 0., 20, -1, 1.);
+          LOG(info) << "Set the resolution function for " << particleNames[i] << " with formula " << mResolution[i]->GetFormula()->GetExpFormula();
+          break;
+        }
+      }
+    }
+    // Print a summary
+    for (int i = 0; i < 9; ++i) {
+      if (!mResolution[i]) {
+        LOG(info) << "Resolution function for " << particleNames[i] << " not provided, using default " << mDefaultResoParams[i];
+        mResolution[i] = new TF2(Form("tofResTrack.%s_Default", particleNames[i]), mDefaultResoParams[i], 0., 20, -1, 1.);
+      }
+      LOG(info) << "Resolution function for " << particleNames[i] << " is " << mResolution[i]->GetName() << " with formula " << mResolution[i]->GetFormula()->GetExpFormula();
+    }
+  }
+
+  template <o2::track::PID::ID pid>
+  float getResolution(const float p, const float eta) const
+  {
+    return mResolution[pid]->Eval(p, eta);
+  }
+
  private:
   // Charge calibration
   int mEtaN = 0; // Number of eta bins, 0 means no correction
   float mEtaStart = 0.f;
   float mEtaStop = 0.f;
   float mInvEtaWidth = 9999.f;
   std::vector<float> mContent;
+  std::array<TF2*, 9> mResolution{nullptr};
+  static constexpr std::array<const char*, 9> mDefaultResoParams{"14.3*TMath::Power((TMath::Max(x-0.319,0.1))*(1-0.4235*y*y),-0.8467)",
+                                                                 "14.3*TMath::Power((TMath::Max(x-0.319,0.1))*(1-0.4235*y*y),-0.8467)",
+                                                                 "14.3*TMath::Power((TMath::Max(x-0.319,0.1))*(1-0.4235*y*y),-0.8467)",
+                                                                 "42.66*TMath::Power((TMath::Max(x-0.417,0.1))*(1-0.4235*y*y),-0.7145)",
+                                                                 "99.46*TMath::Power((TMath::Max(x-0.447,0.1))*(1-0.4235*y*y),-0.8094)",
+                                                                 "216*TMath::Power((TMath::Max(x-0.647,0.1))*(1-0.4235*y*y),-0.76)",
+                                                                 "315*TMath::Power((TMath::Max(x-0.811,0.1))*(1-0.4235*y*y),-0.783)",
+                                                                 "157*TMath::Power((TMath::Max(x-0.556,0.1))*(1-0.4235*y*y),-0.783)",
+                                                                 "216*TMath::Power((TMath::Max(x-0.647,0.1))*(1-0.4235*y*y),-0.76)"};
 
   // Time shift for post calibration
   TGraph* gPosEtaTimeCorr = nullptr; /// Time shift correction for positive tracks
@@ -391,7 +438,7 @@ class ExpTimes
   static constexpr float mMassZSqared = mMassZ * mMassZ;                   /// (M/z)^2
 
   /// Computes the expected time of a track, given it TOF expected momentum
-  static float ComputeExpectedTime(const float tofExpMom, const float length) { return length * sqrt((mMassZSqared) + (tofExpMom * tofExpMom)) / (kCSPEED * tofExpMom); }
+  static float ComputeExpectedTime(const float tofExpMom, const float length) { return length * sqrt((mMassZSqared) + (tofExpMom * tofExpMom)) / (o2::constants::physics::LightSpeedCm2PS * tofExpMom); }
 
   /// Gets the expected signal of the track of interest under the PID assumption
   /// \param track Track of interest
@@ -401,7 +448,7 @@ class ExpTimes
       return defaultReturnValue;
     }
     if (track.trackType() == o2::aod::track::Run2Track) {
-      return ComputeExpectedTime(track.tofExpMom() * kCSPEDDInv, track.length());
+      return ComputeExpectedTime(track.tofExpMom() * o2::constants::physics::invLightSpeedCm2PS, track.length());
     }
     return ComputeExpectedTime(track.tofExpMom(), track.length());
   }
@@ -416,7 +463,7 @@ class ExpTimes
       return defaultReturnValue;
     }
     if (track.trackType() == o2::aod::track::Run2Track) {
-      return ComputeExpectedTime(track.tofExpMom() * kCSPEDDInv / (1.f + track.sign() * parameters.getMomentumChargeShift(track.eta())), track.length());
+      return ComputeExpectedTime(track.tofExpMom() * o2::constants::physics::invLightSpeedCm2PS / (1.f + track.sign() * parameters.getMomentumChargeShift(track.eta())), track.length());
     }
     LOG(debug) << "TOF exp. mom. " << track.tofExpMom() << " shifted = " << track.tofExpMom() / (1.f + track.sign() * parameters.getMomentumChargeShift(track.eta()));
     return ComputeExpectedTime(track.tofExpMom() / (1.f + track.sign() * parameters.getMomentumChargeShift(track.eta())), track.length()) + parameters.getTimeShift(track.eta(), track.sign());
@@ -432,9 +479,14 @@ class ExpTimes
   static float GetExpectedSigma(const ParamType& parameters, const TrackType& track, const float tofSignal, const float collisionTimeRes)
   {
     const float& mom = track.p();
+    const float& eta = track.eta();
     if (mom <= 0) {
       return -999.f;
     }
+    const float reso = parameters.template getResolution<id>(mom, eta);
+    if (reso > 0) {
+      return std::sqrt(reso * reso + parameters[4] * parameters[4] + collisionTimeRes * collisionTimeRes);
+    }
     if constexpr (id <= o2::track::PID::Pion) {
       LOG(debug) << "Using parameters for the pion hypothesis and ID " << id;
       const float dpp = parameters[0] + parameters[1] * mom + parameters[2] * mMassZ / mom; // mean relative pt resolution;