Adapted the tests and benchmarks to the SoA track fit EDM.

Author: krasznaa

benchmarks/cuda/toy_detector_cuda.cpp

@@ -73,10 +73,6 @@ BENCHMARK_DEFINE_F(ToyDetectorBenchmark, CUDA)(benchmark::State& state) {
     // Detector view object
     auto det_view = detray::get_data(det_buffer);
 
-    // D2H copy object
-    traccc::device::container_d2h_copy_alg<traccc::track_state_container_types>
-        track_state_d2h{{device_mr, &host_mr}, copy};
-
     for (auto _ : state) {
 
         state.PauseTiming();
@@ -129,7 +125,7 @@ BENCHMARK_DEFINE_F(ToyDetectorBenchmark, CUDA)(benchmark::State& state) {
                                params_cuda_buffer);
 
             // Run track fitting
-            traccc::track_state_container_types::buffer
+            traccc::edm::track_fit_container<traccc::default_algebra>::buffer
                 track_states_cuda_buffer = device_fitting(
                     det_view, field,
                     {track_candidates_cuda_buffer, measurements_cuda_buffer});

tests/common/tests/kalman_fitting_momentum_resolution_test.cpp

@@ -27,11 +27,13 @@
 namespace traccc {
 
 void KalmanFittingMomentumResolutionTests::consistency_tests(
-    const track_state_collection_types::host& track_states_per_track) const {
+    const edm::track_fit_collection<default_algebra>::host::const_proxy_type&
+        track,
+    const edm::track_state_collection<default_algebra>::host&) const {
 
     // The nubmer of track states is supposed be equal to the number
     // of planes
-    ASSERT_EQ(track_states_per_track.size(), std::get<11>(GetParam()));
+    ASSERT_EQ(track.state_indices().size(), std::get<11>(GetParam()));
 }
 
 void KalmanFittingMomentumResolutionTests::momentum_resolution_tests(

tests/common/tests/kalman_fitting_momentum_resolution_test.hpp

@@ -57,7 +57,10 @@ class KalmanFittingMomentumResolutionTests
         1.f * traccc::unit<scalar>::ns};
 
     void consistency_tests(
-        const track_state_collection_types::host& track_states_per_track) const;
+        const edm::track_fit_collection<
+            default_algebra>::host::const_proxy_type& track,
+        const edm::track_state_collection<default_algebra>::host& track_states)
+        const;
 
     void momentum_resolution_tests(std::string_view file_name) const;
 

tests/common/tests/kalman_fitting_telescope_test.hpp

@@ -77,12 +77,14 @@ class KalmanFittingTelescopeTests
                   std::get<11>(GetParam()));
     }
 
-    void consistency_tests(const track_state_collection_types::host&
-                               track_states_per_track) const {
+    void consistency_tests(
+        const edm::track_fit_collection<
+            default_algebra>::host::const_proxy_type& track,
+        const edm::track_state_collection<default_algebra>::host&) const {
 
         // The nubmer of track states is supposed be equal to the number
         // of planes
-        ASSERT_EQ(track_states_per_track.size(), std::get<11>(GetParam()));
+        ASSERT_EQ(track.state_indices().size(), std::get<11>(GetParam()));
     }
 
     protected:

tests/common/tests/kalman_fitting_test.cpp

@@ -151,31 +151,51 @@ void KalmanFittingTests::ndf_tests(
 }
 
 void KalmanFittingTests::ndf_tests(
-    const fitting_result<traccc::default_algebra>& fit_res,
-    const track_state_collection_types::host& track_states_per_track) {
+    const edm::track_fit_collection<default_algebra>::host::const_proxy_type&
+        track,
+    const edm::track_state_collection<default_algebra>::host& track_states,
+    const measurement_collection_types::host& measurements) {
 
     scalar dim_sum = 0;
     std::size_t n_effective_states = 0;
 
-    for (const auto& state : track_states_per_track) {
+    for (unsigned int state_idx : track.state_indices()) {
+
+        auto state = track_states.at(state_idx);
 
-        if (!state.is_hole && state.is_smoothed) {
+        if (!state.is_hole() && state.is_smoothed()) {
 
-            dim_sum += static_cast<scalar>(state.get_measurement().meas_dim);
+            dim_sum += static_cast<scalar>(
+                measurements.at(state.measurement_index()).meas_dim);
             n_effective_states++;
         }
     }
 
     // Check if the number of degree of freedoms is equal to (the sum of
     // measurement dimensions - 5)
-    ASSERT_FLOAT_EQ(static_cast<float>(fit_res.trk_quality.ndf),
+    ASSERT_FLOAT_EQ(static_cast<float>(track.ndf()),
                     static_cast<float>(dim_sum) - 5.f);
 
     // The number of track states is supposed to be eqaul to the number
     // of measurements unless KF failes in the middle of propagation
-    if (n_effective_states == track_states_per_track.size()) {
+    if (n_effective_states == track.state_indices().size()) {
         n_success++;
     }
 }
 
+std::size_t KalmanFittingTests::count_successfully_fitted_tracks(
+    const edm::track_fit_collection<default_algebra>::host& tracks) const {
+
+    const std::size_t n_tracks = tracks.size();
+    std::size_t n_fitted_tracks = 0u;
+
+    for (std::size_t i = 0; i < n_tracks; ++i) {
+        if (tracks.at(i).fit_outcome() == track_fit_outcome::SUCCESS) {
+            n_fitted_tracks++;
+        }
+    }
+
+    return n_fitted_tracks;
+}
+
 }  // namespace traccc

tests/common/tests/kalman_fitting_test.hpp

@@ -12,6 +12,8 @@
 #include "traccc/definitions/common.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate_collection.hpp"
+#include "traccc/edm/track_fit_collection.hpp"
+#include "traccc/edm/track_state_collection.hpp"
 #include "traccc/fitting/kalman_filter/kalman_fitter.hpp"
 #include "traccc/geometry/detector.hpp"
 #include "traccc/simulation/event_generators.hpp"
@@ -79,12 +81,23 @@ class KalmanFittingTests : public testing::Test {
 
     /// Validadte the NDF for track fitting output
     ///
-    /// @param fit_res Fitting statistics result of a track
-    /// @param track_states_per_track Track states of a track
+    /// @param track Fitting statistics result of a track
+    /// @param track_states All track states in the event
+    /// @param measurements All measurements in the event
     ///
     void ndf_tests(
-        const fitting_result<traccc::default_algebra>& fit_res,
-        const track_state_collection_types::host& track_states_per_track);
+        const edm::track_fit_collection<
+            default_algebra>::host::const_proxy_type& track,
+        const edm::track_state_collection<default_algebra>::host& track_states,
+        const measurement_collection_types::host& measurements);
+
+    /// Count the number of tracks that were successfully fitted
+    ///
+    /// @param tracks The track fit collection to count on
+    /// @return The number of tracks that were successfully fitted
+    ///
+    std::size_t count_successfully_fitted_tracks(
+        const edm::track_fit_collection<default_algebra>::host& tracks) const;
 
     // The number of tracks successful with KF
     std::size_t n_success{0u};

tests/common/tests/kalman_fitting_wire_chamber_test.hpp

@@ -70,12 +70,14 @@ class KalmanFittingWireChamberTests
         0.05f / traccc::unit<scalar>::GeV,
         1.f * traccc::unit<scalar>::ns};
 
-    void consistency_tests(const track_state_collection_types::host&
-                               track_states_per_track) const {
+    void consistency_tests(
+        const edm::track_fit_collection<
+            default_algebra>::host::const_proxy_type& track,
+        const edm::track_state_collection<default_algebra>::host&) const {
 
         // The nubmer of track states is supposed be greater than or
         // equal to the number of layers
-        ASSERT_GE(track_states_per_track.size(), n_wire_layers);
+        ASSERT_GE(track.state_indices().size(), n_wire_layers);
     }
 
     protected:

tests/cpu/test_ckf_sparse_tracks_telescope.cpp

@@ -181,22 +181,22 @@ TEST_P(CkfSparseTrackTelescopeTests, Run) {
                          {vecmem::get_data(track_candidates),
                           vecmem::get_data(measurements_per_event)});
         const std::size_t n_fitted_tracks =
-            count_successfully_fitted_tracks(track_states);
+            count_successfully_fitted_tracks(track_states.tracks);
 
-        ASSERT_EQ(track_states.size(), n_truth_tracks);
-        ASSERT_EQ(track_states.size(), n_fitted_tracks);
+        ASSERT_EQ(track_states.tracks.size(), n_truth_tracks);
+        ASSERT_EQ(track_states.tracks.size(), n_fitted_tracks);
 
         for (unsigned int i_trk = 0; i_trk < n_truth_tracks; i_trk++) {
 
-            const auto& track_states_per_track = track_states[i_trk].items;
-            const auto& fit_res = track_states[i_trk].header;
+            consistency_tests(track_states.tracks.at(i_trk),
+                              track_states.states);
 
-            consistency_tests(track_states_per_track);
+            ndf_tests(track_states.tracks.at(i_trk), track_states.states,
+                      measurements_per_event);
 
-            ndf_tests(fit_res, track_states_per_track);
-
-            fit_performance_writer.write(track_states_per_track, fit_res,
-                                         host_det, evt_data);
+            fit_performance_writer.write(
+                track_states.tracks.at(i_trk), track_states.states,
+                measurements_per_event, host_det, evt_data);
         }
     }
 

tests/cpu/test_kalman_fitter_hole_count.cpp

@@ -7,7 +7,6 @@
 
 // Project include(s).
 #include "traccc/bfield/construct_const_bfield.hpp"
-#include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/kalman_fitting_algorithm.hpp"
 #include "traccc/io/utils.hpp"
 #include "traccc/resolution/fitting_performance_writer.hpp"
@@ -161,19 +160,19 @@ TEST_P(KalmanFittingHoleCountTests, Run) {
                  vecmem::get_data(track_candidates.measurements)});
 
     // A sanity check
-    const std::size_t n_tracks = track_states.size();
+    const std::size_t n_tracks = track_states.tracks.size();
     ASSERT_EQ(n_tracks, n_truth_tracks);
 
     // Check the number of holes
     // The three holes at the end are not counted as KF aborts once it goes
     // through all track candidates
-    const auto& fit_res = track_states.at(0u).header;
-    ASSERT_EQ(fit_res.trk_quality.n_holes, 5u);
+    const auto track = track_states.tracks.at(0u);
+    ASSERT_EQ(track.nholes(), 5u);
 
     // Some sanity checks
     ASSERT_FLOAT_EQ(
-        static_cast<float>(fit_res.trk_quality.ndf),
-        static_cast<float>(track_states.at(0u).items.size()) * 2.f - 5.f);
+        static_cast<float>(track.ndf()),
+        static_cast<float>(track.state_indices().size()) * 2.f - 5.f);
 }
 
 INSTANTIATE_TEST_SUITE_P(

tests/cpu/test_kalman_fitter_momentum_resolution.cpp

@@ -7,7 +7,6 @@
 
 // Project include(s).
 #include "traccc/bfield/construct_const_bfield.hpp"
-#include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/kalman_fitting_algorithm.hpp"
 #include "traccc/io/utils.hpp"
 #include "traccc/resolution/fitting_performance_writer.hpp"
@@ -173,9 +172,9 @@ TEST_P(KalmanFittingMomentumResolutionTests, Run) {
                      vecmem::get_data(track_candidates.measurements)});
 
         // Iterator over tracks
-        const std::size_t n_tracks = track_states.size();
+        const std::size_t n_tracks = track_states.tracks.size();
         const std::size_t n_fitted_tracks =
-            count_successfully_fitted_tracks(track_states);
+            count_successfully_fitted_tracks(track_states.tracks);
 
         // n_trakcs = 100
         ASSERT_GE(static_cast<float>(n_tracks),
@@ -185,17 +184,17 @@ TEST_P(KalmanFittingMomentumResolutionTests, Run) {
 
         for (std::size_t i_trk = 0; i_trk < n_tracks; i_trk++) {
 
-            const auto& track_states_per_track = track_states[i_trk].items;
-            const auto& fit_res = track_states[i_trk].header;
+            consistency_tests(track_states.tracks.at(i_trk),
+                              track_states.states);
 
-            consistency_tests(track_states_per_track);
+            ndf_tests(track_states.tracks.at(i_trk), track_states.states,
+                      track_candidates.measurements);
 
-            ndf_tests(fit_res, track_states_per_track);
+            ASSERT_EQ(track_states.tracks.at(i_trk).nholes(), 0u);
 
-            ASSERT_EQ(fit_res.trk_quality.n_holes, 0u);
-
-            fit_performance_writer.write(track_states_per_track, fit_res,
-                                         host_det, evt_data);
+            fit_performance_writer.write(
+                track_states.tracks.at(i_trk), track_states.states,
+                track_candidates.measurements, host_det, evt_data);
         }
     }