EDM4hep/src/EDM4hepUtil.cpp

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 #include "ActsPlugins/EDM4hep/EDM4hepUtil.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Definitions/TrackParametrization.hpp"
0013 #include "Acts/Definitions/Units.hpp"
0014 #include "Acts/EventData/MultiTrajectory.hpp"
0015 #include "Acts/Propagator/detail/CovarianceEngine.hpp"
0016 #include "Acts/Utilities/UnitVectors.hpp"
0017 #include "Acts/Vertexing/Vertex.hpp"
0018 #include "ActsPodioEdm/TrackerHitLocalCollection.h"
0019 #include "ActsPodioEdm/TrackerHitLocalSimTrackerHitLinkCollection.h"
0020
0021 #include <array>
0022 #include <numbers>
0023
0024 #include <edm4hep/Constants.h>
0025 #include <edm4hep/MCParticle.h>
0026 #include <edm4hep/MutableSimTrackerHit.h>
0027 #include <edm4hep/MutableVertex.h>
0028 #include <edm4hep/SimTrackerHit.h>
0029 #include <edm4hep/TrackState.h>
0030 #include <edm4hep/Vector3f.h>
0031 #include <edm4hep/Vector4f.h>
0032
0033 using namespace Acts;
0034 using namespace Acts::detail;
0035
0036 namespace ActsPlugins::EDM4hepUtil {
0037 namespace detail {
0038
0039 SquareMatrix<6> jacobianToEdm4hep(double theta, double qOverP, double Bz) {
0040   // Calculate jacobian from our internal parametrization (d0, z0, phi, theta,
0041   // q/p) to the LCIO / edm4hep (see:
0042   // https://bib-pubdb1.desy.de/record/81214/files/LC-DET-2006-004%5B1%5D.pdf)
0043   // one (d0, z0, phi, tan(lambda), omega). Top left 3x3 matrix in the
0044   // jacobian is 1. Bottom right 2x2 matrix is:
0045   //
0046   // [  d                                 ]
0047   // [------(cot(theta))         0        ]
0048   // [dtheta                              ]
0049   // [                                    ]
0050   // [  d   /  B*q/p   \   d  /  B*q/p   \]
0051   // [------|----------|  ----|----------|]
0052   // [dtheta\sin(theta)/  dq/p\sin(theta)/]
0053   //
0054   // =
0055   //
0056   // [     2                        ]
0057   // [- csc (theta)           0     ]
0058   // [                              ]
0059   // [-B*q/p*cos(theta)       B     ]
0060   // [------------------  ----------]
0061   // [      2             sin(theta)]
0062   // [   sin (theta)                ]
0063
0064   SquareMatrix<6> J;
0065   J.setIdentity();
0066   double sinTheta = std::sin(theta);
0067   J(3, 3) = -1.0 / (sinTheta * sinTheta);
0068   J(4, 4) = Bz / sinTheta;  // dOmega / d(qop)
0069   J(4, 3) = -Bz * qOverP * std::cos(theta) /
0070             (sinTheta * sinTheta);  // dOmega / dTheta
0071   return J;
0072 }
0073
0074 SquareMatrix<6> jacobianFromEdm4hep(double tanLambda, double omega, double Bz) {
0075   // [     d      /                     pi\                                  ]
0076   // [------------|-atan(\tan\lambda) + --|                 0                ]
0077   // [d\tan\lambda\                     2 /                                  ]
0078   // [                                                                       ]
0079   // [     d      /         \Omega        \     d   /         \Omega        \]
0080   // [------------|-----------------------|  -------|-----------------------|]
0081   // [d\tan\lambda|     __________________|  d\Omega|     __________________|]
0082   // [            |    /            2     |         |    /            2     |]
0083   // [            \B*\/  \tan\lambda  + 1 /         \B*\/  \tan\lambda  + 1 /]
0084   //
0085   // =
0086   //
0087   // [         -1                                     ]
0088   // [   ----------------                 0           ]
0089   // [              2                                 ]
0090   // [   \tan\lambda  + 1                             ]
0091   // [                                                ]
0092   // [  -\Omega*\tan\lambda               1           ]
0093   // [-----------------------  -----------------------]
0094   // [                    3/2       __________________]
0095   // [  /           2    \         /            2     ]
0096   // [B*\\tan\lambda  + 1/     B*\/  \tan\lambda  + 1 ]
0097
0098   SquareMatrix<6> J;
0099   J.setIdentity();
0100   J(3, 3) = -1 / (tanLambda * tanLambda + 1);
0101   J(4, 3) = -1 * omega * tanLambda /
0102             (Bz * std::pow(tanLambda * tanLambda + 1, 3. / 2.));
0103   J(4, 4) = 1 / (Bz * std::hypot(tanLambda, 1));
0104
0105   return J;
0106 }
0107
0108 namespace {
0109
0110 /// Mapping from the parameter order used by @c Parameters::values
0111 /// (d0, z0, phi, tanLambda, omega, time) onto the EDM4hep track parameter
0112 /// order, which is a different permutation
0113 /// (@c edm4hep::TrackParams : d0, phi, omega, z0, tanLambda, time).
0114 ///
0115 /// The two orders are not the same, so the covariance has to be permuted when
0116 /// it is packed into (or read back from) an @c edm4hep::TrackState, otherwise
0117 /// the stored matrix is scrambled with respect to what any other EDM4hep
0118 /// consumer expects.
0119 constexpr std::array<edm4hep::TrackParams, 6> kValuesToEdm4hepParam{
0120     edm4hep::TrackParams::d0,         // values[0], eBoundLoc0
0121     edm4hep::TrackParams::z0,         // values[1], eBoundLoc1
0122     edm4hep::TrackParams::phi,        // values[2]
0123     edm4hep::TrackParams::tanLambda,  // values[3]
0124     edm4hep::TrackParams::omega,      // values[4]
0125     edm4hep::TrackParams::time,       // values[5]
0126 };
0127
0128 }  // namespace
0129
0130 void packCovariance(const SquareMatrix<6>& from, edm4hep::CovMatrix6f& to) {
0131   // Let EDM4hep place the elements: setValue takes the parameters by name, so
0132   // the packed layout stays whatever EDM4hep says it is.
0133   for (std::size_t i = 0; i < 6; ++i) {
0134     for (std::size_t j = 0; j <= i; ++j) {
0135       to.setValue(static_cast<float>(from(i, j)), kValuesToEdm4hepParam.at(i),
0136                   kValuesToEdm4hepParam.at(j));
0137     }
0138   }
0139 }
0140
0141 void unpackCovariance(const edm4hep::CovMatrix6f& from, SquareMatrix<6>& to) {
0142   for (std::size_t i = 0; i < 6; ++i) {
0143     for (std::size_t j = 0; j < 6; ++j) {
0144       to(i, j) = from.getValue(kValuesToEdm4hepParam.at(i),
0145                                kValuesToEdm4hepParam.at(j));
0146     }
0147   }
0148 }
0149
0150 double localBz(const MagneticFieldProvider& magneticField,
0151                MagneticFieldProvider::Cache& fieldCache,
0152                const Vector3& position, const Logger& logger) {
0153   auto field = magneticField.getField(position, fieldCache);
0154   if (!field.ok()) {
0155     ACTS_ERROR("Magnetic field lookup failed at "
0156                << position.transpose() << ": " << field.error().message());
0157     throw std::runtime_error{"Magnetic field lookup failed: " +
0158                              field.error().message()};
0159   }
0160   return (*field).z();
0161 }
0162
0163 Vector3 referencePoint(const edm4hep::TrackState& trackState) {
0164   return Vector3{
0165       trackState.referencePoint.x,
0166       trackState.referencePoint.y,
0167       trackState.referencePoint.z,
0168   };
0169 }
0170
0171 Parameters unpackTrackState(const edm4hep::TrackState& trackState) {
0172   Parameters params;
0173   params.covariance = BoundMatrix::Zero();
0174   params.values = BoundVector::Zero();
0175   unpackCovariance(trackState.covMatrix, params.covariance.value());
0176   params.values[0] = trackState.D0;
0177   params.values[1] = trackState.Z0;
0178   params.values[2] = trackState.phi;
0179   params.values[3] = trackState.tanLambda;
0180   params.values[4] = trackState.omega;
0181   params.values[5] = trackState.time;
0182
0183   params.surface =
0184       Surface::makeShared<PerigeeSurface>(referencePoint(trackState));
0185
0186   return params;
0187 }
0188
0189 Vector3 pointOfClosestApproach(const GeometryContext& gctx,
0190                                const Parameters& params) {
0191   // theta and phi are available without knowing the field, so this does not
0192   // depend on the q/p conversion that needs Bz.
0193   const double theta = std::numbers::pi / 2. - std::atan(params.values[3]);
0194   const Vector3 direction = makeDirectionFromPhiTheta(params.values[2], theta);
0195   return params.surface->localToGlobal(
0196       gctx, Vector2{params.values[0], params.values[1]}, direction);
0197 }
0198
0199 Parameters convertTrackParametersToEdm4hep(const GeometryContext& gctx,
0200                                            double Bz,
0201                                            const BoundTrackParameters& params) {
0202   Vector3 global = params.referenceSurface().localToGlobal(
0203       gctx, params.parameters().template head<2>(), params.direction());
0204
0205   std::shared_ptr<const Surface> refSurface =
0206       params.referenceSurface().getSharedPtr();
0207   BoundVector targetPars = params.parameters();
0208   std::optional<BoundMatrix> targetCov = params.covariance();
0209
0210   // If the reference surface is a perigee surface, we use that. Otherwise
0211   // we create a new perigee surface at the global position of the track
0212   // parameters.
0213   if (dynamic_cast<const PerigeeSurface*>(refSurface.get()) == nullptr) {
0214     refSurface = Surface::makeShared<PerigeeSurface>(global);
0215
0216     // We need to convert to the target parameters
0217     // Keep the free parameters around we might need them for the covariance
0218     // conversion
0219
0220     auto perigeeParams =
0221         boundToBoundConversion(gctx, params, *refSurface, Vector3{0, 0, Bz})
0222             .value();
0223     targetPars = perigeeParams.parameters();
0224     targetCov = perigeeParams.covariance();
0225   }
0226
0227   Parameters result;
0228   result.surface = refSurface;
0229
0230   // Only run covariance conversion if we have a covariance input
0231   if (targetCov) {
0232     SquareMatrix<6> J = jacobianToEdm4hep(targetPars[eBoundTheta],
0233                                           targetPars[eBoundQOverP], Bz);
0234     result.covariance = J * targetCov.value() * J.transpose();
0235   }
0236
0237   result.values[0] = targetPars[eBoundLoc0];
0238   result.values[1] = targetPars[eBoundLoc1];
0239   result.values[2] = targetPars[eBoundPhi];
0240   result.values[3] = std::tan(std::numbers::pi / 2. - targetPars[eBoundTheta]);
0241   result.values[4] =
0242       targetPars[eBoundQOverP] / std::sin(targetPars[eBoundTheta]) * Bz;
0243   result.values[5] = targetPars[eBoundTime];
0244
0245   result.particleHypothesis = params.particleHypothesis();
0246
0247   return result;
0248 }
0249
0250 BoundTrackParameters convertTrackParametersFromEdm4hep(
0251     double Bz, const Parameters& params) {
0252   BoundVector targetPars;
0253
0254   SquareMatrix<6> J =
0255       jacobianFromEdm4hep(params.values[3], params.values[4], Bz);
0256
0257   std::optional<BoundMatrix> cov;
0258   if (params.covariance.has_value()) {
0259     cov = J * params.covariance.value() * J.transpose();
0260   }
0261
0262   targetPars[eBoundLoc0] = params.values[0];
0263   targetPars[eBoundLoc1] = params.values[1];
0264   targetPars[eBoundPhi] = params.values[2];
0265   targetPars[eBoundTheta] = std::numbers::pi / 2. - std::atan(params.values[3]);
0266   targetPars[eBoundQOverP] =
0267       params.values[4] * std::sin(targetPars[eBoundTheta]) / Bz;
0268   targetPars[eBoundTime] = params.values[5];
0269
0270   return {params.surface, targetPars, cov, params.particleHypothesis};
0271 }
0272
0273 }  // namespace detail
0274
0275 edm4hep::TrackState writeTrackState(const GeometryContext& gctx,
0276                                     std::int32_t location,
0277                                     const BoundTrackParameters& params,
0278                                     const MagneticFieldProvider& magneticField,
0279                                     MagneticFieldProvider::Cache& fieldCache,
0280                                     const Logger& logger) {
0281   // The curvature depends on the local field, so evaluate it where the
0282   // parameters actually are.
0283   const double Bz =
0284       detail::localBz(magneticField, fieldCache, params.position(gctx), logger);
0285
0286   // Convert to the LCIO track parametrization expected by EDM4hep. This also
0287   // re-expresses the parameters at a perigee surface if they are not on one
0288   // already.
0289   const detail::Parameters converted =
0290       detail::convertTrackParametersToEdm4hep(gctx, Bz, params);
0291
0292   edm4hep::TrackState trackState{};
0293   trackState.location = location;
0294   trackState.D0 = static_cast<float>(converted.values[0]);
0295   trackState.Z0 = static_cast<float>(converted.values[1]);
0296   trackState.phi = static_cast<float>(converted.values[2]);
0297   trackState.tanLambda = static_cast<float>(converted.values[3]);
0298   trackState.omega = static_cast<float>(converted.values[4]);
0299   trackState.time = static_cast<float>(converted.values[5]);
0300
0301   if (converted.covariance) {
0302     detail::packCovariance(converted.covariance.value(), trackState.covMatrix);
0303   }
0304
0305   // The reference point is the centre of the perigee surface the converted
0306   // parameters are expressed on.
0307   const Vector3 center = converted.surface->center(gctx);
0308   trackState.referencePoint.x = static_cast<float>(center.x());
0309   trackState.referencePoint.y = static_cast<float>(center.y());
0310   trackState.referencePoint.z = static_cast<float>(center.z());
0311
0312   ACTS_VERBOSE("- parameters: " << params.parameters().transpose() << " -> "
0313                                 << converted.values.transpose());
0314   if (converted.covariance) {
0315     ACTS_VERBOSE("- covariance: \n"
0316                  << params.covariance().value_or(BoundMatrix::Zero())
0317                  << "\n->\n"
0318                  << converted.covariance.value());
0319   }
0320   ACTS_VERBOSE("- ref surface ctr: " << center.transpose());
0321
0322   return trackState;
0323 }
0324
0325 edm4hep::TrackState writeTrackState(const GeometryContext& gctx,
0326                                     const MagneticFieldContext& mctx,
0327                                     std::int32_t location,
0328                                     const BoundTrackParameters& params,
0329                                     const MagneticFieldProvider& magneticField,
0330                                     const Logger& logger) {
0331   auto fieldCache = magneticField.makeCache(mctx);
0332   return writeTrackState(gctx, location, params, magneticField, fieldCache,
0333                          logger);
0334 }
0335
0336 BoundTrackParameters readTrackState(const GeometryContext& gctx,
0337                                     const edm4hep::TrackState& trackState,
0338                                     const MagneticFieldProvider& magneticField,
0339                                     MagneticFieldProvider::Cache& fieldCache,
0340                                     const Logger& logger) {
0341   const detail::Parameters params = detail::unpackTrackState(trackState);
0342
0343   // Evaluate the local field where the particle actually is, matching what
0344   // writeTrackState does
0345   const double Bz =
0346       detail::localBz(magneticField, fieldCache,
0347                       detail::pointOfClosestApproach(gctx, params), logger);
0348
0349   BoundTrackParameters converted =
0350       detail::convertTrackParametersFromEdm4hep(Bz, params);
0351
0352   ACTS_VERBOSE("- parameters: " << params.values.transpose() << " -> "
0353                                 << converted.parameters().transpose());
0354   ACTS_VERBOSE(
0355       "- ref surface ctr: " << params.surface->center(gctx).transpose());
0356
0357   return converted;
0358 }
0359
0360 BoundTrackParameters readTrackState(const GeometryContext& gctx,
0361                                     const MagneticFieldContext& mctx,
0362                                     const edm4hep::TrackState& trackState,
0363                                     const MagneticFieldProvider& magneticField,
0364                                     const Logger& logger) {
0365   auto fieldCache = magneticField.makeCache(mctx);
0366   return readTrackState(gctx, trackState, magneticField, fieldCache, logger);
0367 }
0368
0369 edm4hep::MCParticle getParticle(const edm4hep::SimTrackerHit& hit) {
0370   return hit.getParticle();
0371 }
0372
0373 void setParticle(edm4hep::MutableSimTrackerHit& hit,
0374                  const edm4hep::MCParticle& particle) {
0375   hit.setParticle(particle);
0376 }
0377
0378 std::size_t SimHitAssociation::size() const {
0379   return m_internalToEdm4hep.size();
0380 }
0381
0382 void SimHitAssociation::reserve(std::size_t size) {
0383   m_internalToEdm4hep.reserve(size);
0384 }
0385
0386 void SimHitAssociation::add(std::size_t internalIndex,
0387                             const edm4hep::SimTrackerHit& edm4hepHit) {
0388   m_internalToEdm4hep.push_back(edm4hepHit);
0389   // m_edm4hepToInternal.at(edm4hepHit.id()) = internalIndex;
0390   m_edm4hepToInternal.emplace(edm4hepHit.id(), internalIndex);
0391 }
0392
0393 edm4hep::SimTrackerHit SimHitAssociation::lookup(
0394     std::size_t internalIndex) const {
0395   return m_internalToEdm4hep.at(internalIndex);
0396 }
0397
0398 std::size_t SimHitAssociation::lookup(const edm4hep::SimTrackerHit& hit) const {
0399   return m_edm4hepToInternal.at(hit.id());
0400 }
0401
0402 void writeVertex(const Vertex& vertex, edm4hep::MutableVertex to) {
0403   static constexpr std::array<edm4hep::FourMomCoords, 4> toEdm4hep = []() {
0404     std::array<edm4hep::FourMomCoords, 4> values{};
0405     using enum edm4hep::FourMomCoords;
0406     values.at(eFreePos0) = x;
0407     values.at(eFreePos1) = y;
0408     values.at(eFreePos2) = z;
0409     values.at(eFreeTime) = t;
0410     return values;
0411   }();
0412
0413   // Wrap this in a templated lambda so we can use `if constexpr` to select the
0414   // correct write function based on the type of the properties of the `to`
0415   // object.
0416   auto writeVertex = []<typename T>(const Vertex& vertex_, T& to_)
0417     requires(std::is_same_v<T, edm4hep::MutableVertex>)
0418   {
0419     if constexpr (detail::kEdm4hepVertexHasTime) {
0420       Vector4 pos = vertex_.fullPosition();
0421       to_.setPosition({static_cast<float>(pos[eFreePos0]),
0422                        static_cast<float>(pos[eFreePos1]),
0423                        static_cast<float>(pos[eFreePos2]),
0424                        static_cast<float>(pos[eFreeTime])});
0425
0426       edm4hep::CovMatrix4f& cov = to_.getCovMatrix();
0427       std::array coords{eFreePos0, eFreePos1, eFreePos2, eFreeTime};
0428       for (auto i : coords) {
0429         for (auto j : coords) {
0430           cov.setValue(static_cast<float>(vertex_.fullCovariance()(i, j)),
0431                        toEdm4hep.at(i), toEdm4hep.at(j));
0432         }
0433       }
0434     } else {
0435       Vector3 pos = vertex_.position();
0436       to_.setPosition({static_cast<float>(pos[eFreePos0]),
0437                        static_cast<float>(pos[eFreePos1]),
0438                        static_cast<float>(pos[eFreePos2])});
0439       edm4hep::CovMatrix3f& cov = to_.getCovMatrix();
0440       std::array coords{eFreePos0, eFreePos1, eFreePos2};
0441       for (auto i : coords) {
0442         for (auto j : coords) {
0443           cov.setValue(static_cast<float>(vertex_.covariance()(i, j)),
0444                        toEdm4hep.at(i), toEdm4hep.at(j));
0445         }
0446       }
0447     }
0448
0449     to_.setChi2(static_cast<float>(vertex_.fitQuality().first));
0450     to_.setNdf(static_cast<int>(vertex_.fitQuality().second));
0451   };
0452
0453   writeVertex(vertex, to);
0454 }
0455
0456 void writeMeasurement(const GeometryContext& gctx,
0457                       const Eigen::Map<const DynamicVector>& parameters,
0458                       const Eigen::Map<const DynamicMatrix>& covariance,
0459                       std::span<const std::uint8_t> indices,
0460                       std::uint64_t cellId, const Acts::Surface& surface,
0461                       ActsPodioEdm::MutableTrackerHitLocal& to) {
0462   if (parameters.size() != covariance.rows() ||
0463       covariance.rows() != covariance.cols() || parameters.size() < 0 ||
0464       indices.size() != static_cast<std::size_t>(parameters.size())) {
0465     throw std::runtime_error(
0466         "Size mismatch between parameters and covariance matrix");
0467   }
0468
0469   auto dim = static_cast<std::size_t>(parameters.size());
0470
0471   if (cellId != 0) {
0472     to.setCellID(cellId);
0473   }
0474
0475   to.setSubspaceIndices({indices.begin(), indices.end()});
0476
0477   auto loc0 = std::ranges::find(indices, eBoundLoc0);
0478   auto loc1 = std::ranges::find(indices, eBoundLoc1);
0479   auto time = std::ranges::find(indices, eBoundTime);
0480
0481   if (loc0 != indices.end() && loc1 != indices.end()) {
0482     Vector2 loc{parameters[std::distance(indices.begin(), loc0)],
0483                 parameters[std::distance(indices.begin(), loc1)]};
0484     Vector3 global = surface.localToGlobal(gctx, loc, Vector3::UnitZ());
0485     global /= Acts::UnitConstants::mm;
0486     to.setPosition({global.x(), global.y(), global.z()});
0487   }
0488
0489   if (time != indices.end()) {
0490     std::size_t timeOffset = std::distance(indices.begin(), time);
0491     to.setTime(
0492         static_cast<float>(parameters[timeOffset] / Acts::UnitConstants::ns));
0493   }
0494
0495   for (std::size_t i = 0; i < dim; ++i) {
0496     to.setValue(static_cast<float>(parameters[i]), i);
0497   }
0498
0499   for (std::size_t i = 0; i < dim; ++i) {
0500     for (std::size_t j = 0; j < dim; ++j) {
0501       to.setCov(static_cast<float>(covariance(i, j)), i, j);
0502     }
0503   }
0504 }
0505
0506 MeasurementData readMeasurement(const ActsPodioEdm::TrackerHitLocal& from) {
0507   auto indices = from.getSubspaceIndices();
0508   auto meas = from.getMeasurement();
0509   auto cov = from.getCovariance();
0510
0511   const auto dim = indices.size();
0512   if (meas.size() != dim || cov.size() != dim * dim) {
0513     throw std::runtime_error(
0514         std::format("Size mismatch in EDM4hep tracker hit: measurement size "
0515                     "{}, covariance size {}, indices size {}",
0516                     meas.size(), cov.size(), dim));
0517   }
0518
0519   MeasurementData result;
0520   result.cellId = from.getCellID();
0521   result.indices = std::move(indices);
0522
0523   for (std::size_t i = 0; i < dim; ++i) {
0524     result.parameters(result.indices[i]) = meas[i];
0525   }
0526   for (std::size_t i = 0; i < dim; ++i) {
0527     for (std::size_t j = 0; j < dim; ++j) {
0528       result.covariance(result.indices[i], result.indices[j]) =
0529           from.getCov(i, j);
0530     }
0531   }
0532
0533   return result;
0534 }
0535
0536 void writeTrackerHitSimHitLinks(
0537     const ActsPodioEdm::TrackerHitLocalCollection& hits,
0538     ActsPodioEdm::TrackerHitLocalSimTrackerHitLinkCollection& links,
0539     const SimHitForHitIndex& lookup) {
0540   for (std::size_t i = 0; i < static_cast<std::size_t>(hits.size()); ++i) {
0541     auto simHit = lookup(i);
0542     if (!simHit.has_value()) {
0543       continue;
0544     }
0545     auto link = links.create();
0546     link.setFrom(hits.at(i));
0547     link.setTo(simHit.value());
0548   }
0549 }
0550
0551 }  // namespace ActsPlugins::EDM4hepUtil