traccc/finding/track_state_candidate.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 #pragma once
0009
0010 // Project include(s)
0011 #include "traccc/definitions/primitives.hpp"
0012 #include "traccc/edm/measurement_collection.hpp"
0013 #include "traccc/edm/track_parameters.hpp"
0014 #include "traccc/edm/track_state_helpers.hpp"
0015
0016 // detray include(s)
0017 #include <detray/utils/ranges/detail/iterator_functions.hpp>
0018
0019 // System include(s)
0020 #include <limits>
0021
0022 namespace traccc {
0023
0024 /// Data payload that is accumulated during the Kalman track follower
0025 struct track_state_candidate {
0026   // The index of a matched measurement
0027   unsigned int measurement_index{std::numeric_limits<unsigned int>::max()};
0028 };
0029
0030 /// Kalman data payload extended by the filtered track parameters
0031 template <detray::concepts::algebra algebra_t>
0032 struct filtered_track_state_candidate : public track_state_candidate {
0033   using scalar_t = detray::dscalar<algebra_t>;
0034   using bound_parameters_t = traccc::bound_track_parameters<algebra_t>;
0035
0036   constexpr filtered_track_state_candidate() = default;
0037
0038   TRACCC_HOST_DEVICE
0039   filtered_track_state_candidate(unsigned int meas_idx, scalar_t chi2,
0040                                  const bound_parameters_t& flt_params)
0041       : track_state_candidate{meas_idx},
0042         filtered_chi2{chi2},
0043         filtered_params{flt_params} {
0044     assert(!flt_params.is_invalid());
0045   }
0046
0047   // The filtered chi2
0048   scalar_t filtered_chi2;
0049
0050   // The filtered track parameters corresponding to the measurement
0051   bound_parameters_t filtered_params{};
0052 };
0053
0054 /// Full Kalman data payload
0055 template <detray::concepts::algebra algebra_t>
0056 struct full_track_state_candidate
0057     : public filtered_track_state_candidate<algebra_t> {
0058   using scalar_t = detray::dscalar<algebra_t>;
0059   using bound_parameters_t = traccc::bound_track_parameters<algebra_t>;
0060   using full_jacobian_t = traccc::bound_matrix<algebra_t>;
0061
0062   constexpr full_track_state_candidate() = default;
0063
0064   TRACCC_HOST_DEVICE
0065   full_track_state_candidate(unsigned int meas_idx, scalar_t chi2,
0066                              const bound_parameters_t& flt_params,
0067                              const bound_parameters_t& pred_params,
0068                              const full_jacobian_t& jac)
0069       : filtered_track_state_candidate<algebra_t>(meas_idx, chi2, flt_params),
0070         predicted_params{pred_params},
0071         jacobian{jac} {
0072     assert(!flt_params.is_invalid());
0073     assert(!pred_params.is_invalid());
0074     assert(jac != matrix::zero<full_jacobian_t>());
0075   }
0076
0077   // The predicted track parameters at the measurement surface
0078   bound_parameters_t predicted_params{};
0079
0080   // The full Jacobian from the previous sensitive surface to the current one
0081   full_jacobian_t jacobian{};
0082 };
0083
0084 /// Struct to simplify the handling of different data recording modes during
0085 /// track following
0086 template <detray::concepts::algebra algebra_t>
0087 struct track_state_candidate_data {
0088   TRACCC_HOST_DEVICE
0089   track_state_candidate_data(
0090       const smoother_type mode, const unsigned int offset,
0091       vecmem::data::vector_view<track_state_candidate> track_cand_view,
0092       vecmem::data::vector_view<filtered_track_state_candidate<algebra_t>>
0093           filtered_track_cand_view,
0094       vecmem::data::vector_view<full_track_state_candidate<algebra_t>>
0095           full_track_cand_view)
0096       : m_track_cands{track_cand_view},
0097         m_filtered_track_cands{filtered_track_cand_view},
0098         m_full_track_cands{full_track_cand_view} {
0099     // Set the data pointer according to the offset
0100     switch (mode) {
0101       case smoother_type::e_mbf: {
0102         assert(offset < m_full_track_cands.capacity());
0103         m_track_cand_ptr = static_cast<void*>(
0104             detray::ranges::detail::next(m_full_track_cands.data(), offset));
0105         break;
0106       }
0107       case smoother_type::e_kalman: {
0108         assert(offset < m_filtered_track_cands.capacity());
0109         m_track_cand_ptr = static_cast<void*>(detray::ranges::detail::next(
0110             m_filtered_track_cands.data(), offset));
0111         break;
0112       }
0113       case smoother_type::e_none: {
0114         assert(offset < m_track_cands.capacity());
0115         m_track_cand_ptr = static_cast<void*>(
0116             detray::ranges::detail::next(m_track_cands.data(), offset));
0117         break;
0118       }
0119       default: {
0120         TRACCC_ERROR_HOST_DEVICE("Unknown smoother option");
0121       }
0122     }
0123
0124     assert(m_track_cand_ptr);
0125   }
0126
0127   /// @returns the data pointer for the track and data collection mode
0128   TRACCC_HOST_DEVICE
0129   void* ptr() const { return m_track_cand_ptr; }
0130
0131  private:
0132   void* m_track_cand_ptr{nullptr};
0133
0134   /// Underlying data collections
0135   /// @{
0136   vecmem::device_vector<track_state_candidate> m_track_cands;
0137   vecmem::device_vector<filtered_track_state_candidate<algebra_t>>
0138       m_filtered_track_cands;
0139   vecmem::device_vector<full_track_state_candidate<algebra_t>>
0140       m_full_track_cands;
0141   /// @}
0142 };
0143
0144 /// Add a new track state to the track container
0145 template <detray::concepts::algebra algebra_t>
0146 TRACCC_HOST_DEVICE inline void make_track_state_candidate(
0147     void* track_cand_ptr, const smoother_type mode, const int idx,
0148     const candidate_measurement& cand,
0149     const bound_track_parameters<algebra_t>& bound_param) {
0150   switch (mode) {
0151     case smoother_type::e_mbf: {
0152       auto* data_ptr =
0153           static_cast<full_track_state_candidate<algebra_t>*>(track_cand_ptr);
0154       detray::ranges::detail::advance(data_ptr, idx);
0155       assert(data_ptr);
0156
0157       // TODO: Get proper Jacobian
0158       traccc::bound_track_parameters<algebra_t> predicted_params{};
0159       traccc::bound_matrix<algebra_t> full_jac{};
0160       *data_ptr = {cand.meas_idx, cand.chi2, bound_param, predicted_params,
0161                    full_jac};
0162       break;
0163     }
0164     case smoother_type::e_kalman: {
0165       auto* data_ptr = static_cast<filtered_track_state_candidate<algebra_t>*>(
0166           track_cand_ptr);
0167       detray::ranges::detail::advance(data_ptr, idx);
0168       assert(data_ptr);
0169
0170       *data_ptr = {cand.meas_idx, cand.chi2, bound_param};
0171       break;
0172     }
0173     case smoother_type::e_none: {
0174       auto* data_ptr = static_cast<track_state_candidate*>(track_cand_ptr);
0175       detray::ranges::detail::advance(data_ptr, idx);
0176       assert(data_ptr);
0177
0178       *data_ptr = {cand.meas_idx};
0179       break;
0180     }
0181     default: {
0182       TRACCC_FATAL_HOST_DEVICE(
0183           "Unknown data coll. type in measurement updater");
0184     }
0185   }
0186 }
0187
0188 /// Add a new track state to the track container
0189 template <detray::concepts::algebra algebra_t, typename BASE>
0190 TRACCC_HOST_DEVICE inline void track_state_from_candidate(
0191     void* track_cand_ptr, const smoother_type mode, const unsigned int link_idx,
0192     typename edm::measurement_collection::const_device measurements,
0193     edm::track<BASE> track,
0194     typename edm::track_state_collection<algebra_t>::view track_states_view) {
0195   typename edm::track_state_collection<algebra_t>::device track_states(
0196       track_states_view);
0197
0198   // The track_cand_ptr points at the first track state
0199   switch (mode) {
0200     case smoother_type::e_mbf: {
0201       auto* data_ptr =
0202           static_cast<full_track_state_candidate<algebra_t>*>(track_cand_ptr);
0203       detray::ranges::detail::advance(data_ptr, link_idx);
0204       assert(data_ptr);
0205       assert(data_ptr->measurement_index < measurements.size());
0206
0207       TRACCC_VERBOSE_DEVICE("-> Measurement %d (chi2 = %f)",
0208                             data_ptr->measurement_index,
0209                             data_ptr->filtered_chi2);
0210
0211       assert(link_idx < track.constituent_links().size());
0212
0213       const unsigned int track_state_index =
0214           track_states.push_back(edm::make_track_state<algebra_t>(
0215               measurements, data_ptr->measurement_index));
0216       auto track_state = track_states.at(track_state_index);
0217
0218       track.constituent_links().at(link_idx) =
0219           traccc::edm::track_constituent_link{
0220               edm::track_constituent_link::track_state, track_state_index};
0221
0222       track_state.set_hole(false);
0223       track_state.filtered_params() = data_ptr->filtered_params;
0224       track_state.filtered_chi2() = data_ptr->filtered_chi2;
0225       // TODO: Fake it till you make it: Remove again!
0226       track_state.smoothed_params() = data_ptr->filtered_params;
0227       track_state.smoothed_chi2() = data_ptr->filtered_chi2;
0228       track_state.backward_chi2() = data_ptr->filtered_chi2;
0229       track_state.set_smoothed(true);
0230       break;
0231     }
0232     case smoother_type::e_kalman: {
0233       auto* data_ptr = static_cast<filtered_track_state_candidate<algebra_t>*>(
0234           track_cand_ptr);
0235       detray::ranges::detail::advance(data_ptr, link_idx);
0236       assert(data_ptr);
0237       assert(data_ptr->measurement_index < measurements.size());
0238
0239       TRACCC_VERBOSE_DEVICE("-> Measurement %d (chi2 = %f)",
0240                             data_ptr->measurement_index,
0241                             data_ptr->filtered_chi2);
0242
0243       assert(link_idx < track.constituent_links().size());
0244
0245       const unsigned int track_state_index =
0246           track_states.push_back(edm::make_track_state<algebra_t>(
0247               measurements, data_ptr->measurement_index));
0248       auto track_state = track_states.at(track_state_index);
0249
0250       track.constituent_links().at(link_idx) =
0251           traccc::edm::track_constituent_link{
0252               edm::track_constituent_link::track_state, track_state_index};
0253
0254       track_state.set_hole(false);
0255       track_state.filtered_params() = data_ptr->filtered_params;
0256       track_state.filtered_chi2() = data_ptr->filtered_chi2;
0257
0258       break;
0259     }
0260     case smoother_type::e_none: {
0261       auto* data_ptr = static_cast<track_state_candidate*>(track_cand_ptr);
0262       detray::ranges::detail::advance(data_ptr, link_idx);
0263       assert(data_ptr);
0264       assert(data_ptr->measurement_index < measurements.size());
0265
0266       TRACCC_VERBOSE_DEVICE("-> Measurement %d", data_ptr->measurement_index);
0267
0268       assert(link_idx < track.constituent_links().size());
0269
0270       track.constituent_links().at(link_idx) =
0271           traccc::edm::track_constituent_link{
0272               edm::track_constituent_link::measurement,
0273               data_ptr->measurement_index};
0274       break;
0275     }
0276     default: {
0277       TRACCC_ERROR_HOST_DEVICE("Unknown smoother option");
0278     }
0279   }
0280 }
0281
0282 }  // namespace traccc