Try not to use the vector container type of the detector in the KF

niermann999 · niermann999 · commit 715efb8fd94b · 2025-02-05T10:16:58.000Z
diff --git a/benchmarks/common/benchmarks/toy_detector_benchmark.hpp b/benchmarks/common/benchmarks/toy_detector_benchmark.hpp
@@ -22,9 +22,6 @@
 #include "detray/detectors/bfield.hpp"
 #include "detray/io/frontend/detector_reader.hpp"
 #include "detray/io/frontend/detector_writer.hpp"
-#include "detray/navigation/navigator.hpp"
-#include "detray/propagator/propagator.hpp"
-#include "detray/propagator/rk_stepper.hpp"
 #include "detray/test/utils/detectors/build_toy_detector.hpp"
 #include "detray/test/utils/simulation/event_generator/track_generators.hpp"
 #include "detray/tracks/ray.hpp"
diff --git a/benchmarks/cpu/toy_detector_cpu.cpp b/benchmarks/cpu/toy_detector_cpu.cpp
@@ -27,9 +27,6 @@
 #include "detray/core/detector.hpp"
 #include "detray/detectors/bfield.hpp"
 #include "detray/io/frontend/detector_reader.hpp"
-#include "detray/navigation/navigator.hpp"
-#include "detray/propagator/propagator.hpp"
-#include "detray/propagator/rk_stepper.hpp"
 
 // VecMem include(s).
 #include <vecmem/memory/host_memory_resource.hpp>
diff --git a/core/include/traccc/edm/track_state.hpp b/core/include/traccc/edm/track_state.hpp
@@ -15,7 +15,6 @@
 #include "traccc/edm/track_candidate.hpp"
 
 // detray include(s).
-#include "detray/navigation/navigator.hpp"
 #include "detray/tracks/bound_track_parameters.hpp"
 
 namespace traccc {
diff --git a/core/include/traccc/fitting/details/fit_tracks.hpp b/core/include/traccc/fitting/details/fit_tracks.hpp
@@ -58,15 +58,14 @@ track_state_container_types::host fit_tracks(
         }
 
         // Make a fitter state
-        typename fitter_t::state fitter_state(std::move(input_states));
+        typename fitter_t::state fitter_state(vecmem::get_data(input_states));
 
         // Run the fitter.
         fitter.fit(track_candidates.get_headers()[i], fitter_state);
 
         // Save the results into the output container.
-        result.push_back(
-            std::move(fitter_state.m_fit_res),
-            std::move(fitter_state.m_fit_actor_state.m_track_states));
+        result.push_back(std::move(fitter_state.m_fit_res),
+                         std::move(input_states));
     }
 
     // Return the fitted track states.
diff --git a/core/include/traccc/fitting/kalman_filter/gain_matrix_smoother.hpp b/core/include/traccc/fitting/kalman_filter/gain_matrix_smoother.hpp
@@ -12,6 +12,9 @@
 #include "traccc/edm/track_parameters.hpp"
 #include "traccc/edm/track_state.hpp"
 
+// detray include(s)
+#include <detray/geometry/shapes/line.hpp>
+
 namespace traccc {
 
 /// Type unrolling functor to smooth the track parameters after the Kalman
diff --git a/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp b/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp
@@ -12,6 +12,9 @@
 #include "traccc/definitions/track_parametrization.hpp"
 #include "traccc/edm/track_state.hpp"
 
+// detray include(s)
+#include <detray/geometry/shapes/line.hpp>
+
 namespace traccc {
 
 /// Type unrolling functor for Kalman updating
diff --git a/core/include/traccc/fitting/kalman_filter/kalman_actor.hpp b/core/include/traccc/fitting/kalman_filter/kalman_actor.hpp
@@ -15,89 +15,84 @@
 #include "traccc/utils/particle.hpp"
 
 // detray include(s).
-#include "detray/propagator/base_actor.hpp"
+#include <detray/navigation/navigator.hpp>
+#include <detray/propagator/base_actor.hpp>
+
+// vecmem include(s)
+#include <vecmem/containers/device_vector.hpp>
 
 namespace traccc {
 
 /// Detray actor for Kalman filtering
-template <typename algebra_t, template <typename...> class vector_t>
+template <typename algebra_t>
 struct kalman_actor : detray::actor {
 
     // Type declarations
     using track_state_type = track_state<algebra_t>;
 
     // Actor state
     struct state {
+        using iterator_t =
+            typename vecmem::device_vector<track_state_type>::iterator;
 
         /// Constructor with the vector of track states
         TRACCC_HOST_DEVICE
-        state(vector_t<track_state_type>&& track_states)
-            : m_track_states(std::move(track_states)) {
-            m_it = m_track_states.begin();
-            m_it_rev = m_track_states.rbegin();
+        state(vecmem::data::vector_view<track_state_type> track_states,
+              bool backward_mode = false)
+            : m_begin{track_states.ptr()},
+              m_end{track_states.ptr() + track_states.size()} {
+            reset(backward_mode);
         }
 
         /// Constructor with the vector of track states
         TRACCC_HOST_DEVICE
-        state(const vector_t<track_state_type>& track_states)
-            : m_track_states(track_states) {
-            m_it = m_track_states.begin();
-            m_it_rev = m_track_states.rbegin();
+        state(vecmem::device_vector<track_state<algebra_t>>& track_states,
+              bool backward_mode = false)
+            : m_begin{track_states.begin()}, m_end{track_states.end()} {
+            reset(backward_mode);
         }
 
-        /// @return the reference of track state pointed by the iterator
-        TRACCC_HOST_DEVICE
-        track_state_type& operator()() {
-            if (!backward_mode) {
-                return *m_it;
-            } else {
-                return *m_it_rev;
-            }
-        }
+        /// Get range iterators
+        /// @{
+        TRACCC_HOST_DEVICE iterator_t begin() const { return m_begin; }
+        TRACCC_HOST_DEVICE iterator_t rbegin() const { return m_end; }
+        TRACCC_HOST_DEVICE iterator_t end() const { return m_end; }
+        TRACCC_HOST_DEVICE iterator_t rend() const { return m_begin; }
+        TRACCC_HOST_DEVICE iterator_t current() const { return m_it; }
+        /// @}
+
+        /// @return a reference of the track state pointed to by the iterator
+        TRACCC_HOST_DEVICE track_state_type& operator()() { return *m_it; }
 
         /// Reset the iterator
         TRACCC_HOST_DEVICE
-        void reset() {
-            m_it = m_track_states.begin();
-            m_it_rev = m_track_states.rbegin();
+        void reset(bool backward_mode = false) {
+            m_it = backward_mode ? m_begin : m_end;
         }
 
-        /// Advance the iterator
-        TRACCC_HOST_DEVICE
-        void next() {
-            if (!backward_mode) {
-                m_it++;
-            } else {
-                m_it_rev++;
-            }
-        }
+        /// Move the iterator forward
+        TRACCC_HOST_DEVICE void next() { ++m_it; }
+
+        /// Move the iterator backward
+        TRACCC_HOST_DEVICE void previous() { --m_it; }
 
         /// @return true if the iterator reaches the end of vector
         TRACCC_HOST_DEVICE
-        bool is_complete() {
-            if (!backward_mode && m_it == m_track_states.end()) {
-                return true;
-            } else if (backward_mode && m_it_rev == m_track_states.rend()) {
-                return true;
-            }
-            return false;
+        bool is_complete(bool backward_mode = false) {
+            return (!backward_mode && m_it == m_end) ||
+                   (backward_mode && m_it == m_begin);
         }
 
-        // vector of track states
-        vector_t<track_state_type> m_track_states;
-
-        // iterator for forward filtering
-        typename vector_t<track_state_type>::iterator m_it;
-
-        // iterator for backward filtering
-        typename vector_t<track_state_type>::reverse_iterator m_it_rev;
+        /// First track state of the track
+        iterator_t m_begin;
+        /// Last track state of the track
+        iterator_t m_end;
+        /// Current track state
+        iterator_t m_it;
 
         // The number of holes (The number of sensitive surfaces which do not
         // have a measurement for the track pattern)
         unsigned int n_holes{0u};
-
-        // Run back filtering for smoothing, if true
-        bool backward_mode = false;
     };
 
     /// Actor operation to perform the Kalman filtering
@@ -110,9 +105,11 @@ struct kalman_actor : detray::actor {
 
         auto& stepping = propagation._stepping;
         auto& navigation = propagation._navigation;
+        const bool backward_mode{navigation.direction() ==
+                                 detray::navigation::direction::e_backward};
 
         // If the iterator reaches the end, terminate the propagation
-        if (actor_state.is_complete()) {
+        if (actor_state.is_complete(backward_mode)) {
             propagation._heartbeat &= navigation.abort();
             return;
         }
@@ -125,23 +122,22 @@ struct kalman_actor : detray::actor {
             // Increase the hole counts if the propagator fails to find the next
             // measurement
             if (navigation.barcode() != trk_state.surface_link()) {
-                if (!actor_state.backward_mode) {
+                if (!backward_mode) {
                     actor_state.n_holes++;
                 }
                 return;
             }
 
-            // This track state is not a hole
-            if (!actor_state.backward_mode) {
-                trk_state.is_hole = false;
-            }
-
             // Run Kalman Gain Updater
             const auto sf = navigation.get_surface();
 
             bool res = false;
 
-            if (!actor_state.backward_mode) {
+            if (!backward_mode) {
+
+                // This track state is not a hole
+                trk_state.is_hole = false;
+
                 // Forward filter
                 res = sf.template visit_mask<gain_matrix_updater<algebra_t>>(
                     trk_state, propagation._stepping.bound_params());
@@ -151,10 +147,16 @@ struct kalman_actor : detray::actor {
 
                 // Set full jacobian
                 trk_state.jacobian() = stepping.full_jacobian();
+
+                // Update iterator
+                actor_state.next();
             } else {
                 // Backward filter for smoothing
                 res = sf.template visit_mask<two_filters_smoother<algebra_t>>(
                     trk_state, propagation._stepping.bound_params());
+
+                // Update iterator
+                actor_state.previous();
             }
 
             // Abort if the Kalman update fails
@@ -170,9 +172,6 @@ struct kalman_actor : detray::actor {
                 stepping.particle_hypothesis(),
                 propagation._stepping.bound_params()));
 
-            // Update iterator
-            actor_state.next();
-
             // Flag renavigation of the current candidate
             navigation.set_high_trust();
         }
diff --git a/core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp b/core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp
diff --git a/device/common/include/traccc/finding/device/impl/find_tracks.ipp b/device/common/include/traccc/finding/device/impl/find_tracks.ipp
diff --git a/tests/common/tests/kalman_fitting_momentum_resolution_test.hpp b/tests/common/tests/kalman_fitting_momentum_resolution_test.hpp
