feat: template KalmanVertexUpdater to optionally fit time

Core/include/Acts/Vertexing/AdaptiveMultiVertexFitter.ipp

@@ -307,9 +307,16 @@ Acts::Result<void> Acts::
           trkAtVtx.linearizedState = *result;
           trkAtVtx.isLinearized = true;
         }
-        // Update the vertex with the new track
-        KalmanVertexUpdater::updateVertexWithTrack<input_track_t>(*vtx,
-                                                                  trkAtVtx);
+        // Update the vertex with the new track. The second template argument
+        // corresponds to the number of fitted vertex dimensions (i.e., 3 if we
+        // only fit spatial coordinates and 4 if we also fit time).
+        if (m_cfg.useTime) {
+          KalmanVertexUpdater::updateVertexWithTrack<input_track_t, 4>(
+              *vtx, trkAtVtx);
+        } else {
+          KalmanVertexUpdater::updateVertexWithTrack<input_track_t, 3>(
+              *vtx, trkAtVtx);
+        }
       } else {
         ACTS_VERBOSE("Track weight too low. Skip track.");
       }
@@ -369,7 +376,15 @@ void Acts::AdaptiveMultiVertexFitter<
     for (const auto trk : state.vtxInfoMap[vtx].trackLinks) {
       auto& trkAtVtx = state.tracksAtVerticesMap.at(std::make_pair(trk, vtx));
       if (trkAtVtx.trackWeight > m_cfg.minWeight) {
-        KalmanVertexTrackUpdater::update<input_track_t>(trkAtVtx, *vtx);
+        // Update the new track under the assumption that it originates at the
+        // vertex. The second template argument corresponds to the number of
+        // fitted vertex dimensions (i.e., 3 if we only fit spatial coordinates
+        // and 4 if we also fit time).
+        if (m_cfg.useTime) {
+          KalmanVertexTrackUpdater::update<input_track_t, 4>(trkAtVtx, *vtx);
+        } else {
+          KalmanVertexTrackUpdater::update<input_track_t, 3>(trkAtVtx, *vtx);
+        }
       }
     }
   }

Core/include/Acts/Vertexing/KalmanVertexTrackUpdater.hpp

@@ -18,29 +18,44 @@ namespace Acts {
 namespace KalmanVertexTrackUpdater {
 
 /// KalmanVertexTrackUpdater
-///
+/// Based on
+/// Ref. (1):
+/// R. Frühwirth et al.
+/// Vertex reconstruction and track bundling at the lep collider using robust
+/// algorithms
+/// Computer Physics Comm.: 96 (1996) 189
+/// Chapter 2.1
+
 /// @brief Refits a single track with the knowledge of
 /// the vertex it has originated from
 ///
+/// @tparam input_track_t Track object type
+/// @tparam nDimVertex number of dimensions of the vertex. Can be 3 (if we only
+/// fit its spatial coordinates) or 4 (if we also fit time).
+///
 /// @param track Track to update
 /// @param vtx Vertex `track` belongs to
-template <typename input_track_t>
+template <typename input_track_t, unsigned int nDimVertex>
 void update(TrackAtVertex<input_track_t>& track,
             const Vertex<input_track_t>& vtx);
 
 namespace detail {
 
-/// @brief reates a new covariance matrix for the
-/// refitted track parameters
+/// @brief Calculates a covariance matrix for the refitted track parameters
+///
+/// @tparam nDimVertex number of dimensions of the vertex. Can be 3 (if we only
+/// fit its spatial coordinates) or 4 (if we also fit time).
 ///
-/// @param sMat Track ovariance in momentum space
-/// @param newTrkCov New track covariance matrixs
+/// @param wMat W_k matrix from Ref. (1)
+/// @param crossCovVP Cross-covariance matrix between vertex position and track
+/// momentum
 /// @param vtxCov Vertex covariance matrix
-/// @param newTrkParams New track parameter
-inline BoundMatrix createFullTrackCovariance(const SquareMatrix3& wMat,
-                                             const SquareMatrix3& crossCovVP,
-                                             const SquareMatrix3& vtxCov,
-                                             const BoundVector& newTrkParams);
+/// @param newTrkParams Refitted track parameters
+template <unsigned int nDimVertex>
+inline BoundMatrix calculateTrackCovariance(
+    const SquareMatrix3& wMat, const ActsMatrix<nDimVertex, 3>& crossCovVP,
+    const ActsSquareMatrix<nDimVertex>& vtxCov,
+    const BoundVector& newTrkParams);
 
 }  // Namespace detail
 

Core/include/Acts/Vertexing/KalmanVertexTrackUpdater.ipp

@@ -11,40 +11,56 @@
 #include "Acts/Utilities/detail/periodic.hpp"
 #include "Acts/Vertexing/VertexingError.hpp"
 
-template <typename input_track_t>
+template <typename input_track_t, unsigned int nDimVertex>
 void Acts::KalmanVertexTrackUpdater::update(TrackAtVertex<input_track_t>& track,
                                             const Vertex<input_track_t>& vtx) {
+  if constexpr (nDimVertex != 3 && nDimVertex != 4) {
+    throw std::invalid_argument(
+        "The vertex dimension must either be 3 (when fitting the spatial "
+        "coordinates) or 4 (when fitting the spatial coordinates + time).");
+  }
+
+  using VertexVector = ActsVector<nDimVertex>;
+  using VertexMatrix = ActsSquareMatrix<nDimVertex>;
+  constexpr unsigned int nBoundParams = nDimVertex + 2;
+  using ParameterVector = ActsVector<nBoundParams>;
+  using ParameterMatrix = ActsSquareMatrix<nBoundParams>;
   // Check if linearized state exists
   if (!track.isLinearized) {
     throw std::invalid_argument("TrackAtVertex object must be linearized.");
   }
 
   // Extract vertex position and covariance
   // \tilde{x_n}
-  const Vector3& vtxPos = vtx.position();
+  const VertexVector vtxPos = vtx.fullPosition().template head<nDimVertex>();
   // C_n
-  const SquareMatrix3& vtxCov = vtx.covariance();
+  const VertexMatrix vtxCov =
+      vtx.fullCovariance().template block<nDimVertex, nDimVertex>(0, 0);
 
   // Get the linearized track
   const LinearizedTrack& linTrack = track.linearizedState;
-
   // Retrieve variables from the track linearization. The comments indicate the
-  // corresponding symbol used in the Ref. (1).
+  // corresponding symbol used in Ref. (1).
   // A_k
-  const ActsMatrix<5, 3> posJac = linTrack.positionJacobian.block<5, 3>(0, 0);
+  const ActsMatrix<nBoundParams, nDimVertex> posJac =
+      linTrack.positionJacobian.block<nBoundParams, nDimVertex>(0, 0);
   // B_k
-  const ActsMatrix<5, 3> momJac = linTrack.momentumJacobian.block<5, 3>(0, 0);
+  const ActsMatrix<nBoundParams, 3> momJac =
+      linTrack.momentumJacobian.block<nBoundParams, 3>(0, 0);
   // p_k
-  const ActsVector<5> trkParams = linTrack.parametersAtPCA.head<5>();
-  // TODO we could use `linTrack.weightAtPCA` but only if we would use time
-  // G_k
-  const ActsSquareMatrix<5> trkParamWeight =
-      linTrack.covarianceAtPCA.block<5, 5>(0, 0).inverse();
+  const ParameterVector trkParams =
+      linTrack.parametersAtPCA.head<nBoundParams>();
   // c_k
-  const ActsVector<5> constTerm = linTrack.constantTerm.head<5>();
+  const ParameterVector constTerm = linTrack.constantTerm.head<nBoundParams>();
+  // TODO we could use `linTrack.weightAtPCA` but only if we would always fit
+  // time.
+  // G_k
+  const ParameterMatrix trkParamWeight =
+      linTrack.covarianceAtPCA.block<nBoundParams, nBoundParams>(0, 0)
+          .inverse();
 
   // Cache object filled with zeros
-  KalmanVertexUpdater::Cache cache;
+  KalmanVertexUpdater::Cache<nDimVertex> cache;
 
   // Calculate the update of the vertex position when the track is removed. This
   // might be unintuitive, but it is needed to compute a symmetric chi2.
@@ -67,30 +83,34 @@ void Acts::KalmanVertexTrackUpdater::update(TrackAtVertex<input_track_t>& track,
   newTrkParams(BoundIndices::eBoundQOverP) = newTrkMomentum(2);        // qOverP
 
   // E_k^n
-  const SquareMatrix3 crossCovVP =
+  const ActsMatrix<nDimVertex, 3> crossCovVP =
       -vtxCov * posJac.transpose() * trkParamWeight * momJac * cache.wMat;
 
   // Difference in positions. cache.newVertexPos corresponds to \tilde{x_k^{n*}} in Ref. (1).
-  Vector3 posDiff = vtxPos - cache.newVertexPos;
+  VertexVector posDiff =
+      vtxPos - cache.newVertexPos.template head<nDimVertex>();
 
   // r_k^n
-  ActsVector<5> paramDiff =
+  ParameterVector paramDiff =
       trkParams - (constTerm + posJac * vtxPos + momJac * newTrkMomentum);
 
   // New chi2 to be set later
-  double chi2 = posDiff.dot(cache.newVertexWeight * posDiff) +
-                paramDiff.dot(trkParamWeight * paramDiff);
+  double chi2 =
+      posDiff.dot(
+          cache.newVertexWeight.template block<nDimVertex, nDimVertex>(0, 0) *
+          posDiff) +
+      paramDiff.dot(trkParamWeight * paramDiff);
 
-  Acts::BoundMatrix fullPerTrackCov = detail::createFullTrackCovariance(
+  Acts::BoundMatrix newTrackCov = detail::calculateTrackCovariance<nDimVertex>(
       cache.wMat, crossCovVP, vtxCov, newTrkParams);
 
   // Create new refitted parameters
   std::shared_ptr<PerigeeSurface> perigeeSurface =
-      Surface::makeShared<PerigeeSurface>(vtxPos);
+      Surface::makeShared<PerigeeSurface>(vtxPos.template head<3>());
 
-  BoundTrackParameters refittedPerigee = BoundTrackParameters(
-      perigeeSurface, newTrkParams, std::move(fullPerTrackCov),
-      track.fittedParams.particleHypothesis());
+  BoundTrackParameters refittedPerigee =
+      BoundTrackParameters(perigeeSurface, newTrkParams, std::move(newTrackCov),
+                           track.fittedParams.particleHypothesis());
 
   // Set new properties
   track.fittedParams = refittedPerigee;
@@ -100,42 +120,66 @@ void Acts::KalmanVertexTrackUpdater::update(TrackAtVertex<input_track_t>& track,
   return;
 }
 
+template <unsigned int nDimVertex>
 inline Acts::BoundMatrix
-Acts::KalmanVertexTrackUpdater::detail::createFullTrackCovariance(
-    const SquareMatrix3& wMat, const SquareMatrix3& crossCovVP,
-    const SquareMatrix3& vtxCov, const BoundVector& newTrkParams) {
+Acts::KalmanVertexTrackUpdater::detail::calculateTrackCovariance(
+    const SquareMatrix3& wMat, const ActsMatrix<nDimVertex, 3>& crossCovVP,
+    const ActsSquareMatrix<nDimVertex>& vtxCov,
+    const BoundVector& newTrkParams) {
   // D_k^n
   ActsSquareMatrix<3> momCov =
       wMat + crossCovVP.transpose() * vtxCov.inverse() * crossCovVP;
 
-  // Full (x,y,z,phi, theta, q/p) covariance matrix
-  // To be made 7d again after switching to (x,y,z,phi, theta, q/p, t)
-  ActsSquareMatrix<6> fullTrkCov(ActsSquareMatrix<6>::Zero());
-
-  fullTrkCov.block<3, 3>(0, 0) = vtxCov;
-  fullTrkCov.block<3, 3>(0, 3) = crossCovVP;
-  fullTrkCov.block<3, 3>(3, 0) = crossCovVP.transpose();
-  fullTrkCov.block<3, 3>(3, 3) = momCov;
+  // Full x, y, z, phi, theta, q/p, and, optionally, t covariance matrix. Note
+  // that we call this set of parameters "free" in the following even though
+  // that is not quite correct (free parameters correspond to
+  // x, y, z, t, px, py, pz)
+  constexpr unsigned int nFreeParams = nDimVertex + 3;
+  ActsSquareMatrix<nFreeParams> freeTrkCov(
+      ActsSquareMatrix<nFreeParams>::Zero());
+
+  freeTrkCov.template block<3, 3>(0, 0) = vtxCov.template block<3, 3>(0, 0);
+  freeTrkCov.template block<3, 3>(0, 3) = crossCovVP.template block<3, 3>(0, 0);
+  freeTrkCov.template block<3, 3>(3, 0) =
+      crossCovVP.template block<3, 3>(0, 0).transpose();
+  freeTrkCov.template block<3, 3>(3, 3) = momCov;
+
+  // Fill time (cross-)covariances
+  if constexpr (nFreeParams == 7) {
+    freeTrkCov.template block<3, 1>(0, 6) = vtxCov.template block<3, 1>(0, 3);
+    freeTrkCov.template block<3, 1>(3, 6) =
+        crossCovVP.template block<1, 3>(3, 0).transpose();
+    freeTrkCov.template block<1, 3>(6, 0) = vtxCov.template block<1, 3>(3, 0);
+    freeTrkCov.template block<1, 3>(6, 3) =
+        crossCovVP.template block<1, 3>(3, 0);
+    freeTrkCov(6, 6) = vtxCov(3, 3);
+  }
 
-  // Combined track jacobian
-  ActsMatrix<5, 6> trkJac(ActsMatrix<5, 6>::Zero());
+  // Jacobian relating "free" and bound track parameters
+  constexpr unsigned int nBoundParams = nDimVertex + 2;
+  ActsMatrix<nBoundParams, nFreeParams> freeToBoundJac(
+      ActsMatrix<nBoundParams, nFreeParams>::Zero());
 
+  // TODO: Jacobian is not quite correct
   // First row
-  trkJac(0, 0) = -std::sin(newTrkParams[2]);
-  trkJac(0, 1) = std::cos(newTrkParams[2]);
+  freeToBoundJac(0, 0) = -std::sin(newTrkParams[2]);
+  freeToBoundJac(0, 1) = std::cos(newTrkParams[2]);
 
   double tanTheta = std::tan(newTrkParams[3]);
 
   // Second row
-  trkJac(1, 0) = -trkJac(0, 1) / tanTheta;
-  trkJac(1, 1) = trkJac(0, 0) / tanTheta;
+  freeToBoundJac(1, 0) = -freeToBoundJac(0, 1) / tanTheta;
+  freeToBoundJac(1, 1) = freeToBoundJac(0, 0) / tanTheta;
 
-  trkJac.block<4, 4>(1, 2) = ActsMatrix<4, 4>::Identity();
+  // Dimension of the part of the Jacobian that is an identity matrix
+  constexpr unsigned int nDimIdentity = nFreeParams - 2;
+  freeToBoundJac.template block<nDimIdentity, nDimIdentity>(1, 2) =
+      ActsMatrix<nDimIdentity, nDimIdentity>::Identity();
 
   // Full perigee track covariance
-  BoundMatrix fullPerTrackCov(BoundMatrix::Identity());
-  fullPerTrackCov.block<5, 5>(0, 0) =
-      (trkJac * (fullTrkCov * trkJac.transpose()));
+  BoundMatrix boundTrackCov(BoundMatrix::Identity());
+  boundTrackCov.block<nBoundParams, nBoundParams>(0, 0) =
+      (freeToBoundJac * (freeTrkCov * freeToBoundJac.transpose()));
 
-  return fullPerTrackCov;
+  return boundTrackCov;
 }