Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-06-06 08:07:56

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2024 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 
0009 #pragma once
0010 
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Surfaces/PerigeeSurface.hpp"
0013 #include "Acts/Vertexing/LinearizedTrack.hpp"
0014 #include "Acts/Vertexing/TrackAtVertex.hpp"
0015 #include "Acts/Vertexing/Vertex.hpp"
0016 
0017 namespace Acts::KalmanVertexUpdater::detail {
0018 
0019 /// Cache object, the comments indicate the names of the variables in Ref. (1)
0020 /// @tparam nDimVertex number of dimensions of the vertex. Can be 3 (if we only
0021 /// fit its spatial coordinates) or 4 (if we also fit time).
0022 template <unsigned int nDimVertex>
0023 struct Cache {
0024   using VertexVector = ActsVector<nDimVertex>;
0025   using VertexMatrix = ActsSquareMatrix<nDimVertex>;
0026   // \tilde{x_k}
0027   VertexVector newVertexPos = VertexVector::Zero();
0028   // C_k
0029   VertexMatrix newVertexCov = VertexMatrix::Zero();
0030   // C_k^-1
0031   VertexMatrix newVertexWeight = VertexMatrix::Zero();
0032   // C_{k-1}^-1
0033   VertexMatrix oldVertexWeight = VertexMatrix::Zero();
0034   // W_k
0035   SquareMatrix3 wMat = SquareMatrix3::Zero();
0036 };
0037 
0038 /// @brief Calculates updated vertex position and covariance as well as the
0039 /// updated track momentum when adding/removing linTrack. Saves the result in
0040 /// cache.
0041 ///
0042 /// @tparam nDimVertex number of dimensions of the vertex. Can be 3 (if we only
0043 /// fit its spatial coordinates) or 4 (if we also fit time).
0044 ///
0045 /// @param vtx Vertex
0046 /// @param linTrack Linearized track to be added or removed
0047 /// @param trackWeight Track weight
0048 /// @param sign +1 (add track) or -1 (remove track)
0049 /// @note Tracks are removed during the smoothing procedure to compute
0050 /// the chi2 of the track wrt the updated vertex position
0051 /// @param[out] cache A cache to store the results of this function
0052 template <unsigned int nDimVertex>
0053 void calculateUpdate(const Vertex& vtx, const Acts::LinearizedTrack& linTrack,
0054                      const double trackWeight, const int sign,
0055                      Cache<nDimVertex>& cache) {
0056   constexpr unsigned int nBoundParams = nDimVertex + 2;
0057   using ParameterVector = ActsVector<nBoundParams>;
0058   using ParameterMatrix = ActsSquareMatrix<nBoundParams>;
0059   // Retrieve variables from the track linearization. The comments indicate the
0060   // corresponding symbol used in Ref. (1).
0061   // A_k
0062   const ActsMatrix<nBoundParams, nDimVertex> posJac =
0063       linTrack.positionJacobian.block<nBoundParams, nDimVertex>(0, 0);
0064   // B_k
0065   const ActsMatrix<nBoundParams, 3> momJac =
0066       linTrack.momentumJacobian.block<nBoundParams, 3>(0, 0);
0067   // p_k
0068   const ParameterVector trkParams =
0069       linTrack.parametersAtPCA.head<nBoundParams>();
0070   // c_k
0071   const ParameterVector constTerm = linTrack.constantTerm.head<nBoundParams>();
0072   // TODO we could use `linTrack.weightAtPCA` but only if we would always fit
0073   // time.
0074   // G_k
0075   // Note that, when removing a track, G_k -> - G_k, see Ref. (1).
0076   // Further note that, as we use the weighted formalism, the track weight
0077   // matrix (i.e., the inverse track covariance matrix) should be multiplied
0078   // with the track weight from the AMVF formalism. Here, we choose to
0079   // consider these two multiplicative factors directly in the updates of
0080   // newVertexWeight and newVertexPos.
0081   const ParameterMatrix trkParamWeight =
0082       linTrack.covarianceAtPCA.block<nBoundParams, nBoundParams>(0, 0)
0083           .inverse();
0084 
0085   // Retrieve current position of the vertex and its current weight matrix
0086   const ActsVector<nDimVertex> oldVtxPos =
0087       vtx.fullPosition().template head<nDimVertex>();
0088   // C_{k-1}^-1
0089   cache.oldVertexWeight =
0090       (vtx.fullCovariance().template block<nDimVertex, nDimVertex>(0, 0))
0091           .inverse();
0092 
0093   // W_k
0094   cache.wMat = (momJac.transpose() * (trkParamWeight * momJac)).inverse();
0095 
0096   // G_k^B = G_k - G_k*B_k*W_k*B_k^(T)*G_k
0097   ParameterMatrix gBMat = trkParamWeight - trkParamWeight * momJac *
0098                                                cache.wMat * momJac.transpose() *
0099                                                trkParamWeight;
0100 
0101   // C_k^-1
0102   cache.newVertexWeight = cache.oldVertexWeight + sign * trackWeight *
0103                                                       posJac.transpose() *
0104                                                       gBMat * posJac;
0105   // C_k
0106   cache.newVertexCov = cache.newVertexWeight.inverse();
0107 
0108   // \tilde{x_k}
0109   cache.newVertexPos =
0110       cache.newVertexCov * (cache.oldVertexWeight * oldVtxPos +
0111                             sign * trackWeight * posJac.transpose() * gBMat *
0112                                 (trkParams - constTerm));
0113 }
0114 
0115 template <unsigned int nDimVertex>
0116 double vertexPositionChi2Update(const Vector4& oldVtxPos,
0117                                 const Cache<nDimVertex>& cache) {
0118   ActsVector<nDimVertex> posDiff =
0119       cache.newVertexPos - oldVtxPos.template head<nDimVertex>();
0120 
0121   // Calculate and return corresponding chi2
0122   return posDiff.transpose() * (cache.oldVertexWeight * posDiff);
0123 }
0124 
0125 template <unsigned int nDimVertex>
0126 double trackParametersChi2(const LinearizedTrack& linTrack,
0127                            const Cache<nDimVertex>& cache) {
0128   constexpr unsigned int nBoundParams = nDimVertex + 2;
0129   using ParameterVector = ActsVector<nBoundParams>;
0130   using ParameterMatrix = ActsSquareMatrix<nBoundParams>;
0131   // A_k
0132   const ActsMatrix<nBoundParams, nDimVertex> posJac =
0133       linTrack.positionJacobian.block<nBoundParams, nDimVertex>(0, 0);
0134   // B_k
0135   const ActsMatrix<nBoundParams, 3> momJac =
0136       linTrack.momentumJacobian.block<nBoundParams, 3>(0, 0);
0137   // p_k
0138   const ParameterVector trkParams =
0139       linTrack.parametersAtPCA.head<nBoundParams>();
0140   // c_k
0141   const ParameterVector constTerm = linTrack.constantTerm.head<nBoundParams>();
0142   // TODO we could use `linTrack.weightAtPCA` but only if we would always fit
0143   // time.
0144   // G_k
0145   const ParameterMatrix trkParamWeight =
0146       linTrack.covarianceAtPCA.block<nBoundParams, nBoundParams>(0, 0)
0147           .inverse();
0148 
0149   // A_k * \tilde{x_k}
0150   const ParameterVector posJacVtxPos = posJac * cache.newVertexPos;
0151 
0152   // \tilde{q_k}
0153   Vector3 newTrkMom = cache.wMat * momJac.transpose() * trkParamWeight *
0154                       (trkParams - constTerm - posJacVtxPos);
0155 
0156   // \tilde{p_k}
0157   ParameterVector linearizedTrackParameters =
0158       constTerm + posJacVtxPos + momJac * newTrkMom;
0159 
0160   // r_k
0161   ParameterVector paramDiff = trkParams - linearizedTrackParameters;
0162 
0163   // Return chi2
0164   return paramDiff.transpose() * (trkParamWeight * paramDiff);
0165 }
0166 
0167 /// @brief Calculates a covariance matrix for the refitted track parameters
0168 ///
0169 /// @tparam nDimVertex number of dimensions of the vertex. Can be 3 (if we only
0170 /// fit its spatial coordinates) or 4 (if we also fit time).
0171 ///
0172 /// @param wMat W_k matrix from Ref. (1)
0173 /// @param crossCovVP Cross-covariance matrix between vertex position and track
0174 /// momentum
0175 /// @param vtxCov Vertex covariance matrix
0176 /// @param newTrkParams Refitted track parameters
0177 template <unsigned int nDimVertex>
0178 Acts::BoundMatrix calculateTrackCovariance(
0179     const SquareMatrix3& wMat, const ActsMatrix<nDimVertex, 3>& crossCovVP,
0180     const ActsSquareMatrix<nDimVertex>& vtxCov,
0181     const BoundVector& newTrkParams) {
0182   // D_k^n
0183   ActsSquareMatrix<3> momCov =
0184       wMat + crossCovVP.transpose() * vtxCov.inverse() * crossCovVP;
0185 
0186   // Full x, y, z, phi, theta, q/p, and, optionally, t covariance matrix. Note
0187   // that we call this set of parameters "free" in the following even though
0188   // that is not quite correct (free parameters correspond to
0189   // x, y, z, t, px, py, pz)
0190   constexpr unsigned int nFreeParams = nDimVertex + 3;
0191   ActsSquareMatrix<nFreeParams> freeTrkCov(
0192       ActsSquareMatrix<nFreeParams>::Zero());
0193 
0194   freeTrkCov.template block<3, 3>(0, 0) = vtxCov.template block<3, 3>(0, 0);
0195   freeTrkCov.template block<3, 3>(0, 3) = crossCovVP.template block<3, 3>(0, 0);
0196   freeTrkCov.template block<3, 3>(3, 0) =
0197       crossCovVP.template block<3, 3>(0, 0).transpose();
0198   freeTrkCov.template block<3, 3>(3, 3) = momCov;
0199 
0200   // Fill time (cross-)covariances
0201   if constexpr (nFreeParams == 7) {
0202     freeTrkCov.template block<3, 1>(0, 6) = vtxCov.template block<3, 1>(0, 3);
0203     freeTrkCov.template block<3, 1>(3, 6) =
0204         crossCovVP.template block<1, 3>(3, 0).transpose();
0205     freeTrkCov.template block<1, 3>(6, 0) = vtxCov.template block<1, 3>(3, 0);
0206     freeTrkCov.template block<1, 3>(6, 3) =
0207         crossCovVP.template block<1, 3>(3, 0);
0208     freeTrkCov(6, 6) = vtxCov(3, 3);
0209   }
0210 
0211   // Jacobian relating "free" and bound track parameters
0212   constexpr unsigned int nBoundParams = nDimVertex + 2;
0213   ActsMatrix<nBoundParams, nFreeParams> freeToBoundJac(
0214       ActsMatrix<nBoundParams, nFreeParams>::Zero());
0215 
0216   // TODO: Jacobian is not quite correct
0217   // First row
0218   freeToBoundJac(0, 0) = -std::sin(newTrkParams[2]);
0219   freeToBoundJac(0, 1) = std::cos(newTrkParams[2]);
0220 
0221   double tanTheta = std::tan(newTrkParams[3]);
0222 
0223   // Second row
0224   freeToBoundJac(1, 0) = -freeToBoundJac(0, 1) / tanTheta;
0225   freeToBoundJac(1, 1) = freeToBoundJac(0, 0) / tanTheta;
0226 
0227   // Dimension of the part of the Jacobian that is an identity matrix
0228   constexpr unsigned int nDimIdentity = nFreeParams - 2;
0229   freeToBoundJac.template block<nDimIdentity, nDimIdentity>(1, 2) =
0230       ActsMatrix<nDimIdentity, nDimIdentity>::Identity();
0231 
0232   // Full perigee track covariance
0233   BoundMatrix boundTrackCov(BoundMatrix::Identity());
0234   boundTrackCov.block<nBoundParams, nBoundParams>(0, 0) =
0235       (freeToBoundJac * (freeTrkCov * freeToBoundJac.transpose()));
0236 
0237   return boundTrackCov;
0238 }
0239 
0240 template <unsigned int nDimVertex>
0241 void updateVertexWithTrackImpl(Vertex& vtx, TrackAtVertex& trk, int sign) {
0242   if constexpr (nDimVertex != 3 && nDimVertex != 4) {
0243     throw std::invalid_argument(
0244         "The vertex dimension must either be 3 (when fitting the spatial "
0245         "coordinates) or 4 (when fitting the spatial coordinates + time).");
0246   }
0247 
0248   double trackWeight = trk.trackWeight;
0249 
0250   // Set up cache where entire content is set to 0
0251   Cache<nDimVertex> cache;
0252 
0253   // Calculate update and save result in cache
0254   calculateUpdate(vtx, trk.linearizedState, trackWeight, sign, cache);
0255 
0256   // Get fit quality parameters wrt to old vertex
0257   double chi2 = 0.;
0258   double ndf = 0.;
0259   std::tie(chi2, ndf) = vtx.fitQuality();
0260 
0261   // Chi2 of the track parameters
0262   double trkChi2 = trackParametersChi2(trk.linearizedState, cache);
0263 
0264   // Update of the chi2 of the vertex position
0265   double vtxPosChi2Update = vertexPositionChi2Update(vtx.fullPosition(), cache);
0266 
0267   // Calculate new chi2
0268   chi2 += sign * (vtxPosChi2Update + trackWeight * trkChi2);
0269 
0270   // Calculate ndf
0271   ndf += sign * trackWeight * 2.;
0272 
0273   // Updating the vertex
0274   if constexpr (nDimVertex == 3) {
0275     vtx.fullPosition().head<3>() = cache.newVertexPos.template head<3>();
0276     vtx.fullCovariance().template topLeftCorner<3, 3>() =
0277         cache.newVertexCov.template topLeftCorner<3, 3>();
0278   } else if constexpr (nDimVertex == 4) {
0279     vtx.fullPosition() = cache.newVertexPos;
0280     vtx.fullCovariance() = cache.newVertexCov;
0281   }
0282   vtx.setFitQuality(chi2, ndf);
0283 
0284   if (sign == 1) {
0285     // Update track
0286     trk.chi2Track = trkChi2;
0287     trk.ndf = 2 * trackWeight;
0288   }
0289   // Remove trk from current vertex by setting its weight to 0
0290   else if (sign == -1) {
0291     trk.trackWeight = 0.;
0292   } else {
0293     throw std::invalid_argument(
0294         "Sign for adding/removing track must be +1 (add) or -1 (remove).");
0295   }
0296 }
0297 
0298 template <unsigned int nDimVertex>
0299 void updateTrackWithVertexImpl(TrackAtVertex& track, const Vertex& vtx) {
0300   if constexpr (nDimVertex != 3 && nDimVertex != 4) {
0301     throw std::invalid_argument(
0302         "The vertex dimension must either be 3 (when fitting the spatial "
0303         "coordinates) or 4 (when fitting the spatial coordinates + time).");
0304   }
0305 
0306   using VertexVector = ActsVector<nDimVertex>;
0307   using VertexMatrix = ActsSquareMatrix<nDimVertex>;
0308   constexpr unsigned int nBoundParams = nDimVertex + 2;
0309   using ParameterVector = ActsVector<nBoundParams>;
0310   using ParameterMatrix = ActsSquareMatrix<nBoundParams>;
0311   // Check if linearized state exists
0312   if (!track.isLinearized) {
0313     throw std::invalid_argument("TrackAtVertex object must be linearized.");
0314   }
0315 
0316   // Extract vertex position and covariance
0317   // \tilde{x_n}
0318   const VertexVector vtxPos = vtx.fullPosition().template head<nDimVertex>();
0319   // C_n
0320   const VertexMatrix vtxCov =
0321       vtx.fullCovariance().template block<nDimVertex, nDimVertex>(0, 0);
0322 
0323   // Get the linearized track
0324   const LinearizedTrack& linTrack = track.linearizedState;
0325   // Retrieve variables from the track linearization. The comments indicate the
0326   // corresponding symbol used in Ref. (1).
0327   // A_k
0328   const ActsMatrix<nBoundParams, nDimVertex> posJac =
0329       linTrack.positionJacobian.block<nBoundParams, nDimVertex>(0, 0);
0330   // B_k
0331   const ActsMatrix<nBoundParams, 3> momJac =
0332       linTrack.momentumJacobian.block<nBoundParams, 3>(0, 0);
0333   // p_k
0334   const ParameterVector trkParams =
0335       linTrack.parametersAtPCA.head<nBoundParams>();
0336   // c_k
0337   const ParameterVector constTerm = linTrack.constantTerm.head<nBoundParams>();
0338   // TODO we could use `linTrack.weightAtPCA` but only if we would always fit
0339   // time.
0340   // G_k
0341   const ParameterMatrix trkParamWeight =
0342       linTrack.covarianceAtPCA.block<nBoundParams, nBoundParams>(0, 0)
0343           .inverse();
0344 
0345   // Cache object filled with zeros
0346   Cache<nDimVertex> cache;
0347 
0348   // Calculate the update of the vertex position when the track is removed. This
0349   // might be unintuitive, but it is needed to compute a symmetric chi2.
0350   calculateUpdate(vtx, linTrack, track.trackWeight, -1, cache);
0351 
0352   // Refit track momentum with the final vertex position
0353   Vector3 newTrkMomentum = cache.wMat * momJac.transpose() * trkParamWeight *
0354                            (trkParams - constTerm - posJac * vtxPos);
0355 
0356   // Track parameters, impact parameters are set to 0 and momentum corresponds
0357   // to newTrkMomentum. TODO: Make transition fitterParams -> fittedMomentum.
0358   BoundVector newTrkParams(BoundVector::Zero());
0359 
0360   // Get phi and theta and correct for possible periodicity changes
0361   const auto correctedPhiTheta =
0362       Acts::detail::normalizePhiTheta(newTrkMomentum(0), newTrkMomentum(1));
0363   newTrkParams(BoundIndices::eBoundPhi) = correctedPhiTheta.first;     // phi
0364   newTrkParams(BoundIndices::eBoundTheta) = correctedPhiTheta.second;  // theta
0365   newTrkParams(BoundIndices::eBoundQOverP) = newTrkMomentum(2);        // qOverP
0366 
0367   // E_k^n
0368   const ActsMatrix<nDimVertex, 3> crossCovVP =
0369       -vtxCov * posJac.transpose() * trkParamWeight * momJac * cache.wMat;
0370 
0371   // Difference in positions. cache.newVertexPos corresponds to \tilde{x_k^{n*}} in Ref. (1).
0372   VertexVector posDiff =
0373       vtxPos - cache.newVertexPos.template head<nDimVertex>();
0374 
0375   // r_k^n
0376   ParameterVector paramDiff =
0377       trkParams - (constTerm + posJac * vtxPos + momJac * newTrkMomentum);
0378 
0379   // New chi2 to be set later
0380   double chi2 =
0381       posDiff.dot(
0382           cache.newVertexWeight.template block<nDimVertex, nDimVertex>(0, 0) *
0383           posDiff) +
0384       paramDiff.dot(trkParamWeight * paramDiff);
0385 
0386   Acts::BoundMatrix newTrackCov = calculateTrackCovariance<nDimVertex>(
0387       cache.wMat, crossCovVP, vtxCov, newTrkParams);
0388 
0389   // Create new refitted parameters
0390   std::shared_ptr<PerigeeSurface> perigeeSurface =
0391       Surface::makeShared<PerigeeSurface>(vtxPos.template head<3>());
0392 
0393   BoundTrackParameters refittedPerigee =
0394       BoundTrackParameters(perigeeSurface, newTrkParams, std::move(newTrackCov),
0395                            track.fittedParams.particleHypothesis());
0396 
0397   // Set new properties
0398   track.fittedParams = refittedPerigee;
0399   track.chi2Track = chi2;
0400   track.ndf = 2 * track.trackWeight;
0401 
0402   return;
0403 }
0404 
0405 }  // namespace Acts::KalmanVertexUpdater::detail