Methods for Kalman filter linearized wrt reference track

Author: shahor02

DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h

@@ -229,6 +229,7 @@ class TrackParametrization
   // parameters manipulation
   GPUd() bool correctForELoss(value_t xrho, bool anglecorr = false);
   GPUd() bool rotateParam(value_t alpha);
+  GPUd() bool rotateParam(value_t& alpha, value_t& ca, value_t& sa);
   GPUd() bool propagateParamTo(value_t xk, value_t b);
   GPUd() bool propagateParamTo(value_t xk, const dim3_t& b);
   GPUd() void invertParam();

DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h

@@ -89,9 +89,14 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
   // parameters + covmat manipulation
   GPUd() bool testRotate(value_t alpha) const;
   GPUd() bool rotate(value_t alpha);
-  GPUd() bool propagateTo(value_t xk, value_t b);
+  GPUd() bool rotate(value_t alpha, TrackParametrization<value_T>& linRef, value_t bz);
+  GPUd() bool propagateTo(value_t xk, value_t bz);
+  GPUd() bool propagateTo(value_t xk, TrackParametrization<value_T>& linRef, value_t bz);
+  GPUd() bool propagateTo(value_t xk, value_t bz, TrackParametrization<value_T>* linRef) { return linRef ? propagateTo(xk, *linRef, bz) : propagateTo(xk, bz); }
   GPUd() bool propagateTo(value_t xk, const dim3_t& b);
-  GPUd() bool propagateToDCA(const o2::dataformats::VertexBase& vtx, value_t b, o2::dataformats::DCA* dca = nullptr, value_t maxD = 999.f);
+  GPUd() bool propagateTo(value_t xk, TrackParametrization<value_T>& linRef, const dim3_t& b);
+  GPUd() bool propagateTo(value_t xk, const dim3_t& b, TrackParametrization<value_T>* linRef) { return linRef ? propagateTo(xk, *linRef, b) : propagateTo(xk, b); }
+  GPUd() bool propagateToDCA(const o2::dataformats::VertexBase& vtx, value_t bz, o2::dataformats::DCA* dca = nullptr, value_t maxD = 999.f);
   GPUd() void invert();
   GPUd() value_t getPredictedChi2(const dim2_t& p, const dim3_t& cov) const;
   GPUd() value_t getPredictedChi2Quiet(const dim2_t& p, const dim3_t& cov) const;
@@ -118,7 +123,7 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
   GPUd() bool update(const BaseCluster<T>& p);
 
   GPUd() bool correctForMaterial(value_t x2x0, value_t xrho, bool anglecorr = false);
-
+  GPUd() bool correctForMaterial(TrackParametrization<value_T>& linRef, value_t x2x0, value_t xrho, bool anglecorr = false);
   GPUd() void resetCovariance(value_t s2 = 0);
   GPUd() void checkCovariance();
   GPUd() void checkCorrelations();

DataFormats/Reconstruction/src/TrackParametrization.cxx

@@ -188,6 +188,38 @@ GPUd() bool TrackParametrization<value_T>::rotateParam(value_t alpha)
   return true;
 }
 
+//______________________________________________________________
+template <typename value_T>
+GPUd() bool TrackParametrization<value_T>::rotateParam(value_t& alpha, value_t& ca, value_t& sa)
+{
+  // rotate to alpha frame
+  if (gpu::CAMath::Abs(getSnp()) > constants::math::Almost1) {
+    LOGP(debug, "Precondition is not satisfied: |sin(phi)|>1 ! {:f}", getSnp());
+    return false;
+  }
+  //
+  math_utils::detail::bringToPMPi<value_t>(alpha);
+  math_utils::detail::sincos(alpha - getAlpha(), sa, ca);
+  value_t snp = getSnp(), csp = gpu::CAMath::Sqrt((1.f - snp) * (1.f + snp)); // Improve precision
+  // RS: check if rotation does no invalidate track model (cos(local_phi)>=0, i.e. particle direction in local frame is along the X axis
+  if ((csp * ca + snp * sa) < 0) {
+    // LOGF(warning,"Rotation failed: local cos(phi) would become {:.2f}", csp * ca + snp * sa);
+    return false;
+  }
+  //
+  value_t tmp = snp * ca - csp * sa;
+  if (gpu::CAMath::Abs(tmp) > constants::math::Almost1) {
+    LOGP(debug, "Rotation failed: new snp {:.2f}", tmp);
+    return false;
+  }
+  value_t xold = getX(), yold = getY();
+  mAlpha = alpha;
+  mX = xold * ca + yold * sa;
+  mP[kY] = -xold * sa + yold * ca;
+  mP[kSnp] = tmp;
+  return true;
+}
+
 //____________________________________________________________
 template <typename value_T>
 GPUd() bool TrackParametrization<value_T>::propagateParamTo(value_t xk, const dim3_t& b)