Plugins/Podio/PodioTrackContainer.hpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2023 CERN for the benefit of the Acts project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/EventData/MultiTrajectory.hpp"
0012 #include "Acts/EventData/ParticleHypothesis.hpp"
0013 #include "Acts/EventData/TrackContainer.hpp"
0014 #include "Acts/EventData/TrackStateProxy.hpp"
0015 #include "Acts/EventData/detail/DynamicColumn.hpp"
0016 #include "Acts/Plugins/Podio/PodioDynamicColumns.hpp"
0017 #include "Acts/Plugins/Podio/PodioUtil.hpp"
0018
0019 #pragma GCC diagnostic push
0020 #pragma GCC diagnostic ignored "-Wold-style-cast"
0021 #include "ActsPodioEdm/ParticleHypothesis.h"
0022 #include "ActsPodioEdm/Track.h"
0023 #include "ActsPodioEdm/TrackCollection.h"
0024 #include "ActsPodioEdm/TrackInfo.h"
0025 #pragma GCC diagnostic pop
0026
0027 #include <mutex>
0028 #include <stdexcept>
0029 #include <type_traits>
0030
0031 #include <podio/Frame.h>
0032
0033 namespace Acts {
0034
0035 class MutablePodioTrackContainer;
0036 class ConstPodioTrackContainer;
0037
0038 template <>
0039 struct IsReadOnlyTrackContainer<MutablePodioTrackContainer> : std::false_type {
0040 };
0041
0042 template <>
0043 struct IsReadOnlyTrackContainer<ConstPodioTrackContainer> : std::true_type {};
0044
0045 class PodioTrackContainerBase {
0046  public:
0047   using IndexType = MultiTrajectoryTraits::IndexType;
0048   static constexpr auto kInvalid = MultiTrajectoryTraits::kInvalid;
0049   static constexpr auto MeasurementSizeMax =
0050       MultiTrajectoryTraits::MeasurementSizeMax;
0051
0052   using Parameters =
0053       typename detail_lt::FixedSizeTypes<eBoundSize, false>::CoefficientsMap;
0054   using Covariance =
0055       typename detail_lt::FixedSizeTypes<eBoundSize, false>::CovarianceMap;
0056
0057   using ConstParameters =
0058       typename detail_lt::FixedSizeTypes<eBoundSize, true>::CoefficientsMap;
0059   using ConstCovariance =
0060       typename detail_lt::FixedSizeTypes<eBoundSize, true>::CovarianceMap;
0061
0062  protected:
0063   PodioTrackContainerBase(const PodioUtil::ConversionHelper& helper)
0064       : m_helper{helper} {}
0065
0066   template <bool EnsureConst, typename T>
0067   static std::any component_impl(T& instance, HashedString key,
0068                                  IndexType itrack) {
0069     using namespace Acts::HashedStringLiteral;
0070     if constexpr (EnsureConst) {
0071       static_assert(std::is_const_v<std::remove_reference_t<T>>,
0072                     "Is not const");
0073     }
0074
0075     using namespace Acts::HashedStringLiteral;
0076     auto track = instance.m_collection->at(itrack);
0077     std::conditional_t<EnsureConst, const ActsPodioEdm::TrackInfo*,
0078                        ActsPodioEdm::TrackInfo*>
0079         dataPtr;
0080     if constexpr (EnsureConst) {
0081       dataPtr = &track.getData();
0082     } else {
0083       dataPtr = &PodioUtil::getDataMutable(track);
0084     }
0085     auto& data = *dataPtr;
0086     switch (key) {
0087       case "tipIndex"_hash:
0088         return &data.tipIndex;
0089       case "stemIndex"_hash:
0090         return &data.stemIndex;
0091       case "params"_hash:
0092         return data.parameters.data();
0093       case "cov"_hash:
0094         return data.covariance.data();
0095       case "nMeasurements"_hash:
0096         return &data.nMeasurements;
0097       case "nHoles"_hash:
0098         return &data.nHoles;
0099       case "chi2"_hash:
0100         return &data.chi2;
0101       case "ndf"_hash:
0102         return &data.ndf;
0103       case "nOutliers"_hash:
0104         return &data.nOutliers;
0105       case "nSharedHits"_hash:
0106         return &data.nSharedHits;
0107       default:
0108         auto it = instance.m_dynamic.find(key);
0109         if (it == instance.m_dynamic.end()) {
0110           throw std::runtime_error("Unable to handle this component");
0111         }
0112
0113         std::conditional_t<EnsureConst,
0114                            const podio_detail::ConstDynamicColumnBase*,
0115                            podio_detail::DynamicColumnBase*>
0116             col = it->second.get();
0117         assert(col && "Dynamic column is null");
0118         return col->get(itrack);
0119     }
0120   }
0121
0122   template <typename T>
0123   static auto dynamicKeys_impl(T& instance) {
0124     using column_type =
0125         typename decltype(instance.m_dynamic)::mapped_type::element_type;
0126     return detail::DynamicKeyRange<column_type>{instance.m_dynamic.begin(),
0127                                                 instance.m_dynamic.end()};
0128   }
0129
0130   template <typename T>
0131   static ParticleHypothesis particleHypothesis_impl(T& instance,
0132                                                     IndexType itrack) {
0133     auto track = instance.m_collection->at(itrack);
0134     const auto& src = track.getParticleHypothesis();
0135     return ParticleHypothesis{static_cast<PdgParticle>(src.absPdg), src.mass,
0136                               src.absQ};
0137   }
0138
0139   static void populateSurfaceBuffer(
0140       const PodioUtil::ConversionHelper& helper,
0141       const ActsPodioEdm::TrackCollection& collection,
0142       std::vector<std::shared_ptr<const Surface>>& surfaces) noexcept {
0143     surfaces.reserve(collection.size());
0144     for (ActsPodioEdm::Track track : collection) {
0145       surfaces.push_back(PodioUtil::convertSurfaceFromPodio(
0146           helper, track.getReferenceSurface()));
0147     }
0148   }
0149
0150   std::reference_wrapper<const PodioUtil::ConversionHelper> m_helper;
0151   std::vector<std::shared_ptr<const Surface>> m_surfaces;
0152 };
0153
0154 class MutablePodioTrackContainer : public PodioTrackContainerBase {
0155  public:
0156   MutablePodioTrackContainer(const PodioUtil::ConversionHelper& helper)
0157       : PodioTrackContainerBase{helper},
0158         m_collection{std::make_unique<ActsPodioEdm::TrackCollection>()} {
0159     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0160   }
0161
0162   MutablePodioTrackContainer(const MutablePodioTrackContainer& other);
0163   MutablePodioTrackContainer(MutablePodioTrackContainer&& other) = default;
0164
0165   MutablePodioTrackContainer(const ConstPodioTrackContainer& other);
0166
0167   // BEGIN INTERFACE HELPER
0168
0169  private:
0170   std::shared_ptr<const Surface> getOrCreateSurface(IndexType itrack) {
0171     std::shared_ptr<const Surface>& ptr = m_surfaces.at(itrack);
0172     if (!ptr) {
0173       ActsPodioEdm::Track track = m_collection->at(itrack);
0174       ptr = PodioUtil::convertSurfaceFromPodio(m_helper,
0175                                                track.getReferenceSurface());
0176     }
0177     return ptr;
0178   }
0179
0180  public:
0181   std::any component_impl(HashedString key, IndexType itrack) {
0182     return PodioTrackContainerBase::component_impl<false>(*this, key, itrack);
0183   }
0184
0185   std::any component_impl(HashedString key, IndexType itrack) const {
0186     return PodioTrackContainerBase::component_impl<true>(*this, key, itrack);
0187   }
0188
0189   bool hasColumn_impl(HashedString key) const {
0190     return m_dynamic.find(key) != m_dynamic.end();
0191   }
0192
0193   std::size_t size_impl() const { return m_collection->size(); }
0194
0195   void clear() { m_collection->clear(); }
0196
0197   // END INTERFACE HELPER
0198
0199   const Surface* referenceSurface_impl(IndexType itrack) const {
0200     return m_surfaces.at(itrack).get();
0201   }
0202
0203   ParticleHypothesis particleHypothesis_impl(IndexType itrack) const {
0204     return PodioTrackContainerBase::particleHypothesis_impl(*this, itrack);
0205   }
0206
0207   void setReferenceSurface_impl(IndexType itrack,
0208                                 std::shared_ptr<const Surface> surface) {
0209     auto track = m_collection->at(itrack);
0210     track.setReferenceSurface(
0211         PodioUtil::convertSurfaceToPodio(m_helper, *surface));
0212     m_surfaces.at(itrack) = std::move(surface);
0213   }
0214
0215  public:
0216   // BEGIN INTERFACE
0217
0218   IndexType addTrack_impl() {
0219     auto track = m_collection->create();
0220     PodioUtil::getReferenceSurfaceMutable(track).surfaceType =
0221         PodioUtil::kNoSurface;
0222     m_surfaces.emplace_back();
0223     for (const auto& [key, vec] : m_dynamic) {
0224       vec->add();
0225     }
0226     return m_collection->size() - 1;
0227   };
0228
0229   void removeTrack_impl(IndexType itrack);
0230
0231   template <typename T>
0232   constexpr void addColumn_impl(std::string_view key) {
0233     Acts::HashedString hashedKey = hashString(key);
0234     m_dynamic.insert(
0235         {hashedKey, std::make_unique<podio_detail::DynamicColumn<T>>(key)});
0236   }
0237
0238   Parameters parameters(IndexType itrack) {
0239     return Parameters{
0240         PodioUtil::getDataMutable(m_collection->at(itrack)).parameters.data()};
0241   }
0242
0243   ConstParameters parameters(IndexType itrack) const {
0244     return ConstParameters{
0245         m_collection->at(itrack).getData().parameters.data()};
0246   }
0247
0248   Covariance covariance(IndexType itrack) {
0249     return Covariance{
0250         PodioUtil::getDataMutable(m_collection->at(itrack)).covariance.data()};
0251   }
0252
0253   ConstCovariance covariance(IndexType itrack) const {
0254     return ConstCovariance{
0255         m_collection->at(itrack).getData().covariance.data()};
0256   }
0257
0258   void copyDynamicFrom_impl(IndexType dstIdx, HashedString key,
0259                             const std::any& srcPtr) {
0260     auto it = m_dynamic.find(key);
0261     if (it == m_dynamic.end()) {
0262       throw std::invalid_argument{
0263           "Destination container does not have matching dynamic column"};
0264     }
0265
0266     it->second->copyFrom(dstIdx, srcPtr);
0267   }
0268
0269   void ensureDynamicColumns_impl(const MutablePodioTrackContainer& other);
0270
0271   void reserve(IndexType /*size*/) {}
0272
0273   ActsPodioEdm::TrackCollection& trackCollection() { return *m_collection; }
0274
0275   void releaseInto(podio::Frame& frame, const std::string& suffix = "") {
0276     std::string s = suffix;
0277     if (!s.empty()) {
0278       s = "_" + s;
0279     }
0280     frame.put(std::move(m_collection), "tracks" + s);
0281     m_surfaces.clear();
0282
0283     for (const auto& [key, col] : m_dynamic) {
0284       col->releaseInto(frame, "tracks" + s + "_extra__");
0285     }
0286   }
0287
0288   detail::DynamicKeyRange<podio_detail::DynamicColumnBase> dynamicKeys_impl()
0289       const {
0290     return PodioTrackContainerBase::dynamicKeys_impl(*this);
0291   }
0292
0293   void setParticleHypothesis_impl(
0294       IndexType itrack, const ParticleHypothesis& particleHypothesis) {
0295     ActsPodioEdm::ParticleHypothesis pHypo;
0296     pHypo.absPdg = particleHypothesis.absolutePdg();
0297     pHypo.mass = particleHypothesis.mass();
0298     pHypo.absQ = particleHypothesis.absoluteCharge();
0299     m_collection->at(itrack).setParticleHypothesis(pHypo);
0300   }
0301
0302   // END INTERFACE
0303
0304  private:
0305   friend PodioTrackContainerBase;
0306
0307   std::unique_ptr<ActsPodioEdm::TrackCollection> m_collection;
0308   std::vector<HashedString> m_dynamicKeys;
0309   std::unordered_map<HashedString,
0310                      std::unique_ptr<podio_detail::DynamicColumnBase>>
0311       m_dynamic;
0312 };
0313
0314 static_assert(
0315     TrackContainerBackend<MutablePodioTrackContainer>,
0316     "MutablePodioTrackContainer does not fulfill TrackContainerBackend");
0317
0318 class ConstPodioTrackContainer : public PodioTrackContainerBase {
0319  public:
0320   ConstPodioTrackContainer(const PodioUtil::ConversionHelper& helper,
0321                            const ActsPodioEdm::TrackCollection& collection)
0322       : PodioTrackContainerBase{helper}, m_collection{&collection} {
0323     // Not much we can do to recover dynamic columns here
0324     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0325   }
0326
0327   ConstPodioTrackContainer(const PodioUtil::ConversionHelper& helper,
0328                            const podio::Frame& frame,
0329                            const std::string& suffix = "")
0330       : PodioTrackContainerBase{helper} {
0331     std::string s = suffix.empty() ? suffix : "_" + suffix;
0332     std::string tracksKey = "tracks" + s;
0333
0334     std::vector<std::string> available = frame.getAvailableCollections();
0335     if (std::find(available.begin(), available.end(), tracksKey) ==
0336         available.end()) {
0337       throw std::runtime_error{"Track collection '" + tracksKey +
0338                                "' not found in frame"};
0339     }
0340
0341     const auto* collection = frame.get(tracksKey);
0342
0343     if (const auto* d =
0344             dynamic_cast<const ActsPodioEdm::TrackCollection*>(collection);
0345         d != nullptr) {
0346       m_collection = d;
0347     } else {
0348       throw std::runtime_error{"Unable to get collection " + tracksKey};
0349     }
0350
0351     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0352
0353     podio_detail::recoverDynamicColumns(frame, tracksKey, m_dynamic);
0354   }
0355
0356   std::any component_impl(HashedString key, IndexType itrack) const {
0357     return PodioTrackContainerBase::component_impl<true>(*this, key, itrack);
0358   }
0359
0360   bool hasColumn_impl(HashedString key) const {
0361     return m_dynamic.find(key) != m_dynamic.end();
0362   }
0363
0364   std::size_t size_impl() const { return m_collection->size(); }
0365
0366   const Surface* referenceSurface_impl(IndexType itrack) const {
0367     return m_surfaces.at(itrack).get();
0368   }
0369
0370   ParticleHypothesis particleHypothesis_impl(IndexType itrack) const {
0371     return PodioTrackContainerBase::particleHypothesis_impl(*this, itrack);
0372   }
0373
0374   ConstParameters parameters(IndexType itrack) const {
0375     return ConstParameters{
0376         m_collection->at(itrack).getData().parameters.data()};
0377   }
0378
0379   ConstCovariance covariance(IndexType itrack) const {
0380     return ConstCovariance{
0381         m_collection->at(itrack).getData().covariance.data()};
0382   }
0383
0384   const ActsPodioEdm::TrackCollection& trackCollection() {
0385     return *m_collection;
0386   }
0387
0388   detail::DynamicKeyRange<podio_detail::ConstDynamicColumnBase>
0389   dynamicKeys_impl() const {
0390     return PodioTrackContainerBase::dynamicKeys_impl(*this);
0391   }
0392
0393  private:
0394   friend PodioTrackContainerBase;
0395
0396   const ActsPodioEdm::TrackCollection* m_collection;
0397   std::unordered_map<HashedString,
0398                      std::unique_ptr<podio_detail::ConstDynamicColumnBase>>
0399       m_dynamic;
0400   std::vector<HashedString> m_dynamicKeys;
0401 };
0402
0403 static_assert(
0404     ConstTrackContainerBackend<ConstPodioTrackContainer>,
0405     "ConstPodioTrackContainer does not fulfill ConstTrackContainerBackend");
0406
0407 }  //  namespace Acts