Plugins/Podio/PodioTrackContainer.hpp

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