src/components/ParticlesFromTrackFit.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Whitney Armstrong, Sylvester Joosten, Wouter Deconinck
0003
0004 #include <algorithm>
0005
0006 // Gaudi
0007 #include "Gaudi/Algorithm.h"
0008 #include "GaudiKernel/ToolHandle.h"
0009 #include "GaudiKernel/RndmGenerators.h"
0010 #include "Gaudi/Property.h"
0011
0012 #include "DDRec/CellIDPositionConverter.h"
0013 #include "DDRec/SurfaceManager.h"
0014 #include "DDRec/Surface.h"
0015
0016 #include <k4FWCore/DataHandle.h>
0017 #include <k4Interface/IGeoSvc.h>
0018
0019 #include "Acts/EventData/MultiTrajectory.hpp"
0020 #include "Acts/EventData/MultiTrajectoryHelpers.hpp"
0021
0022 // Event Model related classes
0023 #include "edm4eic/EDM4eicVersion.h"
0024 #include "edm4eic/ReconstructedParticleCollection.h"
0025 #include "edm4eic/TrackerHitCollection.h"
0026 #include "edm4eic/TrackParametersCollection.h"
0027 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0028 #include "ActsExamples/EventData/Track.hpp"
0029 #include "ActsExamples/EventData/Trajectories.hpp"
0030
0031 #include "Acts/Utilities/Helpers.hpp"
0032
0033 #include "edm4hep/utils/vector_utils.h"
0034
0035 #include <cmath>
0036
0037 namespace Jug::Reco {
0038
0039   #if EDM4EIC_VERSION_MAJOR >= 5
0040     // This array relates the Acts and EDM4eic covariance matrices, including
0041     // the unit conversion to get from Acts units into EDM4eic units.
0042     //
0043     // Note: std::map is not constexpr, so we use a constexpr std::array
0044     // std::array initialization need double braces since arrays are aggregates
0045     // ref: https://en.cppreference.com/w/cpp/language/aggregate_initialization
0046     static constexpr std::array<std::pair<Acts::BoundIndices, double>, 6> edm4eic_indexed_units{{
0047       {Acts::eBoundLoc0, Acts::UnitConstants::mm},
0048       {Acts::eBoundLoc1, Acts::UnitConstants::mm},
0049       {Acts::eBoundTheta, 1.},
0050       {Acts::eBoundPhi, 1.},
0051       {Acts::eBoundQOverP, 1. / Acts::UnitConstants::GeV},
0052       {Acts::eBoundTime, Acts::UnitConstants::ns}
0053     }};
0054   #endif
0055
0056   /** Extract the particles form fit trajectories.
0057    *
0058    * \ingroup tracking
0059    */
0060    class ParticlesFromTrackFit : public Gaudi::Algorithm {
0061    private:
0062     mutable DataHandle<ActsExamples::TrajectoriesContainer>     m_inputTrajectories{"inputTrajectories", Gaudi::DataHandle::Reader, this};
0063     mutable DataHandle<edm4eic::ReconstructedParticleCollection> m_outputParticles{"outputParticles", Gaudi::DataHandle::Writer, this};
0064     mutable DataHandle<edm4eic::TrackParametersCollection> m_outputTrackParameters{"outputTrackParameters", Gaudi::DataHandle::Writer, this};
0065
0066    public:
0067     //  ill-formed: using Gaudi::Algorithm::GaudiAlgorithm;
0068     ParticlesFromTrackFit(const std::string& name, ISvcLocator* svcLoc)
0069         : Gaudi::Algorithm(name, svcLoc) {
0070           declareProperty("inputTrajectories", m_inputTrajectories,"");
0071           declareProperty("outputParticles", m_outputParticles, "");
0072           declareProperty("outputTrackParameters", m_outputTrackParameters, "Acts Track Parameters");
0073         }
0074
0075     StatusCode initialize() override {
0076       if (Gaudi::Algorithm::initialize().isFailure()) {
0077         return StatusCode::FAILURE;
0078       }
0079       return StatusCode::SUCCESS;
0080     }
0081
0082     StatusCode execute(const EventContext&) const override {
0083       // input collection
0084       const auto* const trajectories = m_inputTrajectories.get();
0085       // create output collections
0086       auto* rec_parts = m_outputParticles.createAndPut();
0087       auto* track_pars = m_outputTrackParameters.createAndPut();
0088
0089       if (msgLevel(MSG::DEBUG)) {
0090         debug() << std::size(*trajectories) << " trajectories " << endmsg;
0091       }
0092
0093       // Loop over the trajectories
0094         for (const auto& traj : *trajectories) {
0095
0096           // Get the entry index for the single trajectory
0097           // The trajectory entry indices and the multiTrajectory
0098           const auto& mj        = traj.multiTrajectory();
0099           const auto& trackTips = traj.tips();
0100           if (trackTips.empty()) {
0101             if (msgLevel(MSG::DEBUG)) {
0102               debug() << "Empty multiTrajectory." << endmsg;
0103             }
0104             continue;
0105           }
0106
0107           const auto& trackTip = trackTips.front();
0108
0109           // Collect the trajectory summary info
0110           auto trajState       = Acts::MultiTrajectoryHelpers::trajectoryState(mj, trackTip);
0111           //int  m_nMeasurements = trajState.nMeasurements;
0112           //int  m_nStates       = trajState.nStates;
0113
0114           // Get the fitted track parameter
0115           //
0116           if (traj.hasTrackParameters(trackTip)) {
0117             const auto& boundParam = traj.trackParameters(trackTip);
0118             const auto& parameter  = boundParam.parameters();
0119             const auto& covariance = *boundParam.covariance();
0120             if (msgLevel(MSG::DEBUG)) {
0121               debug() << "loc 0 = " << parameter[Acts::eBoundLoc0] << endmsg;
0122               debug() << "loc 1 = " << parameter[Acts::eBoundLoc1] << endmsg;
0123               debug() << "phi   = " << parameter[Acts::eBoundPhi] << endmsg;
0124               debug() << "theta = " << parameter[Acts::eBoundTheta] << endmsg;
0125               debug() << "q/p   = " << parameter[Acts::eBoundQOverP] << endmsg;
0126               debug() << "p     = " << 1.0 / parameter[Acts::eBoundQOverP] << endmsg;
0127
0128               debug() << "err phi = " << sqrt(covariance(Acts::eBoundPhi, Acts::eBoundPhi)) << endmsg;
0129               debug() << "err th  = " << sqrt(covariance(Acts::eBoundTheta, Acts::eBoundTheta)) << endmsg;
0130               debug() << "err q/p = " << sqrt(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)) << endmsg;
0131
0132               debug() << " chi2 = " << trajState.chi2Sum << endmsg;
0133             }
0134
0135             auto pars = track_pars->create();
0136             pars.setType(0); // type: track head --> 0
0137             pars.setLoc({
0138               static_cast<float>(parameter[Acts::eBoundLoc0]),
0139               static_cast<float>(parameter[Acts::eBoundLoc1])
0140             });
0141             pars.setTheta(static_cast<float>(parameter[Acts::eBoundTheta]));
0142             pars.setPhi(static_cast<float>(parameter[Acts::eBoundPhi]));
0143             pars.setQOverP(static_cast<float>(parameter[Acts::eBoundQOverP]));
0144             pars.setTime(static_cast<float>(parameter[Acts::eBoundTime]));
0145             #if EDM4EIC_VERSION_MAJOR >= 5
0146               edm4eic::Cov6f cov;
0147               for (size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0148                 for (size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0149                   // FIXME why not pars.getCovariance()(i,j) = covariance(a,b) / x / y;
0150                   cov(i,j) = covariance(a,b) / x / y;
0151                 }
0152               }
0153               pars.setCovariance(cov);
0154             #else
0155               pars.setCharge(static_cast<float>(boundParam.charge()));
0156               pars.setLocError({
0157                     static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc0)),
0158                     static_cast<float>(covariance(Acts::eBoundLoc1, Acts::eBoundLoc1)),
0159                     static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc1))
0160                 });
0161               pars.setMomentumError({
0162                     static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundTheta)),
0163                     static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundPhi)),
0164                     static_cast<float>(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)),
0165                     static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundPhi)),
0166                     static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundQOverP)),
0167                     static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundQOverP))
0168                 });
0169               pars.setTimeError(sqrt(static_cast<float>(covariance(Acts::eBoundTime, Acts::eBoundTime))));
0170             #endif
0171           }
0172
0173           auto tsize = trackTips.size();
0174           if (msgLevel(MSG::DEBUG)) {
0175             debug() << "# fitted parameters : " << tsize << endmsg;
0176           }
0177           if (tsize == 0) {
0178             continue;
0179           }
0180
0181           mj.visitBackwards(tsize - 1, [&](auto&& trackstate) {
0182             // debug() << trackstate.hasPredicted() << endmsg;
0183             // debug() << trackstate.predicted() << endmsg;
0184             auto params = trackstate.predicted(); //<< endmsg;
0185
0186             double p0 = (1.0 / params[Acts::eBoundQOverP]) / Acts::UnitConstants::GeV;
0187             if (msgLevel(MSG::DEBUG)) {
0188               debug() << "track predicted p = " << p0 << " GeV" << endmsg;
0189             }
0190             if (std::abs(p0) > 500) {
0191               if (msgLevel(MSG::DEBUG)) {
0192                 debug() << "skipping" << endmsg;
0193               }
0194               return;
0195             }
0196
0197             auto rec_part = rec_parts->create();
0198             rec_part.setMomentum(
0199               edm4hep::utils::sphericalToVector(
0200                 1.0 / std::abs(params[Acts::eBoundQOverP]),
0201                 params[Acts::eBoundTheta],
0202                 params[Acts::eBoundPhi])
0203             );
0204             rec_part.setCharge(static_cast<int16_t>(std::copysign(1., params[Acts::eBoundQOverP])));
0205           });
0206       }
0207
0208       return StatusCode::SUCCESS;
0209     }
0210
0211   };
0212   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0213   DECLARE_COMPONENT(ParticlesFromTrackFit)
0214
0215 } // namespace Jug::Reco