ActsPlugins/EDM4hep/EDM4hepUtil.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 #pragma once
0009
0010 #include "Acts/Definitions/Algebra.hpp"
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/EventData/Charge.hpp"
0013 #include "Acts/EventData/GenericBoundTrackParameters.hpp"
0014 #include "Acts/EventData/MultiTrajectory.hpp"
0015 #include "Acts/EventData/TrackContainer.hpp"
0016 #include "Acts/EventData/TrackParameters.hpp"
0017 #include "Acts/EventData/TrackProxyConcept.hpp"
0018 #include "Acts/EventData/TrackStatePropMask.hpp"
0019 #include "Acts/Geometry/GeometryContext.hpp"
0020 #include "Acts/Surfaces/PerigeeSurface.hpp"
0021 #include "Acts/Surfaces/Surface.hpp"
0022 #include "Acts/Utilities/Logger.hpp"
0023 #include "Acts/Utilities/UnitVectors.hpp"
0024
0025 #include <stdexcept>
0026
0027 #include <Eigen/src/Core/util/Memory.h>
0028 #include <boost/graph/graph_traits.hpp>
0029 #include <edm4hep/MCParticle.h>
0030 #include <edm4hep/MutableSimTrackerHit.h>
0031 #include <edm4hep/MutableTrack.h>
0032 #include <edm4hep/SimTrackerHit.h>
0033 #include <edm4hep/Track.h>
0034 #include <edm4hep/TrackState.h>
0035 #include <podio/podioVersion.h>
0036
0037 #if podio_VERSION_MAJOR == 0 || \
0038     (podio_VERSION_MAJOR == 1 && podio_VERSION_MINOR <= 2)
0039
0040 template <>
0041 struct std::hash<podio::ObjectID> {
0042   std::size_t operator()(const podio::ObjectID& id) const noexcept {
0043     auto hash_collectionID = std::hash<std::uint32_t>{}(id.collectionID);
0044     auto hash_index = std::hash<int>{}(id.index);
0045
0046     return hash_collectionID ^ hash_index;
0047   }
0048 };
0049
0050 #endif
0051
0052 namespace ActsPlugins::EDM4hepUtil {
0053
0054 static constexpr std::int32_t EDM4HEP_ACTS_POSITION_TYPE = 42;
0055
0056 namespace detail {
0057 struct Parameters {
0058   Acts::ActsVector<6> values{};
0059   // Dummy default
0060   Acts::ParticleHypothesis particleHypothesis =
0061       Acts::ParticleHypothesis::pion();
0062   std::optional<Acts::BoundSquareMatrix> covariance;
0063   std::shared_ptr<const Acts::Surface> surface;
0064 };
0065
0066 Acts::ActsSquareMatrix<6> jacobianToEdm4hep(double theta, double qOverP,
0067                                             double Bz);
0068
0069 Acts::ActsSquareMatrix<6> jacobianFromEdm4hep(double tanLambda, double omega,
0070                                               double Bz);
0071
0072 void unpackCovariance(const float* from, Acts::ActsSquareMatrix<6>& to);
0073 void packCovariance(const Acts::ActsSquareMatrix<6>& from, float* to);
0074
0075 Parameters convertTrackParametersToEdm4hep(
0076     const Acts::GeometryContext& gctx, double Bz,
0077     const Acts::BoundTrackParameters& params);
0078
0079 Acts::BoundTrackParameters convertTrackParametersFromEdm4hep(
0080     double Bz, const Parameters& params);
0081
0082 }  // namespace detail
0083
0084 // Compatibility with EDM4hep < 0.99 and >= 0.99
0085 edm4hep::MCParticle getParticle(const edm4hep::SimTrackerHit& hit);
0086
0087 void setParticle(edm4hep::MutableSimTrackerHit& hit,
0088                  const edm4hep::MCParticle& particle);
0089
0090 template <Acts::TrackProxyConcept track_proxy_t>
0091 void writeTrack(const Acts::GeometryContext& gctx, track_proxy_t track,
0092                 edm4hep::MutableTrack to, double Bz,
0093                 const Acts::Logger& logger = Acts::getDummyLogger()) {
0094   ACTS_VERBOSE("Converting track to EDM4hep");
0095   to.setChi2(track.chi2());
0096   to.setNdf(track.nDoF());
0097
0098   std::vector<edm4hep::TrackState> outTrackStates;
0099   outTrackStates.reserve(track.nTrackStates());
0100
0101   auto setParameters = [](edm4hep::TrackState& trackState,
0102                           const detail::Parameters& params) {
0103     trackState.D0 = params.values[0];
0104     trackState.Z0 = params.values[1];
0105     trackState.phi = params.values[2];
0106     trackState.tanLambda = params.values[3];
0107     trackState.omega = params.values[4];
0108     trackState.time = params.values[5];
0109
0110     if (params.covariance) {
0111       detail::packCovariance(params.covariance.value(),
0112                              trackState.covMatrix.data());
0113     }
0114   };
0115
0116   ACTS_VERBOSE("Converting " << track.nTrackStates() << " track states");
0117
0118   for (const auto& state : track.trackStatesReversed()) {
0119     auto typeFlags = state.typeFlags();
0120     if (!typeFlags.test(Acts::TrackStateFlag::MeasurementFlag)) {
0121       continue;
0122     }
0123
0124     edm4hep::TrackState& trackState = outTrackStates.emplace_back();
0125     trackState.location = edm4hep::TrackState::AtOther;
0126
0127     Acts::BoundTrackParameters params{state.referenceSurface().getSharedPtr(),
0128                                       state.parameters(), state.covariance(),
0129                                       track.particleHypothesis()};
0130
0131     // Convert to LCIO track parametrization expected by EDM4hep
0132     detail::Parameters converted =
0133         detail::convertTrackParametersToEdm4hep(gctx, Bz, params);
0134
0135     // Write the converted parameters to the EDM4hep track state
0136     setParameters(trackState, converted);
0137     ACTS_VERBOSE("- parameters: " << state.parameters().transpose() << " -> "
0138                                   << converted.values.transpose());
0139     ACTS_VERBOSE("- covariance: \n"
0140                  << state.covariance() << "\n->\n"
0141                  << converted.covariance.value());
0142
0143     // Converted parameters are relative to an ad-hoc perigee surface created at
0144     // the hit location
0145     auto center = converted.surface->center(gctx);
0146     trackState.referencePoint.x = center.x();
0147     trackState.referencePoint.y = center.y();
0148     trackState.referencePoint.z = center.z();
0149     ACTS_VERBOSE("- ref surface ctr: " << center.transpose());
0150   }
0151   outTrackStates.front().location = edm4hep::TrackState::AtLastHit;
0152   outTrackStates.back().location = edm4hep::TrackState::AtFirstHit;
0153
0154   // Add a track state that represents the IP parameters
0155   auto& ipState = outTrackStates.emplace_back();
0156
0157   // Convert the track parameters at the IP
0158   Acts::BoundTrackParameters trackParams{
0159       track.referenceSurface().getSharedPtr(), track.parameters(),
0160       track.covariance(), track.particleHypothesis()};
0161
0162   // Convert to LCIO track parametrization expected by EDM4hep
0163   auto converted =
0164       detail::convertTrackParametersToEdm4hep(gctx, Bz, trackParams);
0165   setParameters(ipState, converted);
0166   ipState.location = edm4hep::TrackState::AtIP;
0167   ACTS_VERBOSE("Writing track level quantities as IP track state");
0168   ACTS_VERBOSE("- parameters: " << track.parameters().transpose());
0169   ACTS_VERBOSE("           -> " << converted.values.transpose());
0170   ACTS_VERBOSE("- covariance: \n"
0171                << track.covariance() << "\n->\n"
0172                << converted.covariance.value());
0173
0174   // Write the converted parameters to the EDM4hep track state
0175   // The reference point is at the location of the reference surface of the
0176   // track itself, but if that's not a perigee surface, another ad-hoc perigee
0177   // at the position will be created.
0178   auto center = converted.surface->center(gctx);
0179   ipState.referencePoint.x = center.x();
0180   ipState.referencePoint.y = center.y();
0181   ipState.referencePoint.z = center.z();
0182
0183   ACTS_VERBOSE("- ref surface ctr: " << center.transpose());
0184
0185   for (auto& trackState : outTrackStates) {
0186     to.addToTrackStates(trackState);
0187   }
0188 }
0189
0190 template <Acts::TrackProxyConcept track_proxy_t>
0191 void readTrack(const edm4hep::Track& from, track_proxy_t& track, double Bz,
0192                const Acts::Logger& logger = Acts::getDummyLogger()) {
0193   using namespace Acts;
0194   ACTS_VERBOSE("Reading track from EDM4hep");
0195   TrackStatePropMask mask = TrackStatePropMask::Smoothed;
0196
0197   std::optional<edm4hep::TrackState> ipState;
0198
0199   auto unpack =
0200       [](const edm4hep::TrackState& trackState) -> detail::Parameters {
0201     detail::Parameters params;
0202     params.covariance = BoundMatrix::Zero();
0203     params.values = BoundVector::Zero();
0204     detail::unpackCovariance(trackState.covMatrix.data(),
0205                              params.covariance.value());
0206     params.values[0] = trackState.D0;
0207     params.values[1] = trackState.Z0;
0208     params.values[2] = trackState.phi;
0209     params.values[3] = trackState.tanLambda;
0210     params.values[4] = trackState.omega;
0211     params.values[5] = trackState.time;
0212
0213     Vector3 center = {
0214         trackState.referencePoint.x,
0215         trackState.referencePoint.y,
0216         trackState.referencePoint.z,
0217     };
0218     params.surface = Acts::Surface::makeShared<PerigeeSurface>(center);
0219
0220     return params;
0221   };
0222
0223   ACTS_VERBOSE("Reading " << from.trackStates_size()
0224                           << " track states (including IP state)");
0225   // We write the trackstates out outside in, need to reverse iterate to get the
0226   // same order
0227   for (std::size_t i = from.trackStates_size() - 1;
0228        i <= from.trackStates_size(); i--) {
0229     auto trackState = from.getTrackStates(i);
0230     if (trackState.location == edm4hep::TrackState::AtIP) {
0231       ipState = trackState;
0232       continue;
0233     }
0234
0235     auto params = unpack(trackState);
0236
0237     auto ts = track.appendTrackState(mask);
0238     ts.typeFlags().set(MeasurementFlag);
0239
0240     auto converted = detail::convertTrackParametersFromEdm4hep(Bz, params);
0241
0242     ts.smoothed() = converted.parameters();
0243     ts.smoothedCovariance() =
0244         converted.covariance().value_or(BoundMatrix::Zero());
0245     ts.setReferenceSurface(params.surface);
0246   }
0247
0248   if (!ipState.has_value()) {
0249     ACTS_ERROR("Did not find IP state in edm4hep input");
0250     throw std::runtime_error{"Did not find IP state in edm4hep input"};
0251   }
0252
0253   detail::Parameters params = unpack(ipState.value());
0254
0255   auto converted = detail::convertTrackParametersFromEdm4hep(Bz, params);
0256
0257   ACTS_VERBOSE("IP state parameters: " << converted.parameters().transpose());
0258   ACTS_VERBOSE("-> covariance:\n"
0259                << converted.covariance().value_or(BoundMatrix::Zero()));
0260
0261   track.parameters() = converted.parameters();
0262   track.covariance() = converted.covariance().value_or(BoundMatrix::Zero());
0263   track.setReferenceSurface(params.surface);
0264
0265   track.chi2() = from.getChi2();
0266   track.nDoF() = from.getNdf();
0267   track.nMeasurements() = track.nTrackStates();
0268 }
0269
0270 class SimHitAssociation {
0271  public:
0272   void reserve(std::size_t size);
0273
0274   std::size_t size() const;
0275
0276   void add(std::size_t internalIndex, const edm4hep::SimTrackerHit& edm4hepHit);
0277
0278   [[nodiscard]]
0279   edm4hep::SimTrackerHit lookup(std::size_t internalIndex) const;
0280
0281   std::size_t lookup(const edm4hep::SimTrackerHit& hit) const;
0282
0283  private:
0284   std::vector<edm4hep::SimTrackerHit> m_internalToEdm4hep;
0285   std::unordered_map<podio::ObjectID, std::size_t> m_edm4hepToInternal;
0286 };
0287
0288 }  // namespace ActsPlugins::EDM4hepUtil