src/fitting/kalman_fitting_algorithm.cpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2022-2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // Local include(s).
0009 #include "traccc/fitting/device/kalman_fitting_algorithm.hpp"
0010
0011 // VecMem include(s).
0012 #include <vecmem/containers/data/jagged_vector_buffer.hpp>
0013
0014 // System include(s).
0015 #include <numeric>
0016 #include <vector>
0017
0018 namespace traccc::device {
0019
0020 struct kalman_fitting_algorithm::data {
0021   /// @name Configuration object(s)
0022   /// @{
0023
0024   /// Configuration for the fitting algorithm
0025   fitting_config m_config;
0026
0027   /// @}
0028
0029 };  // struct kalman_fitting_algorithm::data
0030
0031 kalman_fitting_algorithm::fit_payload::fit_payload(const detector_buffer& det,
0032                                                    const magnetic_field& f)
0033     : detector(det), field(f) {}
0034
0035 kalman_fitting_algorithm::kalman_fitting_algorithm(
0036     const config_type& config, const traccc::memory_resource& mr,
0037     const vecmem::copy& copy, std::unique_ptr<const Logger> logger)
0038     : messaging(std::move(logger)),
0039       algorithm_base{mr, copy},
0040       m_data{std::make_unique<data>(config)} {}
0041
0042 kalman_fitting_algorithm::~kalman_fitting_algorithm() = default;
0043
0044 kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()(
0045     const detector_buffer& det, const magnetic_field& field,
0046     const edm::track_container<default_algebra>::const_view& input_tracks)
0047     const {
0048   // Get the number of tracks and the number of constituens (measurements)
0049   // in each track. Note that the number of tracks is not "resizable". That
0050   // we can just get directly from the view object. But the number of
0051   // constituents need to be taken from the buffer itself.
0052   assert(copy().get_size(input_tracks.tracks) ==
0053          input_tracks.tracks.capacity());
0054   const edm::track_collection<default_algebra>::const_view::size_type n_tracks =
0055       input_tracks.tracks.capacity();
0056   if (n_tracks == 0) {
0057     // Return early, if there are no tracks.
0058     return {};
0059   }
0060   std::vector<unsigned int> candidate_sizes;
0061   if (mr().host) {
0062     vecmem::async_sizes sizes =
0063         copy().get_sizes(input_tracks.tracks, *(mr().host));
0064     // Here we could give control back to the caller, once our code allows
0065     // for it. (coroutines...)
0066     auto& temp = sizes.get();
0067     candidate_sizes = {temp.begin(), temp.end()};
0068   } else {
0069     candidate_sizes = copy().get_sizes(input_tracks.tracks);
0070   }
0071
0072   // Get the total number of states (measurements) to fit.
0073   const unsigned int n_states =
0074       std::accumulate(candidate_sizes.begin(), candidate_sizes.end(), 0u);
0075
0076   // Create the result buffer.
0077   edm::track_container<default_algebra>::buffer output_tracks{
0078       {candidate_sizes, mr().main, mr().host,
0079        vecmem::data::buffer_type::resizable},
0080       {n_states, mr().main, vecmem::data::buffer_type::resizable},
0081       input_tracks.measurements};
0082   copy().setup(output_tracks.tracks)->ignore();
0083   copy().setup(output_tracks.states)->ignore();
0084
0085   // Create the order to fit the tracks in.
0086   vecmem::data::vector_buffer<device::sort_key> track_sort_keys(n_tracks,
0087                                                                 mr().main);
0088   vecmem::data::vector_buffer<unsigned int> track_indices{n_tracks, mr().main};
0089   prepare_track_fit_order(input_tracks.tracks, track_sort_keys, track_indices);
0090
0091   // Create the buffer(s) used during the fitting.
0092   vecmem::data::vector_buffer<unsigned int> track_liveness(n_tracks, mr().main);
0093
0094   // Run "fitting prelude" kernel.
0095   fit_prelude_kernel(
0096       {track_indices, input_tracks, output_tracks, track_liveness});
0097
0098   // Calculate the number of surfaces to use during the fit for each track.
0099   // Then create the concrete buffer such that it could be passed to the
0100   // fitting functions through a polymorphic pointer.
0101   std::vector<unsigned int> n_surfaces(candidate_sizes.size());
0102   std::transform(candidate_sizes.begin(), candidate_sizes.end(),
0103                  n_surfaces.begin(), [&](const unsigned int sz) {
0104                    return std::max(
0105                        sz * m_data->m_config.surface_sequence_size_factor,
0106                        m_data->m_config.min_surface_sequence_capacity);
0107                  });
0108
0109   // Prepare the payload for the fitting kernel(s).
0110   const fit_payload payload = prepare_fit_payload(
0111       det, field, n_surfaces, {track_indices, track_liveness, output_tracks});
0112
0113   // Run the iterative track fitting.
0114   for (std::size_t i = 0; i < m_data->m_config.n_iterations; ++i) {
0115     fit_forward_kernel(m_data->m_config, payload);
0116     fit_backward_kernel(m_data->m_config, payload);
0117   }
0118
0119   // Return the fitted tracks.
0120   return output_tracks;
0121 }
0122
0123 }  // namespace traccc::device