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