From 19054c3274f84478e53750cf0df1b9afeef4f3f6 Mon Sep 17 00:00:00 2001
From: beomki-yeo <beomki.yeo@gmail.com>
Date: Thu, 22 Feb 2024 15:21:34 +0100
Subject: [PATCH] Make track parameter estimation work for low curvature

---
 .../track_params_estimation_helper.hpp        | 15 +++++------
 tests/cpu/test_track_params_estimation.cpp    | 27 ++++++++++++++-----
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/core/include/traccc/seeding/track_params_estimation_helper.hpp b/core/include/traccc/seeding/track_params_estimation_helper.hpp
index 5e235557bd..f23c49061d 100644
--- a/core/include/traccc/seeding/track_params_estimation_helper.hpp
+++ b/core/include/traccc/seeding/track_params_estimation_helper.hpp
@@ -81,15 +81,12 @@ inline TRACCC_HOST_DEVICE bound_vector seed_to_bound_vector(
     scalar A = (uv2[1] - uv1[1]) / (uv2[0] - uv1[0]);
     scalar B = uv2[1] - A * uv2[0];
 
-    // Curvature (with a sign) estimate
-    scalar rho = -2.0f * B / getter::perp(vector2{1., A});
-    // The projection of the top space point on the transverse plane of
-    // the new frame
-    scalar rn = local2[0] * local2[0] + local2[1] * local2[1];
-    // The (1/tanTheta) of momentum in the new frame
-    static constexpr scalar G = static_cast<scalar>(1.f / 24.f);
+    // Radius is a signed distance between circleCenter and first sp, which is
+    // at (0, 0) in the new frame. Sign depends on B
+    const scalar signed_R = -0.5f * getter::perp(vector2{1., A}) / B;
     scalar invTanTheta =
-        local2[2] * std::sqrt(1.f / rn) / (1.f + G * rho * rho * rn);
+        local2[2] /
+        (2.f * signed_R * std::asin(getter::perp(local2) / (2.f * signed_R)));
 
     // The momentum direction in the new frame (the center of the circle
     // has the coordinate (-1.*A/(2*B), 1./(2*B)))
@@ -110,7 +107,7 @@ inline TRACCC_HOST_DEVICE bound_vector seed_to_bound_vector(
 
     // The estimated q/pt in [GeV/c]^-1 (note that the pt is the
     // projection of momentum on the transverse plane of the new frame)
-    scalar qOverPt = rho / getter::norm(bfield);
+    scalar qOverPt = 1.f / (signed_R * getter::norm(bfield));
     // The estimated q/p in [GeV/c]^-1
     getter::element(params, e_bound_qoverp, 0) =
         qOverPt / getter::perp(vector2{1., invTanTheta});
diff --git a/tests/cpu/test_track_params_estimation.cpp b/tests/cpu/test_track_params_estimation.cpp
index df9e79dafd..b5095ea7c6 100644
--- a/tests/cpu/test_track_params_estimation.cpp
+++ b/tests/cpu/test_track_params_estimation.cpp
@@ -22,7 +22,14 @@
 
 using namespace traccc;
 
-TEST(track_params_estimation, helix) {
+// Track parameter estimation with 10 X momentum input
+class TrackParameterEstimationTest
+    : public testing::TestWithParam<traccc::scalar> {};
+
+TEST_P(TrackParameterEstimationTest, helix) {
+
+    // Momentum = 0.1 * input parameter
+    const scalar p0 = static_cast<scalar>(GetParam() * 0.1f);
 
     // Memory resource used by the EDM.
     vecmem::host_memory_resource host_mr;
@@ -34,7 +41,9 @@ TEST(track_params_estimation, helix) {
     // Track property
     const point3 pos{0.f, 0.f, 0.f};
     const scalar time{0.f};
-    const vector3 mom{1.f, 0.f, 1.f * unit<scalar>::GeV};
+    vector3 mom{1.f, 0.f, 1.f * unit<scalar>::GeV};
+    mom = p0 * vector::normalize(mom);
+
     const scalar q{-1.f * unit<scalar>::e};
 
     // Make a helix
@@ -42,9 +51,9 @@ TEST(track_params_estimation, helix) {
 
     // Make three spacepoints with the helix
     spacepoint_collection_types::host spacepoints;
-    spacepoints.push_back({hlx(50 * unit<scalar>::mm), {}});
-    spacepoints.push_back({hlx(100 * unit<scalar>::mm), {}});
-    spacepoints.push_back({hlx(150 * unit<scalar>::mm), {}});
+    spacepoints.push_back({hlx(200 * unit<scalar>::mm), {}});
+    spacepoints.push_back({hlx(400 * unit<scalar>::mm), {}});
+    spacepoints.push_back({hlx(600 * unit<scalar>::mm), {}});
 
     // Make a seed from the three spacepoints
     seed_collection_types::host seeds;
@@ -58,4 +67,10 @@ TEST(track_params_estimation, helix) {
     // momentum
     ASSERT_EQ(bound_params.size(), 1u);
     ASSERT_NEAR(bound_params[0].p(), getter::norm(mom), 1e-4);
-}
\ No newline at end of file
+}
+
+// Test from 0.1 to 10 GeV/c
+INSTANTIATE_TEST_SUITE_P(TrackParameterEstimationGroup,
+                         TrackParameterEstimationTest,
+                         testing::Range(1.f, 100.f),
+                         testing::PrintToStringParamName());
\ No newline at end of file