traccc/edm/track_container.hpp

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 #pragma once
0009
0010 // Local include(s).
0011 #include "traccc/definitions/qualifiers.hpp"
0012 #include "traccc/edm/measurement_collection.hpp"
0013 #include "traccc/edm/track_collection.hpp"
0014 #include "traccc/edm/track_state_collection.hpp"
0015
0016 // VecMem include(s).
0017 #include <vecmem/memory/memory_resource.hpp>
0018
0019 namespace traccc::edm {
0020
0021 /// Types collecting all the collections used to describe (fitted) tracks
0022 template <typename ALGEBRA>
0023 struct track_container {
0024   struct host {
0025     /// Constructor using a memory resource
0026     explicit host(vecmem::memory_resource& mr,
0027                   measurement_collection::const_view meas = {})
0028         : tracks{mr}, states{mr}, measurements{meas} {}
0029
0030     /// The tracks
0031     track_collection<ALGEBRA>::host tracks;
0032     /// The track states used by the tracks
0033     track_state_collection<ALGEBRA>::host states;
0034     /// The measurements used by the tracks
0035     measurement_collection::const_view measurements;
0036   };
0037
0038   struct buffer {
0039     /// The tracks
0040     track_collection<ALGEBRA>::buffer tracks;
0041     /// The track states used by the tracks
0042     track_state_collection<ALGEBRA>::buffer states;
0043     /// The measurements used by the tracks
0044     measurement_collection::const_view measurements;
0045   };
0046
0047   struct data {
0048     /// Constructor from a host container
0049     explicit data(host& h)
0050         : tracks{vecmem::get_data(h.tracks)},
0051           states{vecmem::get_data(h.states)},
0052           measurements{h.measurements} {}
0053     /// The tracks
0054     track_collection<ALGEBRA>::data tracks;
0055     /// The track states used by the tracks
0056     track_state_collection<ALGEBRA>::data states;
0057     /// The measurements used by the tracks
0058     measurement_collection::const_view measurements;
0059   };
0060
0061   struct const_data {
0062     /// Constructor from a host container
0063     explicit const_data(const host& h)
0064         : tracks{vecmem::get_data(h.tracks)},
0065           states{vecmem::get_data(h.states)},
0066           measurements{h.measurements} {}
0067     /// The tracks
0068     track_collection<ALGEBRA>::const_data tracks;
0069     /// The track states used by the tracks
0070     track_state_collection<ALGEBRA>::const_data states;
0071     /// The measurements used by the tracks
0072     measurement_collection::const_view measurements;
0073   };
0074
0075   struct view {
0076     /// Default constructor
0077     view() = default;
0078     /// Constructor from a buffer
0079     TRACCC_HOST_DEVICE
0080     view(const buffer& b)
0081         : tracks{b.tracks}, states{b.states}, measurements{b.measurements} {}
0082     /// Constructor from a data object
0083     TRACCC_HOST_DEVICE
0084     view(const data& d)
0085         : tracks{d.tracks}, states{d.states}, measurements{d.measurements} {}
0086     /// The tracks
0087     track_collection<ALGEBRA>::view tracks;
0088     /// The track states used by the tracks
0089     track_state_collection<ALGEBRA>::view states;
0090     /// The measurements used by the tracks
0091     measurement_collection::const_view measurements;
0092   };
0093
0094   struct const_view {
0095     /// Default constructor
0096     const_view() = default;
0097     /// Constructor from a buffer
0098     TRACCC_HOST_DEVICE
0099     const_view(const buffer& b)
0100         : tracks{b.tracks}, states{b.states}, measurements{b.measurements} {}
0101     /// Constructor from a const_data object
0102     TRACCC_HOST_DEVICE
0103     const_view(const const_data& d)
0104         : tracks{d.tracks}, states{d.states}, measurements{d.measurements} {}
0105     /// The tracks
0106     track_collection<ALGEBRA>::const_view tracks;
0107     /// The track states used by the tracks
0108     track_state_collection<ALGEBRA>::const_view states;
0109     /// The measurements used by the tracks
0110     measurement_collection::const_view measurements;
0111   };
0112
0113   struct device {
0114     /// Constructor from a view
0115     TRACCC_HOST_DEVICE
0116     explicit device(const view& v)
0117         : tracks{v.tracks}, states{v.states}, measurements{v.measurements} {}
0118     /// The tracks
0119     track_collection<ALGEBRA>::device tracks;
0120     /// The track states used by the tracks
0121     track_state_collection<ALGEBRA>::device states;
0122     /// The measurements used by the tracks
0123     measurement_collection::const_device measurements;
0124   };
0125
0126   struct const_device {
0127     /// Constructor from a view
0128     TRACCC_HOST_DEVICE
0129     explicit const_device(const const_view& v)
0130         : tracks{v.tracks}, states{v.states}, measurements{v.measurements} {}
0131     /// The tracks
0132     track_collection<ALGEBRA>::const_device tracks;
0133     /// The track states used by the tracks
0134     track_state_collection<ALGEBRA>::const_device states;
0135     /// The measurements used by the tracks
0136     measurement_collection::const_device measurements;
0137   };
0138
0139 };  // struct track_container
0140
0141 /// Convenience function to print the entire track information (const device)
0142 ///
0143 /// @param track_cont the track container, holding track and states/measurements
0144 /// @param trk_idx the index into the track collection to fetch the track that
0145 /// should be printed
0146 template <typename algebra_t>
0147 TRACCC_HOST inline void print_track(
0148     typename edm::track_container<algebra_t>::const_device& track_cont,
0149     unsigned int trk_idx) {
0150   // Print the general track information
0151   const auto trk = track_cont.tracks.at(trk_idx);
0152   std::cout << "TRACK IDX " << trk_idx << ":\n" << trk << std::endl;
0153
0154   if (trk.constituent_links().empty()) {
0155     std::cout << "Track does not contain data" << std::endl;
0156     return;
0157   }
0158
0159   // Print the track states or measurements that belong to the trac
0160   if (trk.constituent_links().at(0).type ==
0161       edm::track_constituent_link::constituent_type::track_state) {
0162     for (const auto& link : trk.constituent_links()) {
0163       std::cout << "-> State " << link.index << ":\n"
0164                 << track_cont.states.at(link.index) << std::endl;
0165     }
0166   } else {
0167     edm::measurement_collection::const_device measurements{
0168         track_cont.measurements};
0169     for (const auto& link : trk.constituent_links()) {
0170       std::cout << "-> Measurement " << link.index << ":\n"
0171                 << measurements.at(link.index) << std::endl;
0172     }
0173   }
0174 }
0175
0176 /// Convenience function to print the entire track information (host)
0177 ///
0178 /// @param track_cont the track container, holding track and states/measurements
0179 /// @param trk_idx the index into the track collection to fetch the track that
0180 /// should be printed
0181 template <typename algebra_t>
0182 TRACCC_HOST inline void print_track(
0183     typename edm::track_container<algebra_t>::host& track_cont,
0184     unsigned int trk_idx) {
0185   typename edm::track_container<algebra_t>::const_data const_tracks_data{
0186       track_cont};
0187   typename edm::track_container<algebra_t>::const_view const_tracks_view{
0188       const_tracks_data};
0189   typename edm::track_container<algebra_t>::const_device const_device_tracks{
0190       const_tracks_view};
0191
0192   edm::print_track<algebra_t>(const_device_tracks, trk_idx);
0193 }
0194
0195 }  // namespace traccc::edm