From adf84eb8ebf0118a4f2f8e110de45988a1699e28 Mon Sep 17 00:00:00 2001
From: Felix Schlepper <felix.schlepper@cern.ch>
Date: Fri, 21 Aug 2026 13:20:46 +0200
Subject: [PATCH] TRK: add nELossSteps

Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>
---
 .../TrackParametrization.h                    |  7 +++++
 .../src/TrackParametrization.cxx              | 31 ++++++++++++++++---
 .../src/TrackParametrizationWithError.cxx     | 10 ++----
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h b/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h
index ec624ed53ca30..4b1fef22a3bb4 100644
--- a/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h
+++ b/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrization.h
@@ -210,6 +210,13 @@ class TrackParametrization
   GPUd() value_t getE() const;
   GPUdi() static value_t getdEdxBB(value_t betagamma) { return BetheBlochSolid(betagamma); }
   GPUdi() static value_t getdEdxBBOpt(value_t betagamma) { return BetheBlochSolidOpt(betagamma); }
+
+  GPUdi() int nELossSteps(value_T dE, value_T ekin) const noexcept
+  {
+    const int n = 1 + int(gpu::CAMath::Abs(dE) / ekin * ELoss2EKinThreshInv);
+    return n > MaxELossIter ? MaxELossIter : n;
+  }
+  GPUd() int getELossSteps(value_t xrho, bool anglecorr) const;
   GPUdi() static value_t getBetheBlochSolidDerivativeApprox(value_T dedx, value_T bg) { return BetheBlochSolidDerivative(dedx, bg); }
 
   GPUd() value_t getTheta() const;
diff --git a/DataFormats/Reconstruction/src/TrackParametrization.cxx b/DataFormats/Reconstruction/src/TrackParametrization.cxx
index 4564e6fc9bfde..fb398bbcf07cf 100644
--- a/DataFormats/Reconstruction/src/TrackParametrization.cxx
+++ b/DataFormats/Reconstruction/src/TrackParametrization.cxx
@@ -867,6 +867,30 @@ GPUd() bool TrackParametrization<value_T>::getXatLabR(value_t r, value_t& x, val
   return true;
 }
 
+//______________________________________________
+template <typename value_T>
+GPUd() int TrackParametrization<value_T>::getELossSteps(value_t xrho, bool anglecorr) const
+{
+  // Copied from correctForMaterial before entering its energy-loss loop
+  const value_t m = getPID().getMass();
+  if (!(m > 0) || xrho == 0.f) {
+    return 0; // correctForMaterial skips the energy-loss block entirely
+  }
+  if (anglecorr) {
+    const value_t csp2 = (1.f - getSnp()) * (1.f + getSnp()); // cos(phi)^2
+    const value_t cst2I = (1.f + getTgl() * getTgl());        // 1/cos(lambda)^2
+    xrho *= gpu::CAMath::Sqrt(cst2I / csp2);
+  }
+  const value_t p = getP(), massInv = 1.f / m;
+  const value_t e = gpu::CAMath::Sqrt(p * p + getPID().getMass2()), ekin = e - m;
+  value_t dedx = getdEdxBBOpt(p * massInv);
+  const int charge2 = getAbsCharge() * getAbsCharge();
+  if (charge2 != 1) {
+    dedx *= charge2;
+  }
+  return nELossSteps(dedx * xrho, ekin);
+}
+
 //______________________________________________
 template <typename value_T>
 GPUd() bool TrackParametrization<value_T>::correctForELoss(value_t xrho, bool anglecorr)
@@ -891,7 +915,7 @@ GPUd() bool TrackParametrization<value_T>::correctForELoss(value_t xrho, bool an
       xrho *= angle;
     }
     int charge2 = getAbsCharge() * getAbsCharge();
-    value_t p = getP(), p0 = p, p2 = p * p, e2 = p2 + getPID().getMass2(), massInv = 1. / m, bg = p * massInv;
+    value_t p = getP(), p0 = p, p2 = p * p, e2 = p2 + getPID().getMass2(), massInv = 1.f / m, bg = p * massInv;
     value_t e = gpu::CAMath::Sqrt(e2), ekin = e - m, dedx = getdEdxBBOpt(bg);
 #ifdef _BB_NONCONST_CORR_
     value_t dedxDer = 0., dedx1 = dedx;
@@ -900,10 +924,7 @@ GPUd() bool TrackParametrization<value_T>::correctForELoss(value_t xrho, bool an
       dedx *= charge2;
     }
     value_t dE = dedx * xrho;
-    int na = 1 + int(gpu::CAMath::Abs(dE) / ekin * ELoss2EKinThreshInv);
-    if (na > MaxELossIter) {
-      na = MaxELossIter;
-    }
+    int na = nELossSteps(dE, ekin);
     if (na > 1) {
       dE /= na;
       xrho /= na;
diff --git a/DataFormats/Reconstruction/src/TrackParametrizationWithError.cxx b/DataFormats/Reconstruction/src/TrackParametrizationWithError.cxx
index e00bde297602b..748cb47094d26 100644
--- a/DataFormats/Reconstruction/src/TrackParametrizationWithError.cxx
+++ b/DataFormats/Reconstruction/src/TrackParametrizationWithError.cxx
@@ -1438,10 +1438,7 @@ GPUd() bool TrackParametrizationWithError<value_T>::correctForMaterial(value_t x
       dedx *= charge2;
     }
     value_t dE = dedx * xrho;
-    int na = 1 + int(gpu::CAMath::Abs(dE) / ekin * ELoss2EKinThreshInv);
-    if (na > MaxELossIter) {
-      na = MaxELossIter;
-    }
+    int na = this->nELossSteps(dE, ekin);
     if (na > 1) {
       dE /= na;
       xrho /= na;
@@ -1573,10 +1570,7 @@ GPUd() bool TrackParametrizationWithError<value_T>::correctForMaterial(TrackPara
       dedx *= charge2;
     }
     value_t dE = dedx * xrho;
-    int na = 1 + int(gpu::CAMath::Abs(dE) / ekin * ELoss2EKinThreshInv);
-    if (na > MaxELossIter) {
-      na = MaxELossIter;
-    }
+    int na = this->nELossSteps(dE, ekin);
     if (na > 1) {
       dE /= na;
       xrho /= na;