Plugins/EDM4hep/EDM4hepUtil.hpp

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