Protect TrackPar against dealing with infinite pT, p

shahor02 · shahor02 · commit e25f6d032ea5 · 2023-11-27T10:50:36.000+01:00
diff --git a/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h b/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h
@@ -115,6 +115,8 @@ GPUconstexpr() int CovarMap[kNParams][kNParams] = {{0, 1, 3, 6, 10},
 GPUconstexpr() int DiagMap[kNParams] = {0, 2, 5, 9, 14};
 
 constexpr float HugeF = o2::constants::math::VeryBig;
+constexpr float MaxPT = 100000.;       // do not allow pTs exceeding this value (to avoid NANs)
+constexpr float MinPTInv = 1. / MaxPT; // do not allow q/pTs less this value (to avoid NANs)
 
 template <typename value_T = float>
 class TrackParametrization
@@ -541,7 +543,10 @@ template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getPtInv() const -> value_t
 {
   // return the inverted track pT
-  const value_t ptInv = gpu::CAMath::Abs(mP[kQ2Pt]);
+  value_t ptInv = gpu::CAMath::Abs(mP[kQ2Pt]);
+  if (ptInv < MinPTInv) {
+    ptInv = MinPTInv;
+  }
   return (mAbsCharge > 1) ? ptInv / mAbsCharge : ptInv;
 }
 
@@ -550,47 +555,40 @@ template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getP2Inv() const -> value_t
 {
   // return the inverted track momentum^2
-  const value_t p2 = mP[kQ2Pt] * mP[kQ2Pt] / (1.f + getTgl() * getTgl());
-  return (mAbsCharge > 1) ? p2 / (mAbsCharge * mAbsCharge) : p2;
+  value_t p2 = getPtInv();
+  return p2 * p2 / (1.f + getTgl() * getTgl());
 }
 
 //____________________________________________________________
 template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getP2() const -> value_t
 {
   // return the track momentum^2
-  const value_t p2inv = getP2Inv();
-  return (p2inv > o2::constants::math::Almost0) ? 1.f / p2inv : o2::constants::math::VeryBig;
+  return 1.f / getP2Inv(); // getP2Inv is protected against being 0, full charge accounted
 }
 
 //____________________________________________________________
 template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getPInv() const -> value_t
 {
-  // return the inverted track momentum^2
-  const value_t pInv = gpu::CAMath::Abs(mP[kQ2Pt]) / gpu::CAMath::Sqrt(1.f + getTgl() * getTgl());
-  return (mAbsCharge > 1) ? pInv / mAbsCharge : pInv;
+  // return the inverted track momentum
+  return getPtInv() / gpu::CAMath::Sqrt(1.f + getTgl() * getTgl()); // getPtInv() is protected against being 0, full charge accounted
 }
 
 //____________________________________________________________
 template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getP() const -> value_t
 {
   // return the track momentum
-  const value_t pInv = getPInv();
-  return (pInv > o2::constants::math::Almost0) ? 1.f / pInv : o2::constants::math::VeryBig;
+  return 1.f / getPInv(); // getPInv is already protected against being 0
 }
 
 //____________________________________________________________
 template <typename value_T>
 GPUdi() auto TrackParametrization<value_T>::getPt() const -> value_t
 {
   // return the track transverse momentum
-  value_t ptI = gpu::CAMath::Abs(mP[kQ2Pt]);
-  if (mAbsCharge > 1) {
-    ptI /= mAbsCharge;
-  }
-  return (ptI > o2::constants::math::Almost0) ? 1.f / ptI : o2::constants::math::VeryBig;
+  return 1.f / getPtInv(); // getPtInv is already protected against being 0
 }
 
 //____________________________________________________________
diff --git a/DataFormats/Reconstruction/src/TrackParametrization.cxx b/DataFormats/Reconstruction/src/TrackParametrization.cxx
@@ -144,9 +144,9 @@ template <typename value_T>
 GPUd() bool TrackParametrization<value_T>::getPosDirGlo(gpu::gpustd::array<value_t, 9>& posdirp) const
 {
   // fill vector with lab x,y,z,px/p,py/p,pz/p,p,sinAlpha,cosAlpha
-  value_t ptI = gpu::CAMath::Abs(getQ2Pt());
+  value_t ptI = getPtInv();
   value_t snp = getSnp();
-  if (ptI < constants::math::Almost0 || gpu::CAMath::Abs(snp) > constants::math::Almost1) {
+  if (gpu::CAMath::Abs(snp) > constants::math::Almost1) {
     return false;
   }
   value_t &sn = posdirp[7], &cs = posdirp[8];

Original file line number	Diff line number	Diff line change
`@@ -144,9 +144,9 @@ template <typename value_T>`
`144`	`144`	`GPUd() bool TrackParametrization<value_T>::getPosDirGlo(gpu::gpustd::array<value_t, 9>& posdirp) const`
`145`	`145`	`{`
`146`	`146`	`// fill vector with lab x,y,z,px/p,py/p,pz/p,p,sinAlpha,cosAlpha`
`147`		`- value_t ptI = gpu::CAMath::Abs(getQ2Pt());`
	`147`	`+ value_t ptI = getPtInv();`
`148`	`148`	`value_t snp = getSnp();`
`149`		`- if (ptI < constants::math::Almost0 \|\| gpu::CAMath::Abs(snp) > constants::math::Almost1) {`
	`149`	`+ if (gpu::CAMath::Abs(snp) > constants::math::Almost1) {`
`150`	`150`	`return false;`
`151`	`151`	`}`
`152`	`152`	`value_t &sn = posdirp[7], &cs = posdirp[8];`