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