ActsPlugins/EDM4hep/PodioTrackStateContainer.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/MultiTrajectoryBackendConcept.hpp"
0013 #include "Acts/EventData/TrackContainer.hpp"
0014 #include "Acts/EventData/TrackStatePropMask.hpp"
0015 #include "Acts/EventData/Types.hpp"
0016 #include "Acts/EventData/detail/DynamicKeyIterator.hpp"
0017 #include "Acts/Utilities/HashedString.hpp"
0018 #include "Acts/Utilities/Helpers.hpp"
0019 #include "ActsPlugins/EDM4hep/PodioDynamicColumns.hpp"
0020 #include "ActsPlugins/EDM4hep/PodioTrackContainer.hpp"
0021 #include "ActsPlugins/EDM4hep/PodioUtil.hpp"
0022
0023 #pragma GCC diagnostic push
0024 #pragma GCC diagnostic ignored "-Wold-style-cast"
0025 #include "ActsPodioEdm/BoundParametersCollection.h"
0026 #include "ActsPodioEdm/JacobianCollection.h"
0027 #include "ActsPodioEdm/TrackStateCollection.h"
0028 #include "ActsPodioEdm/TrackStateInfo.h"
0029 #pragma GCC diagnostic pop
0030
0031 #include <any>
0032 #include <memory>
0033 #include <stdexcept>
0034 #include <tuple>
0035 #include <type_traits>
0036
0037 #include <podio/CollectionBase.h>
0038 #include <podio/Frame.h>
0039
0040 namespace ActsPlugins {
0041
0042 class MutablePodioTrackStateContainer;
0043 class ConstPodioTrackStateContainer;
0044
0045 }  // namespace ActsPlugins
0046
0047 namespace Acts {
0048 template <>
0049 struct IsReadOnlyMultiTrajectory<ActsPlugins::ConstPodioTrackStateContainer>
0050     : std::true_type {};
0051
0052 template <>
0053 struct IsReadOnlyMultiTrajectory<ActsPlugins::MutablePodioTrackStateContainer>
0054     : std::false_type {};
0055 }  // namespace Acts
0056
0057 namespace ActsPlugins {
0058
0059 class PodioTrackStateContainerBase {
0060  public:
0061   using Parameters =
0062       typename Acts::detail_lt::FixedSizeTypes<Acts::eBoundSize,
0063                                                false>::CoefficientsMap;
0064   using Covariance =
0065       typename Acts::detail_lt::FixedSizeTypes<Acts::eBoundSize,
0066                                                false>::CovarianceMap;
0067
0068   using ConstParameters =
0069       typename Acts::detail_lt::FixedSizeTypes<Acts::eBoundSize,
0070                                                true>::CoefficientsMap;
0071   using ConstCovariance =
0072       typename Acts::detail_lt::FixedSizeTypes<Acts::eBoundSize,
0073                                                true>::CovarianceMap;
0074
0075  protected:
0076   template <typename T>
0077   static constexpr bool has_impl(T& instance, Acts::HashedString key,
0078                                  Acts::TrackIndexType istate) {
0079     constexpr auto kInvalid = Acts::MultiTrajectoryTraits::kInvalid;
0080     using namespace Acts::HashedStringLiteral;
0081     auto trackState = instance.m_collection->at(istate);
0082     const auto& data = trackState.getData();
0083     switch (key) {
0084       case "predicted"_hash:
0085         return data.ipredicted != kInvalid;
0086       case "filtered"_hash:
0087         return data.ifiltered != kInvalid;
0088       case "smoothed"_hash:
0089         return data.ismoothed != kInvalid;
0090       case "calibrated"_hash:
0091         return data.measdim != kInvalid;
0092       case "calibratedCov"_hash:
0093         return data.measdim != kInvalid;
0094       case "jacobian"_hash:
0095         return data.ijacobian != kInvalid;
0096       case "projector"_hash:
0097         return data.hasProjector;
0098       case "uncalibratedSourceLink"_hash:
0099         return data.uncalibratedIdentifier != PodioUtil::kNoIdentifier;
0100       case "previous"_hash:
0101       case "next"_hash:
0102       case "measdim"_hash:
0103       case "referenceSurface"_hash:
0104       case "chi2"_hash:
0105       case "pathLength"_hash:
0106       case "typeFlags"_hash:
0107         return true;
0108       default:
0109         return instance.m_dynamic.contains(key);
0110     }
0111
0112     return false;
0113   }
0114
0115   template <bool EnsureConst, typename T>
0116   static std::any component_impl(T& instance, Acts::HashedString key,
0117                                  Acts::TrackIndexType istate) {
0118     if constexpr (EnsureConst) {
0119       static_assert(std::is_const_v<std::remove_reference_t<T>>,
0120                     "Is not const");
0121     }
0122     using namespace Acts::HashedStringLiteral;
0123     auto trackState = instance.m_collection->at(istate);
0124     std::conditional_t<EnsureConst, const ActsPodioEdm::TrackStateInfo*,
0125                        ActsPodioEdm::TrackStateInfo*>
0126         dataPtr;
0127     if constexpr (EnsureConst) {
0128       dataPtr = &trackState.getData();
0129     } else {
0130       dataPtr = &PodioUtil::getDataMutable(trackState);
0131     }
0132     auto& data = *dataPtr;
0133     switch (key) {
0134       case "previous"_hash:
0135         return &data.previous;
0136       case "next"_hash:
0137         return &data.next;
0138       case "predicted"_hash:
0139         return &data.ipredicted;
0140       case "filtered"_hash:
0141         return &data.ifiltered;
0142       case "smoothed"_hash:
0143         return &data.ismoothed;
0144       case "projector"_hash:
0145         return &data.projector;
0146       case "measdim"_hash:
0147         return &data.measdim;
0148       case "chi2"_hash:
0149         return &data.chi2;
0150       case "pathLength"_hash:
0151         return &data.pathLength;
0152       case "typeFlags"_hash:
0153         return &data.typeFlags;
0154       default:
0155         auto it = instance.m_dynamic.find(key);
0156         if (it == instance.m_dynamic.end()) {
0157           throw std::runtime_error("Unable to handle this component");
0158         }
0159         std::conditional_t<EnsureConst,
0160                            const podio_detail::ConstDynamicColumnBase*,
0161                            podio_detail::DynamicColumnBase*>
0162             col = it->second.get();
0163         assert(col && "Dynamic column is null");
0164         return col->get(istate);
0165     }
0166   }
0167
0168   template <typename T>
0169   static constexpr bool hasColumn_impl(T& instance, Acts::HashedString key) {
0170     using namespace Acts::HashedStringLiteral;
0171     switch (key) {
0172       case "predicted"_hash:
0173       case "filtered"_hash:
0174       case "smoothed"_hash:
0175       case "jacobian"_hash:
0176       case "projector"_hash:
0177       case "previous"_hash:
0178       case "next"_hash:
0179       case "uncalibratedSourceLink"_hash:
0180       case "referenceSurface"_hash:
0181       case "measdim"_hash:
0182       case "chi2"_hash:
0183       case "pathLength"_hash:
0184       case "typeFlags"_hash:
0185         return true;
0186       default:
0187         return instance.m_dynamic.contains(key);
0188     }
0189   }
0190
0191   static void populateSurfaceBuffer(
0192       const PodioUtil::ConversionHelper& helper,
0193       const ActsPodioEdm::TrackStateCollection& collection,
0194       std::vector<std::shared_ptr<const Acts::Surface>>& surfaces) noexcept {
0195     surfaces.reserve(collection.size());
0196     for (ActsPodioEdm::TrackState trackState : collection) {
0197       surfaces.push_back(PodioUtil::convertSurfaceFromPodio(
0198           helper, trackState.getReferenceSurface()));
0199     }
0200   }
0201 };
0202
0203 class ConstPodioTrackStateContainer final
0204     : public PodioTrackStateContainerBase,
0205       public Acts::MultiTrajectory<ConstPodioTrackStateContainer> {
0206  public:
0207   ConstPodioTrackStateContainer(
0208       const PodioUtil::ConversionHelper& helper,
0209       const ActsPodioEdm::TrackStateCollection& trackStates,
0210       const ActsPodioEdm::BoundParametersCollection& params,
0211       const ActsPodioEdm::JacobianCollection& jacs)
0212       : m_helper{helper},
0213         m_collection{&trackStates},
0214         m_params{&params},
0215         m_jacs{&jacs} {
0216     // Not much we can do to recover dynamic columns here
0217     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0218   }
0219
0220   /// Construct a const track state container from a mutable
0221   /// @warning If the source mutable container is modified, this container
0222   ///          will be corrupted, as surface buffer and dynamic column state can
0223   ///          not be synchronized!
0224   explicit ConstPodioTrackStateContainer(
0225       const MutablePodioTrackStateContainer& other);
0226
0227   ConstPodioTrackStateContainer(const PodioUtil::ConversionHelper& helper,
0228                                 const podio::Frame& frame,
0229                                 const std::string& suffix = "")
0230       : m_helper{helper},
0231         m_collection{nullptr},
0232         m_params{nullptr},
0233         m_jacs{nullptr} {
0234     std::string s = suffix.empty() ? suffix : "_" + suffix;
0235
0236     std::vector<std::string> available = frame.getAvailableCollections();
0237
0238     std::string trackStatesKey = "trackStates" + s;
0239     std::string paramsKey = "trackStateParameters" + s;
0240     std::string jacsKey = "trackStateJacobians" + s;
0241
0242     if (!Acts::rangeContainsValue(available, trackStatesKey)) {
0243       throw std::runtime_error{"Track state collection '" + trackStatesKey +
0244                                "' not found in frame"};
0245     }
0246
0247     if (!Acts::rangeContainsValue(available, paramsKey)) {
0248       throw std::runtime_error{"Track state parameters collection '" +
0249                                paramsKey + "' not found in frame"};
0250     }
0251
0252     if (!Acts::rangeContainsValue(available, jacsKey)) {
0253       throw std::runtime_error{"Track state jacobian collection '" + jacsKey +
0254                                "' not found in frame"};
0255     }
0256
0257     loadCollection<ActsPodioEdm::TrackStateCollection>(m_collection, frame,
0258                                                        trackStatesKey);
0259     loadCollection<ActsPodioEdm::BoundParametersCollection>(m_params, frame,
0260                                                             paramsKey);
0261     loadCollection<ActsPodioEdm::JacobianCollection>(m_jacs, frame, jacsKey);
0262
0263     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0264
0265     podio_detail::recoverDynamicColumns(frame, trackStatesKey, m_dynamic);
0266   }
0267
0268   Acts::detail::DynamicKeyRange<podio_detail::ConstDynamicColumnBase>
0269   dynamicKeys_impl() const {
0270     return {m_dynamic.begin(), m_dynamic.end()};
0271   }
0272
0273  private:
0274   template <typename collection_t>
0275   static void loadCollection(collection_t const*& dest,
0276                              const podio::Frame& frame,
0277                              const std::string& key) {
0278     const auto* collection = frame.get(key);
0279
0280     if (const auto* d = dynamic_cast<const collection_t*>(collection);
0281         d != nullptr) {
0282       dest = d;
0283     } else {
0284       throw std::runtime_error{"Unable to get collection " + key};
0285     }
0286   }
0287
0288  public:
0289   ConstParameters parameters_impl(IndexType istate) const {
0290     return ConstParameters{m_params->at(istate).getData().values.data()};
0291   }
0292
0293   ConstCovariance covariance_impl(IndexType istate) const {
0294     return ConstCovariance{m_params->at(istate).getData().covariance.data()};
0295   }
0296
0297   ConstCovariance jacobian_impl(IndexType istate) const {
0298     IndexType ijacobian = m_collection->at(istate).getData().ijacobian;
0299     return ConstCovariance{m_jacs->at(ijacobian).getData().values.data()};
0300   }
0301
0302   template <std::size_t measdim>
0303   ConstTrackStateProxy::Calibrated<measdim> calibrated_impl(
0304       IndexType index) const {
0305     return ConstTrackStateProxy::Calibrated<measdim>{
0306         m_collection->at(index).getData().measurement.data()};
0307   }
0308
0309   template <std::size_t measdim>
0310   ConstTrackStateProxy::CalibratedCovariance<measdim> calibratedCovariance_impl(
0311       IndexType index) const {
0312     return ConstTrackStateProxy::CalibratedCovariance<measdim>{
0313         m_collection->at(index).getData().measurementCovariance.data()};
0314   }
0315
0316   IndexType size_impl() const { return m_collection->size(); }
0317
0318   std::any component_impl(Acts::HashedString key, IndexType istate) const {
0319     return PodioTrackStateContainerBase::component_impl<true>(*this, key,
0320                                                               istate);
0321   }
0322
0323   constexpr bool hasColumn_impl(Acts::HashedString key) const {
0324     return PodioTrackStateContainerBase::hasColumn_impl(*this, key);
0325   }
0326
0327   constexpr bool has_impl(Acts::HashedString key, IndexType istate) const {
0328     return PodioTrackStateContainerBase::has_impl(*this, key, istate);
0329   }
0330
0331   Acts::MultiTrajectoryTraits::IndexType calibratedSize_impl(
0332       IndexType istate) const {
0333     return m_collection->at(istate).getData().measdim;
0334   }
0335
0336   Acts::SourceLink getUncalibratedSourceLink_impl(IndexType istate) const {
0337     return m_helper.get().identifierToSourceLink(
0338         m_collection->at(istate).getData().uncalibratedIdentifier);
0339   }
0340
0341   const Acts::Surface* referenceSurface_impl(IndexType istate) const {
0342     return m_surfaces.at(istate).get();
0343   }
0344
0345  private:
0346   friend class PodioTrackStateContainerBase;
0347
0348   std::reference_wrapper<const PodioUtil::ConversionHelper> m_helper;
0349   const ActsPodioEdm::TrackStateCollection* m_collection;
0350   const ActsPodioEdm::BoundParametersCollection* m_params;
0351   const ActsPodioEdm::JacobianCollection* m_jacs;
0352   std::vector<std::shared_ptr<const Acts::Surface>> m_surfaces;
0353
0354   std::unordered_map<Acts::HashedString,
0355                      std::unique_ptr<podio_detail::ConstDynamicColumnBase>>
0356       m_dynamic;
0357   std::vector<Acts::HashedString> m_dynamicKeys;
0358 };
0359
0360 static_assert(
0361     Acts::IsReadOnlyMultiTrajectory<ConstPodioTrackStateContainer>::value,
0362     "MutablePodioTrackStateContainer should not be read-only");
0363
0364 static_assert(
0365     Acts::ConstMultiTrajectoryBackend<ConstPodioTrackStateContainer>,
0366     "ConstPodioTrackStateContainer does not fulfill TrackContainerBackend");
0367
0368 class MutablePodioTrackStateContainer final
0369     : public PodioTrackStateContainerBase,
0370       public Acts::MultiTrajectory<MutablePodioTrackStateContainer> {
0371  public:
0372   explicit MutablePodioTrackStateContainer(PodioUtil::ConversionHelper& helper)
0373       : m_helper{helper} {
0374     m_collection = std::make_unique<ActsPodioEdm::TrackStateCollection>();
0375     m_jacs = std::make_unique<ActsPodioEdm::JacobianCollection>();
0376     m_params = std::make_unique<ActsPodioEdm::BoundParametersCollection>();
0377
0378     populateSurfaceBuffer(m_helper, *m_collection, m_surfaces);
0379   }
0380
0381   ConstParameters parameters_impl(IndexType istate) const {
0382     return ConstParameters{m_params->at(istate).getData().values.data()};
0383   }
0384
0385   Parameters parameters_impl(IndexType istate) {
0386     return Parameters{
0387         PodioUtil::getDataMutable(m_params->at(istate)).values.data()};
0388   }
0389
0390   ConstCovariance covariance_impl(IndexType istate) const {
0391     return ConstCovariance{m_params->at(istate).getData().covariance.data()};
0392   }
0393
0394   Covariance covariance_impl(IndexType istate) {
0395     return Covariance{
0396         PodioUtil::getDataMutable(m_params->at(istate)).covariance.data()};
0397   }
0398
0399   ConstCovariance jacobian_impl(IndexType istate) const {
0400     IndexType ijacobian = m_collection->at(istate).getData().ijacobian;
0401     return ConstCovariance{m_jacs->at(ijacobian).getData().values.data()};
0402   }
0403
0404   Covariance jacobian_impl(IndexType istate) {
0405     IndexType ijacobian = m_collection->at(istate).getData().ijacobian;
0406     return Covariance{
0407         PodioUtil::getDataMutable(m_jacs->at(ijacobian)).values.data()};
0408   }
0409
0410   template <std::size_t measdim>
0411   ConstTrackStateProxy::Calibrated<measdim> calibrated_impl(
0412       IndexType index) const {
0413     return ConstTrackStateProxy::Calibrated<measdim>{
0414         m_collection->at(index).getData().measurement.data()};
0415   }
0416
0417   template <std::size_t measdim>
0418   TrackStateProxy::Calibrated<measdim> calibrated_impl(IndexType index) {
0419     return TrackStateProxy::Calibrated<measdim>{
0420         PodioUtil::getDataMutable(m_collection->at(index)).measurement.data()};
0421   }
0422
0423   template <std::size_t measdim>
0424   ConstTrackStateProxy::CalibratedCovariance<measdim> calibratedCovariance_impl(
0425       IndexType index) const {
0426     return ConstTrackStateProxy::CalibratedCovariance<measdim>{
0427         m_collection->at(index).getData().measurementCovariance.data()};
0428   }
0429
0430   template <std::size_t measdim>
0431   TrackStateProxy::CalibratedCovariance<measdim> calibratedCovariance_impl(
0432       IndexType index) {
0433     return TrackStateProxy::CalibratedCovariance<measdim>{
0434         PodioUtil::getDataMutable(m_collection->at(index))
0435             .measurementCovariance.data()};
0436   }
0437
0438   IndexType size_impl() const { return m_collection->size(); }
0439
0440   std::any component_impl(Acts::HashedString key, IndexType istate) const {
0441     return PodioTrackStateContainerBase::component_impl<true>(*this, key,
0442                                                               istate);
0443   }
0444
0445   std::any component_impl(Acts::HashedString key, IndexType istate) {
0446     return PodioTrackStateContainerBase::component_impl<false>(*this, key,
0447                                                                istate);
0448   }
0449
0450   constexpr bool hasColumn_impl(Acts::HashedString key) const {
0451     return PodioTrackStateContainerBase::hasColumn_impl(*this, key);
0452   }
0453
0454   constexpr bool has_impl(Acts::HashedString key, IndexType istate) const {
0455     return PodioTrackStateContainerBase::has_impl(*this, key, istate);
0456   }
0457
0458   IndexType addTrackState_impl(
0459       Acts::TrackStatePropMask mask = Acts::TrackStatePropMask::All,
0460       Acts::TrackIndexType iprevious = Acts::kTrackIndexInvalid) {
0461     auto trackState = m_collection->create();
0462     auto& data = PodioUtil::getDataMutable(trackState);
0463     data.previous = iprevious;
0464     data.ipredicted = kInvalid;
0465     data.ifiltered = kInvalid;
0466     data.ismoothed = kInvalid;
0467     data.ijacobian = kInvalid;
0468
0469     PodioUtil::getReferenceSurfaceMutable(trackState).surfaceType =
0470         PodioUtil::kNoSurface;
0471
0472     using enum Acts::TrackStatePropMask;
0473
0474     if (ACTS_CHECK_BIT(mask, Predicted)) {
0475       m_params->create();
0476       data.ipredicted = m_params->size() - 1;
0477     }
0478     if (ACTS_CHECK_BIT(mask, Filtered)) {
0479       m_params->create();
0480       data.ifiltered = m_params->size() - 1;
0481     }
0482     if (ACTS_CHECK_BIT(mask, Smoothed)) {
0483       m_params->create();
0484       data.ismoothed = m_params->size() - 1;
0485     }
0486     if (ACTS_CHECK_BIT(mask, Jacobian)) {
0487       m_jacs->create();
0488       data.ijacobian = m_jacs->size() - 1;
0489     }
0490     data.measdim = kInvalid;
0491     data.hasProjector = false;
0492     if (ACTS_CHECK_BIT(mask, Calibrated)) {
0493       data.hasProjector = true;
0494     }
0495     m_surfaces.emplace_back();
0496
0497     data.uncalibratedIdentifier = PodioUtil::kNoIdentifier;
0498     assert(m_collection->size() == m_surfaces.size() &&
0499            "Inconsistent surface buffer");
0500
0501     for (const auto& [key, vec] : m_dynamic) {
0502       vec->add();
0503     }
0504
0505     return m_collection->size() - 1;
0506   }
0507
0508   void addTrackStateComponents_impl(IndexType istate,
0509                                     Acts::TrackStatePropMask mask) {
0510     auto& data = PodioUtil::getDataMutable(m_collection->at(istate));
0511
0512     using enum Acts::TrackStatePropMask;
0513
0514     if (ACTS_CHECK_BIT(mask, Predicted) && data.ipredicted == kInvalid) {
0515       m_params->create();
0516       data.ipredicted = m_params->size() - 1;
0517     }
0518
0519     if (ACTS_CHECK_BIT(mask, Filtered) && data.ifiltered == kInvalid) {
0520       m_params->create();
0521       data.ifiltered = m_params->size() - 1;
0522     }
0523
0524     if (ACTS_CHECK_BIT(mask, Smoothed) && data.ismoothed == kInvalid) {
0525       m_params->create();
0526       data.ismoothed = m_params->size() - 1;
0527     }
0528
0529     if (ACTS_CHECK_BIT(mask, Jacobian) && data.ijacobian == kInvalid) {
0530       m_jacs->create();
0531       data.ijacobian = m_jacs->size() - 1;
0532     }
0533
0534     if (ACTS_CHECK_BIT(mask, Calibrated) && !data.hasProjector) {
0535       data.hasProjector = true;
0536     }
0537   }
0538
0539   void shareFrom_impl(Acts::TrackIndexType iself, Acts::TrackIndexType iother,
0540                       Acts::TrackStatePropMask shareSource,
0541                       Acts::TrackStatePropMask shareTarget) {
0542     auto& self = PodioUtil::getDataMutable(m_collection->at(iself));
0543     auto& other = PodioUtil::getDataMutable(m_collection->at(iother));
0544
0545     assert(ACTS_CHECK_BIT(getTrackState(iother).getMask(), shareSource) &&
0546            "Source has incompatible allocation");
0547
0548     using enum Acts::TrackStatePropMask;
0549
0550     IndexType sourceIndex{kInvalid};
0551     switch (shareSource) {
0552       case Predicted:
0553         sourceIndex = other.ipredicted;
0554         break;
0555       case Filtered:
0556         sourceIndex = other.ifiltered;
0557         break;
0558       case Smoothed:
0559         sourceIndex = other.ismoothed;
0560         break;
0561       case Jacobian:
0562         sourceIndex = other.ijacobian;
0563         break;
0564       default:
0565         throw std::domain_error{"Unable to share this component"};
0566     }
0567
0568     assert(sourceIndex != kInvalid);
0569
0570     switch (shareTarget) {
0571       case Predicted:
0572         assert(shareSource != Jacobian);
0573         self.ipredicted = sourceIndex;
0574         break;
0575       case Filtered:
0576         assert(shareSource != Jacobian);
0577         self.ifiltered = sourceIndex;
0578         break;
0579       case Smoothed:
0580         assert(shareSource != Jacobian);
0581         self.ismoothed = sourceIndex;
0582         break;
0583       case Jacobian:
0584         assert(shareSource == Jacobian);
0585         self.ijacobian = sourceIndex;
0586         break;
0587       default:
0588         throw std::domain_error{"Unable to share this component"};
0589     }
0590   }
0591
0592   void unset_impl(Acts::TrackStatePropMask target,
0593                   Acts::TrackIndexType istate) {
0594     auto& data = PodioUtil::getDataMutable(m_collection->at(istate));
0595
0596     switch (target) {
0597       using enum Acts::TrackStatePropMask;
0598
0599       case Predicted:
0600         data.ipredicted = kInvalid;
0601         break;
0602       case Filtered:
0603         data.ifiltered = kInvalid;
0604         break;
0605       case Smoothed:
0606         data.ismoothed = kInvalid;
0607         break;
0608       case Jacobian:
0609         data.ijacobian = kInvalid;
0610         break;
0611       case Calibrated:
0612         data.measdim = kInvalid;
0613         break;
0614       default:
0615         throw std::domain_error{"Unable to unset this component"};
0616     }
0617   }
0618
0619   void clear_impl() {
0620     m_collection->clear();
0621     m_params->clear();
0622     m_surfaces.clear();
0623     for (const auto& [key, vec] : m_dynamic) {
0624       vec->clear();
0625     }
0626   }
0627
0628   template <typename T>
0629   constexpr void addColumn_impl(std::string_view key) {
0630     Acts::HashedString hashedKey = Acts::hashStringDynamic(key);
0631     m_dynamic.insert(
0632         {hashedKey, std::make_unique<podio_detail::DynamicColumn<T>>(key)});
0633   }
0634
0635   template <typename val_t, typename cov_t>
0636   void allocateCalibrated_impl(IndexType istate,
0637                                const Eigen::DenseBase<val_t>& val,
0638                                const Eigen::DenseBase<cov_t>& cov)
0639     requires(Acts::Concepts::eigen_base_is_fixed_size<val_t> &&
0640              Eigen::PlainObjectBase<val_t>::RowsAtCompileTime <=
0641                  toUnderlying(Acts::eBoundSize) &&
0642              Acts::Concepts::eigen_bases_have_same_num_rows<val_t, cov_t> &&
0643              Acts::Concepts::eigen_base_is_square<cov_t>)
0644   {
0645     constexpr std::size_t measdim = val_t::RowsAtCompileTime;
0646
0647     auto& data = PodioUtil::getDataMutable(m_collection->at(istate));
0648
0649     if (data.measdim != kInvalid && data.measdim != measdim) {
0650       throw std::invalid_argument{
0651           "Measurement dimension does not match the allocated dimension"};
0652     }
0653
0654     data.measdim = measdim;
0655
0656     Eigen::Map<Acts::ActsVector<measdim>> valMap(data.measurement.data());
0657     valMap = val;
0658
0659     Eigen::Map<Acts::ActsSquareMatrix<measdim>> covMap(
0660         data.measurementCovariance.data());
0661     covMap = cov;
0662   }
0663
0664   void setUncalibratedSourceLink_impl(IndexType istate,
0665                                       const Acts::SourceLink& sourceLink) {
0666     PodioUtil::Identifier id =
0667         m_helper.get().sourceLinkToIdentifier(sourceLink);
0668     auto& data = PodioUtil::getDataMutable(m_collection->at(istate));
0669     data.uncalibratedIdentifier = id;
0670   }
0671
0672   void setReferenceSurface_impl(IndexType istate,
0673                                 std::shared_ptr<const Acts::Surface> surface) {
0674     auto trackState = m_collection->at(istate);
0675     trackState.setReferenceSurface(
0676         PodioUtil::convertSurfaceToPodio(m_helper, *surface));
0677     m_surfaces.at(istate) = std::move(surface);
0678   }
0679
0680   Acts::MultiTrajectoryTraits::IndexType calibratedSize_impl(
0681       IndexType istate) const {
0682     return m_collection->at(istate).getData().measdim;
0683   }
0684
0685   Acts::SourceLink getUncalibratedSourceLink_impl(IndexType istate) const {
0686     return m_helper.get().identifierToSourceLink(
0687         m_collection->at(istate).getData().uncalibratedIdentifier);
0688   }
0689
0690   const Acts::Surface* referenceSurface_impl(IndexType istate) const {
0691     return m_surfaces.at(istate).get();
0692   }
0693
0694   void releaseInto(podio::Frame& frame, const std::string& suffix = "") {
0695     std::string s = suffix;
0696     if (!s.empty()) {
0697       s = "_" + s;
0698     }
0699     frame.put(std::move(m_collection), "trackStates" + s);
0700     frame.put(std::move(m_params), "trackStateParameters" + s);
0701     frame.put(std::move(m_jacs), "trackStateJacobians" + s);
0702     m_surfaces.clear();
0703
0704     for (const auto& [key, col] : m_dynamic) {
0705       col->releaseInto(frame, "trackStates" + s + "_extra__");
0706     }
0707
0708     m_dynamic.clear();
0709   }
0710
0711   Acts::detail::DynamicKeyRange<podio_detail::DynamicColumnBase>
0712   dynamicKeys_impl() const {
0713     return {m_dynamic.begin(), m_dynamic.end()};
0714   }
0715
0716   void copyDynamicFrom_impl(IndexType dstIdx, Acts::HashedString key,
0717                             const std::any& srcPtr) {
0718     auto it = m_dynamic.find(key);
0719     if (it == m_dynamic.end()) {
0720       throw std::invalid_argument{
0721           "Destination container does not have matching dynamic column"};
0722     }
0723
0724     it->second->copyFrom(dstIdx, srcPtr);
0725   }
0726
0727  private:
0728   friend class PodioTrackStateContainerBase;
0729   friend class ConstPodioTrackStateContainer;
0730
0731   std::reference_wrapper<PodioUtil::ConversionHelper> m_helper;
0732   std::unique_ptr<ActsPodioEdm::TrackStateCollection> m_collection;
0733   std::unique_ptr<ActsPodioEdm::BoundParametersCollection> m_params;
0734   std::unique_ptr<ActsPodioEdm::JacobianCollection> m_jacs;
0735   std::vector<std::shared_ptr<const Acts::Surface>> m_surfaces;
0736
0737   std::unordered_map<Acts::HashedString,
0738                      std::unique_ptr<podio_detail::DynamicColumnBase>>
0739       m_dynamic;
0740   std::vector<Acts::HashedString> m_dynamicKeys;
0741 };
0742
0743 static_assert(
0744     !Acts::IsReadOnlyMultiTrajectory<MutablePodioTrackStateContainer>::value,
0745     "MutablePodioTrackStateContainer should not be read-only");
0746
0747 static_assert(
0748     Acts::MutableMultiTrajectoryBackend<MutablePodioTrackStateContainer>,
0749     "MutablePodioTrackStateContainer does not fulfill "
0750     "TrackStateContainerBackend");
0751
0752 inline ConstPodioTrackStateContainer::ConstPodioTrackStateContainer(
0753     const MutablePodioTrackStateContainer& other)
0754     : m_helper{other.m_helper},
0755       m_collection{other.m_collection.get()},
0756       m_params{other.m_params.get()},
0757       m_jacs{other.m_jacs.get()},
0758       m_surfaces{other.m_surfaces} {
0759   for (const auto& [key, col] : other.m_dynamic) {
0760     m_dynamic.insert({key, col->asConst()});
0761   }
0762 }
0763
0764 }  // namespace ActsPlugins