Kalman Fitter on CPU

CMakeLists.txt

@@ -230,7 +230,7 @@ add_subdirectory( plugins )
 
 if ( TRACCC_BUILD_EXAMPLES )
    # Find Boost.
-   find_package( Boost REQUIRED COMPONENTS program_options)
+   find_package( Boost REQUIRED COMPONENTS program_options filesystem)
    # Find ROOT.
    find_package( ROOT COMPONENTS Core RIO Tree Hist )
    add_subdirectory( performance )

core/CMakeLists.txt

@@ -24,6 +24,8 @@ traccc_add_library( traccc_core core TYPE SHARED
   "include/traccc/edm/internal_spacepoint.hpp"
   "include/traccc/edm/seed.hpp"
   "include/traccc/edm/cell.hpp"
+  "include/traccc/edm/track_candidate.hpp"
+  "include/traccc/edm/track_state.hpp"
   # Geometry description.
   "include/traccc/geometry/digitization_config.hpp"
   "include/traccc/geometry/module_map.hpp"
@@ -45,6 +47,12 @@ traccc_add_library( traccc_core core TYPE SHARED
   "src/clusterization/spacepoint_formation.cpp"
   "include/traccc/clusterization/measurement_creation.hpp"
   "src/clusterization/measurement_creation.cpp"
+  # Fitting algorithmic code
+  "include/traccc/fitting/kalman_filter/gain_matrix_smoother.hpp"
+  "include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
+  "include/traccc/fitting/kalman_filter/kalman_actor.hpp"
+  "include/traccc/fitting/kalman_filter/kalman_fitter.hpp"
+  "include/traccc/fitting/fitting_algorithm.hpp"
   # Seed finding algorithmic code.
   "include/traccc/seeding/detail/lin_circle.hpp"
   "include/traccc/seeding/detail/doublet.hpp"

core/include/traccc/edm/track_candidate.hpp

@@ -0,0 +1,26 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2022 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/edm/measurement.hpp"
+#include "traccc/edm/track_parameters.hpp"
+
+namespace traccc {
+
+/// Track candidate contains the measurement object and its surface link
+struct track_candidate {
+    geometry_id surface_link;
+    measurement meas;
+};
+
+/// Declare a track candidates container type
+using track_candidate_container_types =
+    container_types<bound_track_parameters, track_candidate>;
+
+}  // namespace traccc

core/include/traccc/edm/track_parameters.hpp

@@ -11,6 +11,7 @@
 #include "traccc/definitions/common.hpp"
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/definitions/track_parametrization.hpp"
+#include "traccc/edm/container.hpp"
 #include "traccc/utils/unit_vectors.hpp"
 
 // detray include(s).

core/include/traccc/edm/track_state.hpp

@@ -0,0 +1,166 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2022 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/definitions/qualifiers.hpp"
+#include "traccc/edm/container.hpp"
+#include "traccc/edm/measurement.hpp"
+#include "traccc/edm/track_candidate.hpp"
+
+// detray include(s).
+#include "detray/propagator/navigator.hpp"
+#include "detray/tracks/bound_track_parameters.hpp"
+
+namespace traccc {
+
+/// Fitting result per track
+template <typename algebra_t>
+struct fitter_info {
+    using scalar_type = typename algebra_t::scalar_type;
+
+    /// Seed track parameter
+    detray::bound_track_parameters<algebra_t> seed_params;
+
+    /// Fitted track parameter
+    detray::bound_track_parameters<algebra_t> fit_params;
+
+    /// Number of degree of freedoms of fitted track
+    scalar_type ndf;
+
+    /// Chi square of fitted track
+    scalar_type chi2;
+
+    /// P-value of fitted track
+    scalar_type p_val;
+};
+
+/// Fitting result per measurement
+template <typename algebra_t>
+struct track_state {
+
+    using bound_track_parameters_type =
+        detray::bound_track_parameters<algebra_t>;
+    using bound_matrix = typename bound_track_parameters_type::covariance_type;
+    using scalar_type = typename algebra_t::scalar_type;
+    using matrix_operator = typename algebra_t::matrix_actor;
+    using size_type = typename matrix_operator::size_ty;
+    template <size_type ROWS, size_type COLS>
+    using matrix_type =
+        typename matrix_operator::template matrix_type<ROWS, COLS>;
+
+    /// Construction with track candidate
+    TRACCC_HOST_DEVICE
+    track_state(const track_candidate& trk_cand)
+        : m_surface_link(trk_cand.surface_link), m_measurement(trk_cand.meas) {
+        m_predicted.set_surface_link(m_surface_link);
+        m_filtered.set_surface_link(m_surface_link);
+        m_smoothed.set_surface_link(m_surface_link);
+    }
+
+    /// @return the surface link
+    TRACCC_HOST_DEVICE
+    inline std::size_t surface_link() const { return m_surface_link; }
+
+    /// @return the measurement
+    TRACCC_HOST_DEVICE
+    inline const measurement& get_measurement() const { return m_measurement; }
+
+    /// @return the local position of measurement with 2 X 1 matrix
+    // FIXME: The conversion from vector to matrix is inefficient
+    TRACCC_HOST_DEVICE
+    inline matrix_type<2, 1> measurement_local() const {
+        matrix_type<2, 1> ret = matrix_operator().template zero<2, 1>();
+        matrix_operator().element(ret, 0, 0) = m_measurement.local[0];
+        matrix_operator().element(ret, 1, 0) = m_measurement.local[1];
+        return ret;
+    }
+
+    /// @return the covariance of local position of measurement
+    TRACCC_HOST_DEVICE
+    inline matrix_type<2, 2> measurement_covariance() const {
+        matrix_type<2, 2> ret = matrix_operator().template zero<2, 2>();
+        matrix_operator().element(ret, 0, 0) = m_measurement.variance[0];
+        matrix_operator().element(ret, 1, 1) = m_measurement.variance[1];
+        return ret;
+    }
+
+    /// @return the non-const reference of predicted track state
+    TRACCC_HOST_DEVICE
+    inline bound_track_parameters_type& predicted() { return m_predicted; }
+
+    /// @return the const reference of predicted track state
+    TRACCC_HOST_DEVICE
+    inline const bound_track_parameters_type& predicted() const {
+        return m_predicted;
+    }
+
+    /// @return the non-const transport jacobian
+    TRACCC_HOST_DEVICE
+    inline bound_matrix& jacobian() { return m_jacobian; }
+
+    /// @return the const transport jacobian
+    TRACCC_HOST_DEVICE
+    inline const bound_matrix& jacobian() const { return m_jacobian; }
+
+    /// @return the non-const chi square of filtered parameter
+    TRACCC_HOST_DEVICE
+    inline scalar_type& filtered_chi2() { return m_filtered_chi2; }
+
+    /// @return the const chi square of filtered parameter
+    TRACCC_HOST_DEVICE
+    inline const scalar_type& filtered_chi2() const { return m_filtered_chi2; }
+
+    /// @return the non-const filtered parameter
+    TRACCC_HOST_DEVICE
+    inline bound_track_parameters_type& filtered() { return m_filtered; }
+
+    /// @return the const filtered parameter
+    TRACCC_HOST_DEVICE
+    inline const bound_track_parameters_type& filtered() const {
+        return m_filtered;
+    }
+
+    /// @return the non-const chi square of smoothed parameter
+    TRACCC_HOST_DEVICE
+    inline scalar_type& smoothed_chi2() { return m_smoothed_chi2; }
+
+    /// @return the const chi square of smoothed parameter
+    TRACCC_HOST_DEVICE
+    inline const scalar_type& smoothed_chi2() const { return m_smoothed_chi2; }
+
+    /// @return the non-const smoothed parameter
+    TRACCC_HOST_DEVICE
+    inline bound_track_parameters_type& smoothed() { return m_smoothed; }
+
+    /// @return the const smoothed parameter
+    TRACCC_HOST_DEVICE
+    inline const bound_track_parameters_type& smoothed() const {
+        return m_smoothed;
+    }
+
+    private:
+    std::size_t m_surface_link;
+    measurement m_measurement;
+    bound_matrix m_jacobian =
+        matrix_operator().template zero<e_bound_size, e_bound_size>();
+    bound_track_parameters_type m_predicted;
+    scalar_type m_filtered_chi2;
+    bound_track_parameters_type m_filtered;
+    scalar_type m_smoothed_chi2;
+    bound_track_parameters_type m_smoothed;
+};
+
+/// Declare all track_state collection types
+using track_state_collection_types = collection_types<track_state<transform3>>;
+
+/// Declare all track_state container types
+using track_state_container_types =
+    container_types<fitter_info<transform3>, track_state<transform3>>;
+
+}  // namespace traccc

core/include/traccc/fitting/fitting_algorithm.hpp

@@ -0,0 +1,73 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2022 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/definitions/qualifiers.hpp"
+#include "traccc/edm/track_candidate.hpp"
+#include "traccc/edm/track_state.hpp"
+#include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
+#include "traccc/utils/algorithm.hpp"
+
+namespace traccc {
+
+/// Fitting algorithm for a set of tracks
+template <typename fitter_t>
+class fitting_algorithm
+    : public algorithm<track_state_container_types::host(
+          const typename fitter_t::detector_type&,
+          const typename track_candidate_container_types::host&)> {
+
+    public:
+    using transform3_type = typename fitter_t::transform3_type;
+
+    /// Run the algorithm
+    ///
+    /// @param track_candidates the candidate measurements from track finding
+    /// @return the container of the fitted track parameters
+    track_state_container_types::host operator()(
+        const typename fitter_t::detector_type& det,
+        const typename track_candidate_container_types::host& track_candidates)
+        const override {
+
+        fitter_t fitter(det);
+
+        track_state_container_types::host output_states;
+
+        // The number of tracks
+        std::size_t n_tracks = track_candidates.size();
+
+        // Iterate over tracks
+        for (std::size_t i = 0; i < n_tracks; i++) {
+
+            // Seed parameter
+            const auto& seed_param = track_candidates[i].header;
+
+            // Make a vector of track state
+            auto& cands = track_candidates[i].items;
+            vecmem::vector<track_state<transform3_type>> input_states;
+            for (auto& cand : cands) {
+                input_states.emplace_back(cand);
+            }
+
+            // Make a fitter state
+            typename fitter_t::state fitter_state(std::move(input_states));
+
+            // Run fitter
+            fitter.fit(seed_param, fitter_state);
+
+            output_states.push_back(
+                std::move(fitter_state.m_fit_info),
+                std::move(fitter_state.m_fit_actor_state.m_track_states));
+        }
+
+        return output_states;
+    }
+};
+
+}  // namespace traccc