Introduced track_candidates_container::device and track_fit_container::device.

krasznaa · krasznaa · commit a4a69001bd4c · 2025-08-11T13:46:17.000+02:00
To make it a bit easier to deal with the slightly complex EDMs
of track finding and fitting in the code.
diff --git a/core/include/traccc/edm/track_candidate_container.hpp b/core/include/traccc/edm/track_candidate_container.hpp
@@ -8,6 +8,7 @@
 #pragma once
 
 // Local include(s).
+#include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate_collection.hpp"
 
@@ -71,6 +72,28 @@ struct track_candidate_container {
         measurement_collection_types::const_view measurements;
     };
 
+    struct device {
+        /// Constructor from a view
+        TRACCC_HOST_DEVICE
+        explicit device(const view& v)
+            : tracks{v.tracks}, measurements{v.measurements} {}
+        /// The track candidates
+        track_candidate_collection<ALGEBRA>::device tracks;
+        /// Measurements referenced by the tracks
+        measurement_collection_types::const_device measurements;
+    };
+
+    struct const_device {
+        /// Constructor from a view
+        TRACCC_HOST_DEVICE
+        explicit const_device(const const_view& v)
+            : tracks{v.tracks}, measurements{v.measurements} {}
+        /// The track candidates
+        track_candidate_collection<ALGEBRA>::const_device tracks;
+        /// Measurements referenced by the tracks
+        measurement_collection_types::const_device measurements;
+    };
+
 };  // struct track_candidate_container
 
 }  // namespace traccc::edm
diff --git a/core/include/traccc/edm/track_fit_container.hpp b/core/include/traccc/edm/track_fit_container.hpp
@@ -8,6 +8,7 @@
 #pragma once
 
 // Local include(s).
+#include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_fit_collection.hpp"
 #include "traccc/edm/track_state_collection.hpp"
@@ -70,6 +71,36 @@ struct track_fit_container {
         track_state_collection<ALGEBRA>::const_data states;
     };
 
+    struct device {
+        /// Constructor from a view
+        TRACCC_HOST_DEVICE
+        explicit device(const view& v)
+            : tracks{v.tracks},
+              states{v.states},
+              measurements{v.measurements} {}
+        /// The fitted tracks
+        track_fit_collection<ALGEBRA>::device tracks;
+        /// The track states used for the fit
+        track_state_collection<ALGEBRA>::device states;
+        /// The measurements used for the fit
+        measurement_collection_types::const_device measurements;
+    };
+
+    struct const_device {
+        /// Constructor from a view
+        TRACCC_HOST_DEVICE
+        explicit const_device(const const_view& v)
+            : tracks{v.tracks},
+              states{v.states},
+              measurements{v.measurements} {}
+        /// The fitted tracks
+        track_fit_collection<ALGEBRA>::const_device tracks;
+        /// The track states used for the fit
+        track_state_collection<ALGEBRA>::const_device states;
+        /// The measurements used for the fit
+        measurement_collection_types::const_device measurements;
+    };
+
 };  // struct track_fit_container
 
 }  // namespace traccc::edm
diff --git a/device/common/include/traccc/fitting/device/fit_backward.hpp b/device/common/include/traccc/fitting/device/fit_backward.hpp
@@ -22,28 +22,23 @@ TRACCC_HOST_DEVICE inline void fit_backward(
     vecmem::device_vector<const unsigned int> param_ids(payload.param_ids_view);
     vecmem::device_vector<unsigned int> param_liveness(
         payload.param_liveness_view);
-    typename edm::track_fit_collection<
+    typename edm::track_fit_container<
         typename fitter_t::detector_type::algebra_type>::device
-        tracks(payload.tracks_view.tracks);
-    typename edm::track_state_collection<
-        typename fitter_t::detector_type::algebra_type>::device
-        track_states(payload.tracks_view.states);
-    measurement_collection_types::const_device measurements{
-        payload.tracks_view.measurements};
+        tracks(payload.tracks_view);
 
-    if (globalIndex >= tracks.size()) {
+    if (globalIndex >= tracks.tracks.size()) {
         return;
     }
 
     const unsigned int param_id = param_ids.at(globalIndex);
-    auto track = tracks.at(param_id);
+    auto track = tracks.tracks.at(param_id);
 
     // Run fitting
     fitter_t fitter(det, payload.field_data, cfg);
 
     if (param_liveness.at(param_id) > 0u) {
         typename fitter_t::state fitter_state(
-            track, track_states, measurements,
+            track, tracks.states, tracks.measurements,
             *(payload.barcodes_view.ptr() + param_id));
 
         kalman_fitter_status fit_status = fitter.smooth(fitter_state);
diff --git a/device/common/include/traccc/fitting/device/fit_forward.hpp b/device/common/include/traccc/fitting/device/fit_forward.hpp
@@ -22,31 +22,26 @@ TRACCC_HOST_DEVICE inline void fit_forward(
     vecmem::device_vector<const unsigned int> param_ids(payload.param_ids_view);
     vecmem::device_vector<unsigned int> param_liveness(
         payload.param_liveness_view);
-    typename edm::track_fit_collection<
+    typename edm::track_fit_container<
         typename fitter_t::detector_type::algebra_type>::device
-        tracks(payload.tracks_view.tracks);
-    typename edm::track_state_collection<
-        typename fitter_t::detector_type::algebra_type>::device
-        track_states(payload.tracks_view.states);
-    measurement_collection_types::const_device measurements{
-        payload.tracks_view.measurements};
+        tracks(payload.tracks_view);
 
-    if (globalIndex >= tracks.size()) {
+    if (globalIndex >= tracks.tracks.size()) {
         return;
     }
 
     const unsigned int param_id = param_ids.at(globalIndex);
 
     fitter_t fitter(det, payload.field_data, cfg);
 
-    auto track = tracks.at(param_id);
+    auto track = tracks.tracks.at(param_id);
     auto params = track.params();
 
     // TODO: Merge into filter?
     inflate_covariance(params, fitter.config().covariance_inflation_factor);
 
     typename fitter_t::state fitter_state(
-        track, track_states, measurements,
+        track, tracks.states, tracks.measurements,
         *(payload.barcodes_view.ptr() + param_id));
 
     kalman_fitter_status fit_status = fitter.filter(params, fitter_state);
diff --git a/examples/run/cpu/seeding_example.cpp b/examples/run/cpu/seeding_example.cpp
@@ -243,12 +243,14 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
                 vecmem::get_data(measurements_per_event), evt_data);
 
             find_performance_writer.write(
-                vecmem::get_data(track_candidates),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             ar_performance_writer.write(
-                vecmem::get_data(track_candidates_ar),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates_ar),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             for (unsigned int i = 0; i < track_states.tracks.size(); i++) {
                 fit_performance_writer.write(
diff --git a/examples/run/cpu/seq_example.cpp b/examples/run/cpu/seq_example.cpp
@@ -341,11 +341,13 @@ int seq_run(const traccc::opts::input_data& input_opts,
                 vecmem::get_data(spacepoints_per_event),
                 vecmem::get_data(measurements_per_event), evt_data);
             find_performance_writer.write(
-                vecmem::get_data(track_candidates),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
             ar_performance_writer.write(
-                vecmem::get_data(resolved_track_candidates),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(resolved_track_candidates),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             for (unsigned int i = 0; i < track_states.tracks.size(); i++) {
                 fit_performance_writer.write(
diff --git a/examples/run/cpu/truth_finding_example.cpp b/examples/run/cpu/truth_finding_example.cpp
@@ -169,8 +169,9 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
 
         if (performance_opts.run) {
             find_performance_writer.write(
-                vecmem::get_data(track_candidates),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             for (std::size_t i = 0; i < n_fitted_tracks; i++) {
                 fit_performance_writer.write(
diff --git a/examples/run/cuda/seeding_example_cuda.cpp b/examples/run/cuda/seeding_example_cuda.cpp
@@ -432,8 +432,9 @@ int seq_run(const traccc::opts::track_seeding& seeding_opts,
                 vecmem::get_data(measurements_per_event), evt_data);
 
             find_performance_writer.write(
-                vecmem::get_data(track_candidates_cuda),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates_cuda),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             for (unsigned int i = 0; i < track_states_cuda.tracks.size(); i++) {
                 fit_performance_writer.write(
diff --git a/examples/run/cuda/truth_finding_example_cuda.cpp b/examples/run/cuda/truth_finding_example_cuda.cpp
@@ -299,8 +299,9 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
 
         if (performance_opts.run) {
             find_performance_writer.write(
-                vecmem::get_data(track_candidates_cuda),
-                vecmem::get_data(measurements_per_event), evt_data);
+                {vecmem::get_data(track_candidates_cuda),
+                 vecmem::get_data(measurements_per_event)},
+                evt_data);
 
             for (unsigned int i = 0; i < track_states_cuda.tracks.size(); i++) {
                 fit_performance_writer.write(
diff --git a/performance/include/traccc/efficiency/finding_performance_writer.hpp b/performance/include/traccc/efficiency/finding_performance_writer.hpp
@@ -15,10 +15,8 @@
 #include "traccc/utils/truth_matching_config.hpp"
 
 // Project include(s).
-#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/edm/track_candidate_container.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 #include "traccc/utils/event_data.hpp"
 
 // System include(s).
@@ -71,19 +69,13 @@ class finding_performance_writer : public messaging {
     /// Destructor
     ~finding_performance_writer();
 
-    void write(
-        const edm::track_candidate_collection<default_algebra>::const_view&
-            track_candidates_view,
-        const measurement_collection_types::const_view& measurements_view,
-        const event_data& evt_data);
-
-    void write(
-        const edm::track_fit_collection<default_algebra>::const_view&
-            track_fit_view,
-        const edm::track_state_collection<default_algebra>::const_view&
-            track_states_view,
-        const measurement_collection_types::const_view& measurements_view,
-        const event_data& evt_data);
+    void write(const edm::track_candidate_container<
+                   default_algebra>::const_view& track_candidates_view,
+               const event_data& evt_data);
+
+    void write(const edm::track_fit_container<default_algebra>::const_view&
+                   track_fit_view,
+               const event_data& evt_data);
 
     void finalize();
 
diff --git a/performance/src/efficiency/finding_performance_writer.cpp b/performance/src/efficiency/finding_performance_writer.cpp
@@ -91,28 +91,25 @@ namespace {
  * with its corresponding measurements.
  */
 std::vector<std::vector<measurement>> prepare_data(
-    const edm::track_candidate_collection<default_algebra>::const_view&
-        track_candidates_view,
-    const measurement_collection_types::const_view& measurements_view) {
+    const edm::track_candidate_container<default_algebra>::const_view&
+        track_candidates_view) {
     std::vector<std::vector<measurement>> result;
 
     // Iterate over the tracks.
-    const edm::track_candidate_collection<default_algebra>::const_device
+    const edm::track_candidate_container<default_algebra>::const_device
         track_candidates(track_candidates_view);
-    const measurement_collection_types::const_device measurements{
-        measurements_view};
 
-    const unsigned int n_tracks = track_candidates.size();
+    const unsigned int n_tracks = track_candidates.tracks.size();
     result.reserve(n_tracks);
 
     for (unsigned int i = 0; i < n_tracks; i++) {
         std::vector<measurement> m;
-        m.reserve(track_candidates.at(i).measurement_indices().size());
+        m.reserve(track_candidates.tracks.at(i).measurement_indices().size());
         const edm::track_candidate_collection<
             default_algebra>::const_device::const_proxy_type track =
-            track_candidates.at(i);
+            track_candidates.tracks.at(i);
         for (unsigned int midx : track.measurement_indices()) {
-            m.push_back(measurements.at(midx));
+            m.push_back(track_candidates.measurements.at(midx));
         }
         result.push_back(std::move(m));
     }
@@ -129,31 +126,25 @@ std::vector<std::vector<measurement>> prepare_data(
  * with its corresponding measurements.
  */
 std::vector<std::vector<measurement>> prepare_data(
-    const edm::track_fit_collection<default_algebra>::const_view&
-        track_fit_view,
-    const edm::track_state_collection<default_algebra>::const_view&
-        track_states_view,
-    const measurement_collection_types::const_view& measurements_view) {
+    const edm::track_fit_container<default_algebra>::const_view&
+        track_fit_view) {
     std::vector<std::vector<measurement>> result;
 
     // Set up the input containers.
-    const edm::track_fit_collection<default_algebra>::const_device track_fits(
+    const edm::track_fit_container<default_algebra>::const_device track_fits(
         track_fit_view);
-    const edm::track_state_collection<default_algebra>::const_device
-        track_states(track_states_view);
-    const measurement_collection_types::const_device measurements{
-        measurements_view};
 
     // Iterate over the tracks.
-    const unsigned int n_tracks = track_fits.size();
+    const unsigned int n_tracks = track_fits.tracks.size();
     result.reserve(n_tracks);
 
     for (unsigned int i = 0; i < n_tracks; i++) {
         std::vector<measurement> result_measurements;
-        result_measurements.reserve(track_fits.at(i).state_indices().size());
-        for (unsigned int st_idx : track_fits.state_indices().at(i)) {
-            result_measurements.push_back(
-                measurements.at(track_states.at(st_idx).measurement_index()));
+        result_measurements.reserve(
+            track_fits.tracks.at(i).state_indices().size());
+        for (unsigned int st_idx : track_fits.tracks.state_indices().at(i)) {
+            result_measurements.push_back(track_fits.measurements.at(
+                track_fits.states.at(st_idx).measurement_index()));
         }
         result.push_back(std::move(result_measurements));
     }
@@ -322,14 +313,13 @@ void finding_performance_writer::write_common(
 
 /// For track finding
 void finding_performance_writer::write(
-    const edm::track_candidate_collection<default_algebra>::const_view&
+    const edm::track_candidate_container<default_algebra>::const_view&
         track_candidates_view,
-    const measurement_collection_types::const_view& measurements_view,
     const event_data& evt_data) {
 
     // Set up the input containers.
     const edm::track_candidate_collection<default_algebra>::const_device
-        track_candidates(track_candidates_view);
+        track_candidates(track_candidates_view.tracks);
 
     const unsigned int n_tracks = track_candidates.size();
 
@@ -341,21 +331,16 @@ void finding_performance_writer::write(
     }
 
     std::vector<std::vector<measurement>> tracks =
-        prepare_data(track_candidates_view, measurements_view);
+        prepare_data(track_candidates_view);
     write_common(tracks, evt_data);
 }
 
 /// For ambiguity resolution
 void finding_performance_writer::write(
-    const edm::track_fit_collection<default_algebra>::const_view&
-        track_fit_view,
-    const edm::track_state_collection<default_algebra>::const_view&
-        track_states_view,
-    const measurement_collection_types::const_view& measurements_view,
+    const edm::track_fit_container<default_algebra>::const_view& track_fit_view,
     const event_data& evt_data) {
 
-    std::vector<std::vector<measurement>> tracks =
-        prepare_data(track_fit_view, track_states_view, measurements_view);
+    std::vector<std::vector<measurement>> tracks = prepare_data(track_fit_view);
     write_common(tracks, evt_data);
 }
 
