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