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/Vertexing/Vertex.hpp"
0017 #include "ActsPodioEdm/TrackerHitLocalCollection.h"
0018 #include "ActsPodioEdm/TrackerHitLocalSimTrackerHitLinkCollection.h"
0019
0020 #include <numbers>
0021
0022 #include <edm4hep/EDM4hepVersion.h>
0023 #include <edm4hep/MCParticle.h>
0024 #include <edm4hep/MutableSimTrackerHit.h>
0025 #include <edm4hep/MutableVertex.h>
0026 #include <edm4hep/SimTrackerHit.h>
0027 #include <edm4hep/TrackState.h>
0028 #include <edm4hep/Vector3f.h>
0029 #include <edm4hep/Vector4f.h>
0030
0031 using namespace Acts;
0032 using namespace Acts::detail;
0033
0034 namespace ActsPlugins::EDM4hepUtil {
0035 namespace detail {
0036
0037 SquareMatrix<6> jacobianToEdm4hep(double theta, double qOverP, double Bz) {
0038   // Calculate jacobian from our internal parametrization (d0, z0, phi, theta,
0039   // q/p) to the LCIO / edm4hep (see:
0040   // https://bib-pubdb1.desy.de/record/81214/files/LC-DET-2006-004%5B1%5D.pdf)
0041   // one (d0, z0, phi, tan(lambda), omega). Top left 3x3 matrix in the
0042   // jacobian is 1. Bottom right 2x2 matrix is:
0043   //
0044   // [  d                                 ]
0045   // [------(cot(theta))         0        ]
0046   // [dtheta                              ]
0047   // [                                    ]
0048   // [  d   /  B*q/p   \   d  /  B*q/p   \]
0049   // [------|----------|  ----|----------|]
0050   // [dtheta\sin(theta)/  dq/p\sin(theta)/]
0051   //
0052   // =
0053   //
0054   // [     2                        ]
0055   // [- csc (theta)           0     ]
0056   // [                              ]
0057   // [-B*q/p*cos(theta)       B     ]
0058   // [------------------  ----------]
0059   // [      2             sin(theta)]
0060   // [   sin (theta)                ]
0061
0062   SquareMatrix<6> J;
0063   J.setIdentity();
0064   double sinTheta = std::sin(theta);
0065   J(3, 3) = -1.0 / (sinTheta * sinTheta);
0066   J(4, 4) = Bz / sinTheta;  // dOmega / d(qop)
0067   J(4, 3) = -Bz * qOverP * std::cos(theta) /
0068             (sinTheta * sinTheta);  // dOmega / dTheta
0069   return J;
0070 }
0071
0072 SquareMatrix<6> jacobianFromEdm4hep(double tanLambda, double omega, double Bz) {
0073   // [     d      /                     pi\                                  ]
0074   // [------------|-atan(\tan\lambda) + --|                 0                ]
0075   // [d\tan\lambda\                     2 /                                  ]
0076   // [                                                                       ]
0077   // [     d      /         \Omega        \     d   /         \Omega        \]
0078   // [------------|-----------------------|  -------|-----------------------|]
0079   // [d\tan\lambda|     __________________|  d\Omega|     __________________|]
0080   // [            |    /            2     |         |    /            2     |]
0081   // [            \B*\/  \tan\lambda  + 1 /         \B*\/  \tan\lambda  + 1 /]
0082   //
0083   // =
0084   //
0085   // [         -1                                     ]
0086   // [   ----------------                 0           ]
0087   // [              2                                 ]
0088   // [   \tan\lambda  + 1                             ]
0089   // [                                                ]
0090   // [  -\Omega*\tan\lambda               1           ]
0091   // [-----------------------  -----------------------]
0092   // [                    3/2       __________________]
0093   // [  /           2    \         /            2     ]
0094   // [B*\\tan\lambda  + 1/     B*\/  \tan\lambda  + 1 ]
0095
0096   SquareMatrix<6> J;
0097   J.setIdentity();
0098   J(3, 3) = -1 / (tanLambda * tanLambda + 1);
0099   J(4, 3) = -1 * omega * tanLambda /
0100             (Bz * std::pow(tanLambda * tanLambda + 1, 3. / 2.));
0101   J(4, 4) = 1 / (Bz * std::hypot(tanLambda, 1));
0102
0103   return J;
0104 }
0105
0106 void packCovariance(const SquareMatrix<6>& from, float* to) {
0107   for (int i = 0; i < from.rows(); i++) {
0108     for (int j = 0; j <= i; j++) {
0109       std::size_t k = (i + 1) * i / 2 + j;
0110       to[k] = static_cast<float>(from(i, j));
0111     }
0112   }
0113 }
0114
0115 void unpackCovariance(const float* from, SquareMatrix<6>& to) {
0116   auto k = [](std::size_t i, std::size_t j) { return (i + 1) * i / 2 + j; };
0117   for (int i = 0; i < to.rows(); i++) {
0118     for (int j = 0; j < to.cols(); j++) {
0119       to(i, j) = from[j <= i ? k(i, j) : k(j, i)];
0120     }
0121   }
0122 }
0123
0124 Parameters convertTrackParametersToEdm4hep(const GeometryContext& gctx,
0125                                            double Bz,
0126                                            const BoundTrackParameters& params) {
0127   Vector3 global = params.referenceSurface().localToGlobal(
0128       gctx, params.parameters().template head<2>(), params.direction());
0129
0130   std::shared_ptr<const Surface> refSurface =
0131       params.referenceSurface().getSharedPtr();
0132   BoundVector targetPars = params.parameters();
0133   std::optional<BoundMatrix> targetCov = params.covariance();
0134
0135   // If the reference surface is a perigee surface, we use that. Otherwise
0136   // we create a new perigee surface at the global position of the track
0137   // parameters.
0138   if (dynamic_cast<const PerigeeSurface*>(refSurface.get()) == nullptr) {
0139     refSurface = Surface::makeShared<PerigeeSurface>(global);
0140
0141     // We need to convert to the target parameters
0142     // Keep the free parameters around we might need them for the covariance
0143     // conversion
0144
0145     auto perigeeParams =
0146         boundToBoundConversion(gctx, params, *refSurface, Vector3{0, 0, Bz})
0147             .value();
0148     targetPars = perigeeParams.parameters();
0149     targetCov = perigeeParams.covariance();
0150   }
0151
0152   Parameters result;
0153   result.surface = refSurface;
0154
0155   // Only run covariance conversion if we have a covariance input
0156   if (targetCov) {
0157     SquareMatrix<6> J = jacobianToEdm4hep(targetPars[eBoundTheta],
0158                                           targetPars[eBoundQOverP], Bz);
0159     result.covariance = J * targetCov.value() * J.transpose();
0160   }
0161
0162   result.values[0] = targetPars[eBoundLoc0];
0163   result.values[1] = targetPars[eBoundLoc1];
0164   result.values[2] = targetPars[eBoundPhi];
0165   result.values[3] = std::tan(std::numbers::pi / 2. - targetPars[eBoundTheta]);
0166   result.values[4] =
0167       targetPars[eBoundQOverP] / std::sin(targetPars[eBoundTheta]) * Bz;
0168   result.values[5] = targetPars[eBoundTime];
0169
0170   result.particleHypothesis = params.particleHypothesis();
0171
0172   return result;
0173 }
0174
0175 BoundTrackParameters convertTrackParametersFromEdm4hep(
0176     double Bz, const Parameters& params) {
0177   BoundVector targetPars;
0178
0179   SquareMatrix<6> J =
0180       jacobianFromEdm4hep(params.values[3], params.values[4], Bz);
0181
0182   std::optional<BoundMatrix> cov;
0183   if (params.covariance.has_value()) {
0184     cov = J * params.covariance.value() * J.transpose();
0185   }
0186
0187   targetPars[eBoundLoc0] = params.values[0];
0188   targetPars[eBoundLoc1] = params.values[1];
0189   targetPars[eBoundPhi] = params.values[2];
0190   targetPars[eBoundTheta] = std::numbers::pi / 2. - std::atan(params.values[3]);
0191   targetPars[eBoundQOverP] =
0192       params.values[4] * std::sin(targetPars[eBoundTheta]) / Bz;
0193   targetPars[eBoundTime] = params.values[5];
0194
0195   return {params.surface, targetPars, cov, params.particleHypothesis};
0196 }
0197
0198 }  // namespace detail
0199
0200 #if EDM4HEP_VERSION_MAJOR >= 1 || \
0201     (EDM4HEP_VERSION_MAJOR == 0 && EDM4HEP_VERSION_MINOR == 99)
0202 edm4hep::MCParticle getParticle(const edm4hep::SimTrackerHit& hit) {
0203   return hit.getParticle();
0204 }
0205
0206 void setParticle(edm4hep::MutableSimTrackerHit& hit,
0207                  const edm4hep::MCParticle& particle) {
0208   hit.setParticle(particle);
0209 }
0210 #else
0211 edm4hep::MCParticle getParticle(const edm4hep::SimTrackerHit& hit) {
0212   return hit.getMCParticle();
0213 }
0214
0215 void setParticle(edm4hep::MutableSimTrackerHit& hit,
0216                  const edm4hep::MCParticle& particle) {
0217   hit.setMCParticle(particle);
0218 }
0219 #endif
0220
0221 std::size_t SimHitAssociation::size() const {
0222   return m_internalToEdm4hep.size();
0223 }
0224
0225 void SimHitAssociation::reserve(std::size_t size) {
0226   m_internalToEdm4hep.reserve(size);
0227 }
0228
0229 void SimHitAssociation::add(std::size_t internalIndex,
0230                             const edm4hep::SimTrackerHit& edm4hepHit) {
0231   m_internalToEdm4hep.push_back(edm4hepHit);
0232   // m_edm4hepToInternal.at(edm4hepHit.id()) = internalIndex;
0233   m_edm4hepToInternal.emplace(edm4hepHit.id(), internalIndex);
0234 }
0235
0236 edm4hep::SimTrackerHit SimHitAssociation::lookup(
0237     std::size_t internalIndex) const {
0238   return m_internalToEdm4hep.at(internalIndex);
0239 }
0240
0241 std::size_t SimHitAssociation::lookup(const edm4hep::SimTrackerHit& hit) const {
0242   return m_edm4hepToInternal.at(hit.id());
0243 }
0244
0245 void writeVertex(const Vertex& vertex, edm4hep::MutableVertex to) {
0246   static constexpr std::array<edm4hep::FourMomCoords, 4> toEdm4hep = []() {
0247     std::array<edm4hep::FourMomCoords, 4> values{};
0248     using enum edm4hep::FourMomCoords;
0249     values.at(eFreePos0) = x;
0250     values.at(eFreePos1) = y;
0251     values.at(eFreePos2) = z;
0252     values.at(eFreeTime) = t;
0253     return values;
0254   }();
0255
0256   // Wrap this in a templated lambda so we can use `if constexpr` to select the
0257   // correct write function based on the type of the properties of the `to`
0258   // object.
0259   auto writeVertex = []<typename T>(const Vertex& vertex_, T& to_)
0260     requires(std::is_same_v<T, edm4hep::MutableVertex>)
0261   {
0262     if constexpr (detail::kEdm4hepVertexHasTime) {
0263       Vector4 pos = vertex_.fullPosition();
0264       to_.setPosition({static_cast<float>(pos[eFreePos0]),
0265                        static_cast<float>(pos[eFreePos1]),
0266                        static_cast<float>(pos[eFreePos2]),
0267                        static_cast<float>(pos[eFreeTime])});
0268
0269       edm4hep::CovMatrix4f& cov = to_.getCovMatrix();
0270       std::array coords{eFreePos0, eFreePos1, eFreePos2, eFreeTime};
0271       for (auto i : coords) {
0272         for (auto j : coords) {
0273           cov.setValue(static_cast<float>(vertex_.fullCovariance()(i, j)),
0274                        toEdm4hep.at(i), toEdm4hep.at(j));
0275         }
0276       }
0277     } else {
0278       Vector3 pos = vertex_.position();
0279       to_.setPosition({static_cast<float>(pos[eFreePos0]),
0280                        static_cast<float>(pos[eFreePos1]),
0281                        static_cast<float>(pos[eFreePos2])});
0282       edm4hep::CovMatrix3f& cov = to_.getCovMatrix();
0283       std::array coords{eFreePos0, eFreePos1, eFreePos2};
0284       for (auto i : coords) {
0285         for (auto j : coords) {
0286           cov.setValue(static_cast<float>(vertex_.covariance()(i, j)),
0287                        toEdm4hep.at(i), toEdm4hep.at(j));
0288         }
0289       }
0290     }
0291
0292     to_.setChi2(static_cast<float>(vertex_.fitQuality().first));
0293     to_.setNdf(static_cast<int>(vertex_.fitQuality().second));
0294   };
0295
0296   writeVertex(vertex, to);
0297 }
0298
0299 void writeMeasurement(const GeometryContext& gctx,
0300                       const Eigen::Map<const DynamicVector>& parameters,
0301                       const Eigen::Map<const DynamicMatrix>& covariance,
0302                       std::span<const std::uint8_t> indices,
0303                       std::uint64_t cellId, const Acts::Surface& surface,
0304                       ActsPodioEdm::MutableTrackerHitLocal& to) {
0305   if (parameters.size() != covariance.rows() ||
0306       covariance.rows() != covariance.cols() || parameters.size() < 0 ||
0307       indices.size() != static_cast<std::size_t>(parameters.size())) {
0308     throw std::runtime_error(
0309         "Size mismatch between parameters and covariance matrix");
0310   }
0311
0312   auto dim = static_cast<std::size_t>(parameters.size());
0313
0314   if (cellId != 0) {
0315     to.setCellID(cellId);
0316   }
0317
0318   to.setSubspaceIndices({indices.begin(), indices.end()});
0319
0320   auto loc0 = std::ranges::find(indices, eBoundLoc0);
0321   auto loc1 = std::ranges::find(indices, eBoundLoc1);
0322   auto time = std::ranges::find(indices, eBoundTime);
0323
0324   if (loc0 != indices.end() && loc1 != indices.end()) {
0325     Vector2 loc{parameters[std::distance(indices.begin(), loc0)],
0326                 parameters[std::distance(indices.begin(), loc1)]};
0327     Vector3 global = surface.localToGlobal(gctx, loc, Vector3::UnitZ());
0328     global /= Acts::UnitConstants::mm;
0329     to.setPosition({global.x(), global.y(), global.z()});
0330   }
0331
0332   if (time != indices.end()) {
0333     std::size_t timeOffset = std::distance(indices.begin(), time);
0334     to.setTime(
0335         static_cast<float>(parameters[timeOffset] / Acts::UnitConstants::ns));
0336   }
0337
0338   for (std::size_t i = 0; i < dim; ++i) {
0339     to.setValue(static_cast<float>(parameters[i]), i);
0340   }
0341
0342   for (std::size_t i = 0; i < dim; ++i) {
0343     for (std::size_t j = 0; j < dim; ++j) {
0344       to.setCov(static_cast<float>(covariance(i, j)), i, j);
0345     }
0346   }
0347 }
0348
0349 MeasurementData readMeasurement(const ActsPodioEdm::TrackerHitLocal& from) {
0350   auto indices = from.getSubspaceIndices();
0351   auto meas = from.getMeasurement();
0352   auto cov = from.getCovariance();
0353
0354   const auto dim = indices.size();
0355   if (meas.size() != dim || cov.size() != dim * dim) {
0356     throw std::runtime_error(
0357         std::format("Size mismatch in EDM4hep tracker hit: measurement size "
0358                     "{}, covariance size {}, indices size {}",
0359                     meas.size(), cov.size(), dim));
0360   }
0361
0362   MeasurementData result;
0363   result.cellId = from.getCellID();
0364   result.indices = std::move(indices);
0365
0366   for (std::size_t i = 0; i < dim; ++i) {
0367     result.parameters(result.indices[i]) = meas[i];
0368   }
0369   for (std::size_t i = 0; i < dim; ++i) {
0370     for (std::size_t j = 0; j < dim; ++j) {
0371       result.covariance(result.indices[i], result.indices[j]) =
0372           from.getCov(i, j);
0373     }
0374   }
0375
0376   return result;
0377 }
0378
0379 void writeTrackerHitSimHitLinks(
0380     const ActsPodioEdm::TrackerHitLocalCollection& hits,
0381     ActsPodioEdm::TrackerHitLocalSimTrackerHitLinkCollection& links,
0382     const SimHitForHitIndex& lookup) {
0383   for (std::size_t i = 0; i < static_cast<std::size_t>(hits.size()); ++i) {
0384     auto simHit = lookup(i);
0385     if (!simHit.has_value()) {
0386       continue;
0387     }
0388     auto link = links.create();
0389     link.setFrom(hits.at(i));
0390     link.setTo(simHit.value());
0391   }
0392 }
0393
0394 }  // namespace ActsPlugins::EDM4hepUtil