Merge pull request #1145 from stephenswat/refactor/kalman_actor_state

Simplify iterators in Kálmán actor

Author: stephenswat

core/include/traccc/fitting/kalman_filter/kalman_actor.hpp

@@ -54,54 +54,55 @@ struct kalman_actor_state {
     TRACCC_HOST_DEVICE
     typename edm::track_state_collection<algebra_t>::device::proxy_type
     operator()() {
-        if (!backward_mode) {
-            return m_track_states.at(*m_it);
-        } else {
-            return m_track_states.at(*m_it_rev);
-        }
+        return m_track_states.at(m_track.state_indices().at(m_idx));
     }
 
     /// Reset the iterator
     TRACCC_HOST_DEVICE
     void reset() {
-        m_it = m_track.state_indices().begin();
-        m_it_rev = m_track.state_indices().rbegin();
+        if (!backward_mode) {
+            m_idx = 0;
+        } else {
+            m_idx = m_track.state_indices().size() - 1;
+        }
     }
 
     /// Advance the iterator
     TRACCC_HOST_DEVICE
     void next() {
         if (!backward_mode) {
-            m_it++;
+            m_idx++;
         } else {
-            m_it_rev++;
+            m_idx--;
         }
     }
 
     /// @return true if the iterator reaches the end of vector
     TRACCC_HOST_DEVICE
     bool is_complete() {
-        if (!backward_mode && m_it == m_track.state_indices().end()) {
+        if (!backward_mode && m_idx == m_track.state_indices().size()) {
             return true;
-        } else if (backward_mode &&
-                   m_it_rev == m_track.state_indices().rend()) {
+        } else if (backward_mode && m_idx > m_track.state_indices().size()) {
             return true;
         }
         return false;
     }
 
+    TRACCC_HOST_DEVICE
+    bool is_state() {
+        return m_track.state_indices().at(m_idx) !=
+               std::numeric_limits<unsigned int>::max();
+    }
+
     /// Object describing the track fit
     typename edm::track_fit_collection<algebra_t>::device::proxy_type m_track;
     /// All track states in the event
     typename edm::track_state_collection<algebra_t>::device m_track_states;
     /// All measurements in the event
     measurement_collection_types::const_device m_measurements;
 
-    /// Iterator for forward filtering over the track states
-    vecmem::device_vector<unsigned int>::iterator m_it;
-
-    /// Iterator for backward filtering over the track states
-    vecmem::device_vector<unsigned int>::reverse_iterator m_it_rev;
+    /// Index of the current track state
+    unsigned int m_idx;
 
     // The number of holes (The number of sensitive surfaces which do not
     // have a measurement for the track pattern)

core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp

@@ -259,22 +259,23 @@ class kalman_fitter {
             return kalman_fitter_status::ERROR_BARCODE_SEQUENCE_OVERFLOW;
         }
 
-        auto& track = fitter_state.m_fit_actor_state.m_track;
-        auto& track_states = fitter_state.m_fit_actor_state.m_track_states;
+        fitter_state.m_fit_actor_state.backward_mode = true;
+        fitter_state.m_fit_actor_state.reset();
 
         // Since the smoothed track parameter of the last surface can be
         // considered to be the filtered one, we can reversly iterate the
         // algorithm to obtain the smoothed parameter of other surfaces
-        for (auto it = track.state_indices().rbegin();
-             it != track.state_indices().rend(); ++it) {
-            if (!track_states.at(*it).is_hole()) {
-                fitter_state.m_fit_actor_state.m_it_rev = it;
-                break;
-            }
-            // TODO: Return false because there is no valid track state
-            // return false;
+        while (!fitter_state.m_fit_actor_state.is_complete() &&
+               (!fitter_state.m_fit_actor_state.is_state() ||
+                fitter_state.m_fit_actor_state().is_hole())) {
+            fitter_state.m_fit_actor_state.next();
         }
-        auto last = track_states.at(*fitter_state.m_fit_actor_state.m_it_rev);
+
+        if (fitter_state.m_fit_actor_state.is_complete()) {
+            return kalman_fitter_status::SUCCESS;
+        }
+
+        auto last = fitter_state.m_fit_actor_state();
 
         const scalar theta = last.filtered_params().theta();
         if (theta <= 0.f || theta >= constant<traccc::scalar>::pi) {
@@ -322,7 +323,6 @@ class kalman_fitter {
 
         propagation._navigation.set_direction(
             detray::navigation::direction::e_backward);
-        fitter_state.m_fit_actor_state.backward_mode = true;
 
         // Synchronize the current barcode with the input track parameter
         while (propagation._navigation.get_target_barcode() !=