Update track covariance estimation logic

This commit reworks the covariance estimation logic to match that used
in ACTS. It greatly increases the covariances of tracks, hopefully
increasing the tracking efficiency.

Author: stephenswat

benchmarks/cpu/toy_detector_cpu.cpp

@@ -51,7 +51,9 @@ BENCHMARK_DEFINE_F(ToyDetectorBenchmark, CPU)(benchmark::State& state) {
     // Algorithms
     traccc::host::seeding_algorithm sa(seeding_cfg, grid_cfg, filter_cfg,
                                        host_mr);
-    traccc::host::track_params_estimation tp(host_mr);
+    traccc::track_params_estimation_config track_params_estimation_config;
+    traccc::host::track_params_estimation tp(track_params_estimation_config,
+                                             host_mr);
     traccc::host::combinatorial_kalman_filter_algorithm host_finding(
         finding_cfg, host_mr);
     traccc::host::kalman_fitting_algorithm host_fitting(fitting_cfg, host_mr,

benchmarks/cuda/toy_detector_cuda.cpp

@@ -61,7 +61,9 @@ BENCHMARK_DEFINE_F(ToyDetectorBenchmark, CUDA)(benchmark::State& state) {
     // Algorithms
     traccc::cuda::seeding_algorithm sa_cuda(seeding_cfg, grid_cfg, filter_cfg,
                                             mr, async_copy, stream);
-    traccc::cuda::track_params_estimation tp_cuda(mr, async_copy, stream);
+    traccc::track_params_estimation_config track_params_estimation_config;
+    traccc::cuda::track_params_estimation tp_cuda(
+        track_params_estimation_config, mr, async_copy, stream);
     traccc::cuda::combinatorial_kalman_filter_algorithm device_finding(
         finding_cfg, mr, async_copy, stream);
     traccc::cuda::kalman_fitting_algorithm device_fitting(fitting_cfg, mr,

core/include/traccc/seeding/detail/track_params_estimation_config.hpp

@@ -0,0 +1,35 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2021-2022 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+#include <array>
+
+#include "traccc/definitions/common.hpp"
+#include "traccc/definitions/primitives.hpp"
+#include "traccc/definitions/track_parametrization.hpp"
+
+namespace traccc {
+
+struct track_params_estimation_config {
+    std::array<scalar, e_bound_size> initial_sigma = {
+        1.f * unit<scalar>::mm,
+        1.f * unit<scalar>::mm,
+        1.f * unit<scalar>::degree,
+        1.f * unit<scalar>::degree,
+        0.f * unit<scalar>::e / unit<scalar>::GeV,
+        1.f * unit<scalar>::ns};
+
+    scalar initial_sigma_qopt = 0.1f * unit<scalar>::e / unit<scalar>::GeV;
+
+    scalar initial_sigma_pt_rel = 0.1f;
+
+    std::array<scalar, e_bound_size> initial_inflation = {1.f, 1.f, 1.f,
+                                                          1.f, 1.f, 100.f};
+};
+
+}  // namespace traccc

core/include/traccc/seeding/track_params_estimation.hpp

@@ -12,6 +12,7 @@
 #include "traccc/edm/seed_collection.hpp"
 #include "traccc/edm/spacepoint_collection.hpp"
 #include "traccc/edm/track_parameters.hpp"
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 #include "traccc/utils/algorithm.hpp"
 #include "traccc/utils/messaging.hpp"
 
@@ -31,15 +32,15 @@ class track_params_estimation
     : public algorithm<bound_track_parameters_collection_types::host(
           const measurement_collection_types::const_view&,
           const edm::spacepoint_collection::const_view&,
-          const edm::seed_collection::const_view&, const vector3&,
-          const std::array<traccc::scalar, traccc::e_bound_size>&)>,
+          const edm::seed_collection::const_view&, const vector3&)>,
       public messaging {
 
     public:
     /// Constructor for track_params_estimation
     ///
     /// @param mr is the memory resource
     track_params_estimation(
+        const track_params_estimation_config& config,
         vecmem::memory_resource& mr,
         std::unique_ptr<const Logger> logger = getDummyLogger().clone());
 
@@ -56,16 +57,11 @@ class track_params_estimation
     output_type operator()(
         const measurement_collection_types::const_view& measurements,
         const edm::spacepoint_collection::const_view& spacepoints,
-        const edm::seed_collection::const_view& seeds, const vector3& bfield,
-        const std::array<traccc::scalar, traccc::e_bound_size>& stddev = {
-            0.02f * traccc::unit<traccc::scalar>::mm,
-            0.03f * traccc::unit<traccc::scalar>::mm,
-            1.f * traccc::unit<traccc::scalar>::degree,
-            1.f * traccc::unit<traccc::scalar>::degree,
-            0.01f / traccc::unit<traccc::scalar>::GeV,
-            1.f * traccc::unit<traccc::scalar>::ns}) const override;
+        const edm::seed_collection::const_view& seeds,
+        const vector3& bfield) const override;
 
     private:
+    const track_params_estimation_config m_config;
     /// The memory resource to use in the algorithm
     std::reference_wrapper<vecmem::memory_resource> m_mr;
 };  // class track_params_estimation

core/src/seeding/track_params_estimation.cpp

@@ -16,14 +16,15 @@
 namespace traccc::host {
 
 track_params_estimation::track_params_estimation(
-    vecmem::memory_resource& mr, std::unique_ptr<const Logger> logger)
-    : messaging(std::move(logger)), m_mr(mr) {}
+    const track_params_estimation_config& config, vecmem::memory_resource& mr,
+    std::unique_ptr<const Logger> logger)
+    : messaging(std::move(logger)), m_config(config), m_mr(mr) {}
 
 track_params_estimation::output_type track_params_estimation::operator()(
     const measurement_collection_types::const_view& measurements_view,
     const edm::spacepoint_collection::const_view& spacepoints_view,
-    const edm::seed_collection::const_view& seeds_view, const vector3& bfield,
-    const std::array<traccc::scalar, traccc::e_bound_size>& stddev) const {
+    const edm::seed_collection::const_view& seeds_view,
+    const vector3& bfield) const {
 
     // Set up the input / output objects.
     const measurement_collection_types::const_device measurements(
@@ -61,8 +62,27 @@ track_params_estimation::output_type track_params_estimation::operator()(
 
         // Set Covariance
         for (std::size_t j = 0; j < e_bound_size; ++j) {
-            getter::element(track_params.covariance(), j, j) =
-                stddev[j] * stddev[j];
+            scalar var = m_config.initial_sigma[i] * m_config.initial_sigma[i];
+
+            if (i == e_bound_qoverp) {
+                scalar var_theta = getter::element(
+                    track_params.covariance(), e_bound_theta, e_bound_theta);
+
+                var += math::pow(m_config.initial_sigma_qopt *
+                                     math::sin(track_params[e_bound_theta]),
+                                 2.f);
+                var += math::pow(m_config.initial_sigma_pt_rel *
+                                     track_params[e_bound_qoverp],
+                                 2.f);
+                var += var_theta *
+                       math::pow(track_params[e_bound_qoverp] /
+                                     math::tan(track_params[e_bound_theta]),
+                                 2.f);
+            }
+
+            var *= m_config.initial_inflation[i];
+
+            getter::element(track_params.covariance(), i, i) = var;
         }
     }
 

device/common/include/traccc/seeding/device/estimate_track_params.hpp

@@ -15,6 +15,7 @@
 #include "traccc/edm/seed_collection.hpp"
 #include "traccc/edm/spacepoint_collection.hpp"
 #include "traccc/edm/track_parameters.hpp"
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 
 namespace traccc::device {
 
@@ -30,11 +31,10 @@ namespace traccc::device {
 ///
 TRACCC_HOST_DEVICE
 inline void estimate_track_params(
-    global_index_t globalIndex,
+    global_index_t globalIndex, const track_params_estimation_config& config,
     const measurement_collection_types::const_view& measurements_view,
     const edm::spacepoint_collection::const_view& spacepoints_view,
     const edm::seed_collection::const_view& seeds_view, const vector3& bfield,
-    const std::array<traccc::scalar, traccc::e_bound_size>& stddev,
     bound_track_parameters_collection_types::view params_view);
 
 }  // namespace traccc::device

device/common/include/traccc/seeding/device/impl/estimate_track_params.ipp

@@ -8,6 +8,7 @@
 #pragma once
 
 // Project include(s).
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 #include "traccc/seeding/device/estimate_track_params.hpp"
 #include "traccc/seeding/track_params_estimation_helper.hpp"
 
@@ -19,10 +20,10 @@ namespace traccc::device {
 TRACCC_HOST_DEVICE
 inline void estimate_track_params(
     const global_index_t globalIndex,
+    const track_params_estimation_config& config,
     const measurement_collection_types::const_view& measurements_view,
     const edm::spacepoint_collection::const_view& spacepoints_view,
     const edm::seed_collection::const_view& seeds_view, const vector3& bfield,
-    const std::array<traccc::scalar, traccc::e_bound_size>& stddev,
     bound_track_parameters_collection_types::view params_view) {
 
     // Check if anything needs to be done.
@@ -46,10 +47,33 @@ inline void estimate_track_params(
     seed_to_bound_param_vector(track_params, measurements_device,
                                spacepoints_device, this_seed, bfield);
 
+    // NOTE: The code below uses the covariance of theta in the calculation of
+    // the calculation of q/p. Thus, theta must be computed first.
+    static_assert(e_bound_qoverp > e_bound_theta);
+
     // Set Covariance
     for (std::size_t i = 0; i < e_bound_size; i++) {
-        getter::element(track_params.covariance(), i, i) =
-            stddev[i] * stddev[i];
+        scalar var = config.initial_sigma[i] * config.initial_sigma[i];
+
+        if (i == e_bound_qoverp) {
+            scalar var_theta = getter::element(track_params.covariance(),
+                                               e_bound_theta, e_bound_theta);
+
+            var += math::pow(config.initial_sigma_qopt *
+                                 math::sin(track_params[e_bound_theta]),
+                             2.f);
+            var += math::pow(
+                config.initial_sigma_pt_rel * track_params[e_bound_qoverp],
+                2.f);
+            var += var_theta *
+                   math::pow(track_params[e_bound_qoverp] /
+                                 math::tan(track_params[e_bound_theta]),
+                             2.f);
+        }
+
+        var *= config.initial_inflation[i];
+
+        getter::element(track_params.covariance(), i, i) = var;
     }
 }
 

device/cuda/include/traccc/cuda/seeding/track_params_estimation.hpp

@@ -13,6 +13,7 @@
 #include "traccc/edm/seed_collection.hpp"
 #include "traccc/edm/spacepoint_collection.hpp"
 #include "traccc/edm/track_parameters.hpp"
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 #include "traccc/utils/algorithm.hpp"
 #include "traccc/utils/memory_resource.hpp"
 #include "traccc/utils/messaging.hpp"
@@ -31,8 +32,7 @@ struct track_params_estimation
     : public algorithm<bound_track_parameters_collection_types::buffer(
           const measurement_collection_types::const_view&,
           const edm::spacepoint_collection::const_view&,
-          const edm::seed_collection::const_view&, const vector3&,
-          const std::array<traccc::scalar, traccc::e_bound_size>&)>,
+          const edm::seed_collection::const_view&, const vector3&)>,
       public messaging {
 
     public:
@@ -43,6 +43,7 @@ struct track_params_estimation
     ///             and host memory blocks
     /// @param str The CUDA stream to perform the operations in
     track_params_estimation(
+        const track_params_estimation_config& config,
         const traccc::memory_resource& mr, vecmem::copy& copy, stream& str,
         std::unique_ptr<const Logger> logger = getDummyLogger().clone());
 
@@ -59,16 +60,12 @@ struct track_params_estimation
     output_type operator()(
         const measurement_collection_types::const_view& measurements,
         const edm::spacepoint_collection::const_view& spacepoints,
-        const edm::seed_collection::const_view& seeds, const vector3& bfield,
-        const std::array<traccc::scalar, traccc::e_bound_size>& = {
-            0.02f * traccc::unit<traccc::scalar>::mm,
-            0.03f * traccc::unit<traccc::scalar>::mm,
-            1.f * traccc::unit<traccc::scalar>::degree,
-            1.f * traccc::unit<traccc::scalar>::degree,
-            0.01f / traccc::unit<traccc::scalar>::GeV,
-            1.f * traccc::unit<traccc::scalar>::ns}) const override;
+        const edm::seed_collection::const_view& seeds,
+        const vector3& bfield) const override;
 
     private:
+    const track_params_estimation_config m_config;
+
     /// Memory resource used by the algorithm
     traccc::memory_resource m_mr;
     /// The copy object to use

device/cuda/src/seeding/track_params_estimation.cu

@@ -23,22 +23,24 @@ namespace traccc::cuda {
 namespace kernels {
 /// CUDA kernel for running @c traccc::device::estimate_track_params
 __global__ void estimate_track_params(
+    const track_params_estimation_config config,
     measurement_collection_types::const_view measurements_view,
     edm::spacepoint_collection::const_view spacepoints_view,
     edm::seed_collection::const_view seed_view, const vector3 bfield,
-    const std::array<traccc::scalar, traccc::e_bound_size> stddev,
     bound_track_parameters_collection_types::view params_view) {
 
-    device::estimate_track_params(details::global_index1(), measurements_view,
-                                  spacepoints_view, seed_view, bfield, stddev,
-                                  params_view);
+    device::estimate_track_params(details::global_index1(), config,
+                                  measurements_view, spacepoints_view,
+                                  seed_view, bfield, params_view);
 }
 }  // namespace kernels
 
 track_params_estimation::track_params_estimation(
+    const track_params_estimation_config& config,
     const traccc::memory_resource& mr, vecmem::copy& copy, stream& str,
     std::unique_ptr<const Logger> logger)
     : messaging(std::move(logger)),
+      m_config(config),
       m_mr(mr),
       m_copy(copy),
       m_stream(str),
@@ -47,8 +49,8 @@ track_params_estimation::track_params_estimation(
 track_params_estimation::output_type track_params_estimation::operator()(
     const measurement_collection_types::const_view& measurements_view,
     const edm::spacepoint_collection::const_view& spacepoints_view,
-    const edm::seed_collection::const_view& seeds_view, const vector3& bfield,
-    const std::array<traccc::scalar, traccc::e_bound_size>& stddev) const {
+    const edm::seed_collection::const_view& seeds_view,
+    const vector3& bfield) const {
 
     // Get a convenience variable for the stream that we'll be using.
     cudaStream_t stream = details::get_stream(m_stream);
@@ -82,7 +84,7 @@ track_params_estimation::output_type track_params_estimation::operator()(
 
     // run the kernel
     kernels::estimate_track_params<<<num_blocks, num_threads, 0, stream>>>(
-        measurements_view, spacepoints_view, seeds_view, bfield, stddev,
+        m_config, measurements_view, spacepoints_view, seeds_view, bfield,
         params_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 

examples/run/common/throughput_mt.ipp

@@ -13,6 +13,7 @@
 // Project include(s)
 #include "traccc/geometry/detector.hpp"
 #include "traccc/geometry/host_detector.hpp"
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 
 // Command line option include(s).
 #include "traccc/options/clusterization.hpp"
@@ -37,6 +38,7 @@
 #include "traccc/performance/throughput.hpp"
 #include "traccc/performance/timer.hpp"
 #include "traccc/performance/timing_info.hpp"
+#include "traccc/seeding/detail/track_params_estimation_config.hpp"
 
 // VecMem include(s).
 #include <vecmem/memory/host_memory_resource.hpp>
@@ -142,6 +144,7 @@ int throughput_mt(std::string_view description, int argc, char* argv[]) {
     const traccc::seedfinder_config seedfinder_config(seeding_opts);
     const traccc::seedfilter_config seedfilter_config(seeding_opts);
     const traccc::spacepoint_grid_config spacepoint_grid_config(seeding_opts);
+    const traccc::track_params_estimation_config track_params_estimation_config;
     detray::propagation::config propagation_config(propagation_opts);
     typename FULL_CHAIN_ALG::finding_algorithm::config_type finding_cfg(
         finding_opts);
@@ -155,10 +158,10 @@ int throughput_mt(std::string_view description, int argc, char* argv[]) {
     std::vector<FULL_CHAIN_ALG> algs;
     algs.reserve(threading_opts.threads + 1);
     for (std::size_t i = 0; i < threading_opts.threads + 1; ++i) {
-        algs.push_back({host_mr, clustering_cfg, seedfinder_config,
-                        spacepoint_grid_config, seedfilter_config, finding_cfg,
-                        fitting_cfg, det_descr, field, &detector,
-                        logger().clone()});
+        algs.push_back(
+            {host_mr, clustering_cfg, seedfinder_config, spacepoint_grid_config,
+             seedfilter_config, track_params_estimation_config, finding_cfg,
+             fitting_cfg, det_descr, field, &detector, logger().clone()});
     }
 
     // Set up a lambda that calls the correct function on the algorithms.