algorithms/fardetectors/FarDetectorLinearTracking.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023 - 2025, Simon Gardner
0003
0004 #include <DD4hep/VolumeManager.h>
0005 #include <Evaluator/DD4hepUnits.h>
0006 #include <Math/GenVector/Cartesian3D.h>
0007 #include <Math/GenVector/DisplacementVector3D.h>
0008 #include <algorithms/geo.h>
0009 #include <edm4eic/Cov6f.h>
0010 #include <edm4eic/EDM4eicVersion.h>
0011 #include <edm4eic/MCRecoTrackParticleAssociationCollection.h>
0012 #include <edm4eic/MCRecoTrackerHitAssociationCollection.h>
0013 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0014 #include <edm4eic/MCRecoTrackerHitLinkCollection.h>
0015 #endif
0016 #include <edm4eic/Measurement2DCollection.h>
0017 #include <edm4eic/RawTrackerHit.h>
0018 #include <edm4eic/TrackCollection.h>
0019 #include <edm4eic/TrackerHit.h>
0020 #include <edm4hep/MCParticle.h>
0021 #include <edm4hep/SimTrackerHit.h>
0022 #include <edm4hep/Vector2f.h>
0023 #include <edm4hep/Vector3d.h>
0024 #include <edm4hep/Vector3f.h>
0025 #include <edm4hep/utils/vector_utils.h>
0026 #include <podio/LinkNavigator.h>
0027 #include <podio/RelationRange.h>
0028 #include <podio/detail/Link.h>
0029 #include <Eigen/Geometry>
0030 #include <Eigen/Householder>
0031 #include <Eigen/Jacobi>
0032 #include <Eigen/SVD>
0033 #include <algorithm>
0034 #include <cmath>
0035 #include <cstddef>
0036 #include <cstdint>
0037 #include <memory>
0038 #include <new>
0039 #include <tuple>
0040 #include <unordered_map>
0041 #include <utility>
0042
0043 #include "FarDetectorLinearTracking.h"
0044 #include "algorithms/fardetectors/FarDetectorLinearTrackingConfig.h"
0045
0046 namespace eicrecon {
0047
0048 void FarDetectorLinearTracking::init() {
0049
0050   // For changing how strongly each layer hit is in contributing to the fit
0051   m_layerWeights = Eigen::VectorXd::Constant(m_cfg.n_layer, 1);
0052
0053   for (std::size_t i = 0; i < std::min(m_cfg.layer_weights.size(), m_cfg.n_layer); i++) {
0054     m_layerWeights(i) = m_cfg.layer_weights[i];
0055   }
0056
0057   // For checking the direction of the track from theta and phi angles
0058   m_optimumDirection = Eigen::Vector3d::UnitZ();
0059   m_optimumDirection =
0060       Eigen::AngleAxisd(m_cfg.optimum_theta, Eigen::Vector3d::UnitY()) * m_optimumDirection;
0061   m_optimumDirection =
0062       Eigen::AngleAxisd(m_cfg.optimum_phi, Eigen::Vector3d::UnitZ()) * m_optimumDirection;
0063
0064   m_cellid_converter = algorithms::GeoSvc::instance().cellIDPositionConverter();
0065 }
0066
0067 void FarDetectorLinearTracking::process(const FarDetectorLinearTracking::Input& input,
0068                                         const FarDetectorLinearTracking::Output& output) const {
0069
0070 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0071   const auto [inputhits, hitLinks, assocHits]  = input;
0072   auto [outputTracks, trackLinks, assocTracks] = output;
0073 #else
0074   const auto [inputhits, assocHits] = input;
0075   auto [outputTracks, assocTracks]  = output;
0076 #endif
0077
0078   // Check the number of input collections is correct
0079   std::size_t nCollections = inputhits.size();
0080   if (nCollections != m_cfg.n_layer) {
0081     error("Wrong number of input collections passed to algorithm");
0082     return;
0083   }
0084
0085 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0086   // Check if truth associations are possible
0087   const bool do_assoc = hitLinks != nullptr && !hitLinks->empty();
0088   if (!do_assoc) {
0089     debug("Provided MCRecoTrackerHitLink collection is empty. No truth associations "
0090           "will be performed.");
0091   }
0092   // Build fast lookup once per event using podio::LinkNavigator
0093   std::optional<podio::LinkNavigator<edm4eic::MCRecoTrackerHitLinkCollection>> link_nav;
0094   if (do_assoc) {
0095     link_nav.emplace(*hitLinks);
0096   }
0097 #else
0098   const bool do_assoc = assocHits != nullptr && !assocHits->empty();
0099   if (!do_assoc) {
0100     debug("Provided MCRecoTrackerHitAssociation collection is empty. No truth associations "
0101           "will be performed.");
0102   }
0103 #endif
0104
0105   std::vector<std::vector<Eigen::Vector3d>> convertedHits;
0106   std::vector<std::vector<edm4hep::MCParticle>> assocParts;
0107   convertedHits.reserve(m_cfg.n_layer);
0108   assocParts.reserve(m_cfg.n_layer);
0109
0110   // Check there aren't too many hits in any layer to handle
0111   // Temporary limit of number of hits per layer before Kalman filtering/GNN implemented
0112   // TODO - Implement more sensible solution
0113   for (const auto& layerHits : inputhits) {
0114     if ((*layerHits).size() > m_cfg.layer_hits_max) {
0115       info("Too many hits in layer");
0116       return;
0117     }
0118     if ((*layerHits).empty()) {
0119       trace("No hits in layer");
0120       return;
0121     }
0122 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0123     ConvertClusters(*layerHits, *link_nav, *assocHits, convertedHits, assocParts);
0124 #else
0125     ConvertClusters(*layerHits, *assocHits, convertedHits, assocParts);
0126 #endif
0127   }
0128
0129   // Create a matrix to store the hit positions
0130   Eigen::MatrixXd hitMatrix(3, m_cfg.n_layer);
0131
0132   // Create vector to store indexes of hits in the track
0133   std::vector<std::size_t> layerHitIndex(m_cfg.n_layer, 0);
0134
0135   int layer = 0;
0136
0137   // Iterate over all combinations of measurements in the layers without recursion
0138   while (true) {
0139     hitMatrix.col(layer) << convertedHits[layer][layerHitIndex[layer]];
0140
0141     bool isValid = true;
0142     // Check the last two hits are within a certain angle of the optimum direction
0143     if (layer > 0 && m_cfg.restrict_direction) {
0144       isValid = checkHitPair(hitMatrix.col(layer - 1), hitMatrix.col(layer));
0145     }
0146
0147     // If valid hit combination, move to the next layer or check the combination
0148     if (isValid) {
0149       if (layer == static_cast<long>(m_cfg.n_layer) - 1) {
0150         // Check the combination, if chi2 limit is passed, add the track to the output
0151         checkHitCombination(&hitMatrix, outputTracks,
0152 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0153                             trackLinks,
0154 #endif
0155                             assocTracks, inputhits, assocParts, layerHitIndex);
0156       } else {
0157         layer++;
0158         continue;
0159       }
0160     }
0161
0162     // Iterate current layer
0163     layerHitIndex[layer]++;
0164
0165     bool doBreak = false;
0166     // Set up next combination to check
0167     while (layerHitIndex[layer] >= convertedHits[layer].size()) {
0168       layerHitIndex[layer] = 0;
0169       if (layer == 0) {
0170         doBreak = true;
0171         break;
0172       }
0173       layer--;
0174       // Iterate previous layer
0175       layerHitIndex[layer]++;
0176     }
0177     if (doBreak) {
0178       break;
0179     }
0180   }
0181 }
0182
0183 void FarDetectorLinearTracking::checkHitCombination(
0184     Eigen::MatrixXd* hitMatrix, edm4eic::TrackCollection* outputTracks,
0185 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0186     edm4eic::MCRecoTrackParticleLinkCollection* trackLinks,
0187 #endif
0188     edm4eic::MCRecoTrackParticleAssociationCollection* assocTracks,
0189     const std::vector<gsl::not_null<const edm4eic::Measurement2DCollection*>>& inputHits,
0190     const std::vector<std::vector<edm4hep::MCParticle>>& assocParts,
0191     const std::vector<std::size_t>& layerHitIndex) const {
0192
0193   Eigen::Vector3d weightedAnchor = (*hitMatrix) * m_layerWeights / (m_layerWeights.sum());
0194
0195   auto localMatrix = (*hitMatrix).colwise() - weightedAnchor;
0196
0197   Eigen::BDCSVD<Eigen::MatrixXd> svd(localMatrix.transpose(),
0198                                      Eigen::ComputeThinU | Eigen::ComputeThinV);
0199
0200   auto V = svd.matrixV();
0201
0202   // Rotate into principle components and calculate chi2/ndf
0203   auto rotatedMatrix = localMatrix.transpose() * V;
0204   auto residuals     = rotatedMatrix.rightCols(2);
0205   double chi2        = (residuals.array() * residuals.array()).sum() / (2 * m_cfg.n_layer);
0206
0207   if (chi2 > m_cfg.chi2_max) {
0208     return;
0209   }
0210
0211   edm4hep::Vector3d outPos = weightedAnchor.data();
0212   edm4hep::Vector3d outVec = V.col(0).data();
0213
0214   // Make sure fit was pointing in the right direction
0215   if (outVec.z > 0) {
0216     outVec = outVec * -1;
0217   }
0218
0219   int32_t type{0};                                          // Type of track
0220   edm4hep::Vector3f position(outPos.x, outPos.y, outPos.z); // Position of the trajectory point [mm]
0221   edm4hep::Vector3f momentum(outVec.x, outVec.y, outVec.z); // 3-momentum at the point [GeV]
0222   edm4eic::Cov6f positionMomentumCovariance;                // Error on the position
0223   float time{0};                                            // Time at this point [ns]
0224   float timeError{0};                                       // Error on the time at this point
0225   float charge{-1};                                         // Charge of the particle
0226   int32_t ndf{static_cast<int32_t>(m_cfg.n_layer) - 1};     // Number of degrees of freedom
0227   int32_t pdg{11};                                          // PDG code of the particle
0228
0229   // Create the track
0230   auto track = (*outputTracks)
0231                    .create(type, position, momentum, positionMomentumCovariance, time, timeError,
0232                            charge, chi2, ndf, pdg);
0233
0234   // Add Measurement2D relations and count occurrence of particles contributing to the track
0235   std::unordered_map<edm4hep::MCParticle, int> particleCount;
0236   for (std::size_t layer = 0; layer < layerHitIndex.size(); layer++) {
0237     track.addToMeasurements((*inputHits[layer])[layerHitIndex[layer]]);
0238     const auto& assocParticle = assocParts[layer][layerHitIndex[layer]];
0239     particleCount[assocParticle]++;
0240   }
0241
0242   // Create track associations for each particle
0243   for (const auto& [particle, count] : particleCount) {
0244 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0245     auto trackLink = trackLinks->create();
0246     trackLink.setFrom(track);
0247     trackLink.setTo(particle);
0248     trackLink.setWeight(count / static_cast<double>(m_cfg.n_layer));
0249 #endif
0250     auto trackAssoc = assocTracks->create();
0251     trackAssoc.setRec(track);
0252     trackAssoc.setSim(particle);
0253     trackAssoc.setWeight(count / static_cast<double>(m_cfg.n_layer));
0254   }
0255 }
0256
0257 // Check if a pair of hits lies close to the optimum direction
0258 bool FarDetectorLinearTracking::checkHitPair(const Eigen::Vector3d& hit1,
0259                                              const Eigen::Vector3d& hit2) const {
0260
0261   Eigen::Vector3d hitDiff = hit2 - hit1;
0262   hitDiff.normalize();
0263
0264   double angle = std::acos(hitDiff.dot(m_optimumDirection));
0265
0266   debug("Vector: x={}, y={}, z={}", hitDiff.x(), hitDiff.y(), hitDiff.z());
0267   debug("Optimum: x={}, y={}, z={}", m_optimumDirection.x(), m_optimumDirection.y(),
0268         m_optimumDirection.z());
0269   debug("Angle: {}, Tolerance {}", angle, m_cfg.step_angle_tolerance);
0270
0271   return angle <= m_cfg.step_angle_tolerance;
0272 }
0273
0274 // Convert measurements into global coordinates
0275 void FarDetectorLinearTracking::ConvertClusters(
0276     const edm4eic::Measurement2DCollection& clusters,
0277 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0278     const podio::LinkNavigator<edm4eic::MCRecoTrackerHitLinkCollection>& link_nav,
0279 #endif
0280     [[maybe_unused]] const edm4eic::MCRecoTrackerHitAssociationCollection& assoc_hits,
0281     std::vector<std::vector<Eigen::Vector3d>>& pointPositions,
0282     std::vector<std::vector<edm4hep::MCParticle>>& assoc_parts) const {
0283
0284   // Get context of first hit
0285   const dd4hep::VolumeManagerContext* context =
0286       m_cellid_converter->findContext(clusters[0].getSurface());
0287
0288   std::vector<Eigen::Vector3d> layerPositions;
0289   std::vector<edm4hep::MCParticle> assocParticles;
0290
0291   for (auto cluster : clusters) {
0292
0293     auto globalPos = context->localToWorld({cluster.getLoc()[0], cluster.getLoc()[1], 0});
0294     layerPositions.emplace_back(globalPos.x() / dd4hep::mm, globalPos.y() / dd4hep::mm,
0295                                 globalPos.z() / dd4hep::mm);
0296
0297     // Determine the MCParticle associated with this measurement based on the weights
0298     // Get hit in measurement with max weight
0299     float maxWeight      = 0;
0300     std::size_t maxIndex = cluster.getWeights().size();
0301     for (std::size_t i = 0; i < cluster.getWeights().size(); ++i) {
0302       if (cluster.getWeights()[i] > maxWeight) {
0303         maxWeight = cluster.getWeights()[i];
0304         maxIndex  = i;
0305       }
0306     }
0307     if (maxIndex == cluster.getWeights().size()) {
0308       // no maximum found (e.g. all weights zero, cluster size zero)
0309       continue;
0310     }
0311     auto maxHit = cluster.getHits()[maxIndex];
0312     // Get associated raw hit
0313     auto rawHit = maxHit.getRawHit();
0314
0315 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0316     const auto sim_hits = link_nav.getLinked(rawHit);
0317     if (!sim_hits.empty()) {
0318       auto particle = sim_hits[0].o.getParticle();
0319       assocParticles.push_back(particle);
0320     }
0321 #else
0322     // Fallback: linear search through associations
0323     for (const auto& assoc : assoc_hits) {
0324       if (assoc.getRawHit() == rawHit) {
0325         auto particle = assoc.getSimHit().getParticle();
0326         assocParticles.push_back(particle);
0327         break;
0328       }
0329     }
0330 #endif
0331   }
0332
0333   pointPositions.push_back(layerPositions);
0334   assoc_parts.push_back(assocParticles);
0335 }
0336
0337 } // namespace eicrecon