refactor: GainMatrixUpdater uses AnyMutableTrackState (#4890)

Blocked by:
- #4889

Author: paulgessinger

Core/include/Acts/TrackFitting/GainMatrixSmoother.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Acts/EventData/AnyTrackStateProxy.hpp"
 #include "Acts/EventData/MultiTrajectory.hpp"
 #include "Acts/Geometry/GeometryContext.hpp"
 #include "Acts/Utilities/Delegate.hpp"
@@ -45,41 +46,6 @@ class GainMatrixSmoother {
                           const Logger& logger = getDummyLogger()) const {
     (void)gctx;
 
-    using TrackStateProxy = typename traj_t::TrackStateProxy;
-
-    GetParameters filtered;
-    GetCovariance filteredCovariance;
-    GetParameters smoothed;
-    GetParameters predicted;
-    GetCovariance predictedCovariance;
-    GetCovariance smoothedCovariance;
-    GetCovariance jacobian;
-
-    filtered.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->filtered();
-    });
-    filteredCovariance.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->filteredCovariance();
-    });
-
-    smoothed.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->smoothed();
-    });
-    smoothedCovariance.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->smoothedCovariance();
-    });
-
-    predicted.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->predicted();
-    });
-    predictedCovariance.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->predictedCovariance();
-    });
-
-    jacobian.connect([](const void*, void* ts) {
-      return static_cast<TrackStateProxy*>(ts)->jacobian();
-    });
-
     ACTS_VERBOSE("Invoked GainMatrixSmoother on entry index: " << entryIndex);
 
     // For the last state: smoothed is filtered - also: switch to next
@@ -118,9 +84,8 @@ class GainMatrixSmoother {
       // ensure the track state has a smoothed component
       ts.addComponents(TrackStatePropMask::Smoothed);
 
-      if (auto res = calculate(&ts, &prev_ts, filtered, filteredCovariance,
-                               smoothed, predicted, predictedCovariance,
-                               smoothedCovariance, jacobian, logger);
+      if (auto res = calculate(AnyMutableTrackStateProxy{ts},
+                               AnyConstTrackStateProxy{prev_ts}, logger);
           !res.ok()) {
         error = res.error();
         return false;
@@ -146,23 +111,11 @@ class GainMatrixSmoother {
   /// formalism.
   ///
   /// @param ts Current track state to be smoothed
-  /// @param prev_ts Previous track state (already smoothed)
-  /// @param filtered Delegate to get filtered parameters
-  /// @param filteredCovariance Delegate to get filtered covariance
-  /// @param smoothed Delegate to get smoothed parameters
-  /// @param predicted Delegate to get predicted parameters
-  /// @param predictedCovariance Delegate to get predicted covariance
-  /// @param smoothedCovariance Delegate to get smoothed covariance
-  /// @param jacobian Delegate to get Jacobian matrix
+  /// @param prev_ts Previous track state (in forward direction)
   /// @param logger Logger for verbose output
   /// @return Success or failure of the smoothing calculation
-  Result<void> calculate(void* ts, void* prev_ts, const GetParameters& filtered,
-                         const GetCovariance& filteredCovariance,
-                         const GetParameters& smoothed,
-                         const GetParameters& predicted,
-                         const GetCovariance& predictedCovariance,
-                         const GetCovariance& smoothedCovariance,
-                         const GetCovariance& jacobian,
+  Result<void> calculate(AnyMutableTrackStateProxy ts,
+                         AnyConstTrackStateProxy prev_ts,
                          const Logger& logger) const;
 };
 

Core/include/Acts/TrackFitting/GainMatrixUpdater.hpp

@@ -8,7 +8,7 @@
 
 #pragma once
 
-#include "Acts/EventData/MultiTrajectory.hpp"
+#include "Acts/EventData/AnyTrackStateProxy.hpp"
 #include "Acts/EventData/Types.hpp"
 #include "Acts/Geometry/GeometryContext.hpp"
 #include "Acts/Utilities/Logger.hpp"
@@ -23,20 +23,6 @@ namespace Acts {
 /// Kalman update step using the gain matrix formalism.
 /// @ingroup track_fitting
 class GainMatrixUpdater {
-  struct InternalTrackState {
-    unsigned int calibratedSize;
-    // This is used to build a covariance matrix view in the .cpp file
-    const double* calibrated;
-    const double* calibratedCovariance;
-    BoundSubspaceIndices projector;
-
-    TrackStateTraits<kMeasurementSizeMax, false>::Parameters predicted;
-    TrackStateTraits<kMeasurementSizeMax, false>::Covariance
-        predictedCovariance;
-    TrackStateTraits<kMeasurementSizeMax, false>::Parameters filtered;
-    TrackStateTraits<kMeasurementSizeMax, false>::Covariance filteredCovariance;
-  };
-
  public:
   /// Run the Kalman update step for a single trajectory state.
   ///
@@ -70,20 +56,8 @@ class GainMatrixUpdater {
     // auto filtered = trackState.filtered();
     // auto filteredCovariance = trackState.filteredCovariance();
 
-    auto [chi2, error] = visitMeasurement(
-        InternalTrackState{
-            trackState.calibratedSize(),
-            // Note that we pass raw pointers here which are used in the correct
-            // shape later
-            trackState.effectiveCalibrated().data(),
-            trackState.effectiveCalibratedCovariance().data(),
-            trackState.projectorSubspaceIndices(),
-            trackState.predicted(),
-            trackState.predictedCovariance(),
-            trackState.filtered(),
-            trackState.filteredCovariance(),
-        },
-        logger);
+    auto [chi2, error] =
+        visitMeasurement(AnyMutableTrackStateProxy{trackState}, logger);
 
     trackState.chi2() = chi2;
 
@@ -92,11 +66,11 @@ class GainMatrixUpdater {
 
  private:
   std::tuple<double, std::error_code> visitMeasurement(
-      InternalTrackState trackState, const Logger& logger) const;
+      AnyMutableTrackStateProxy trackState, const Logger& logger) const;
 
   template <std::size_t N>
   std::tuple<double, std::error_code> visitMeasurementImpl(
-      InternalTrackState trackState, const Logger& logger) const;
+      AnyMutableTrackStateProxy trackState, const Logger& logger) const;
 };
 
 }  // namespace Acts

Core/include/Acts/TrackFitting/MbfSmoother.hpp

@@ -9,13 +9,15 @@
 #pragma once
 
 #include "Acts/Definitions/TrackParametrization.hpp"
+#include "Acts/EventData/AnyTrackStateProxy.hpp"
 #include "Acts/EventData/MultiTrajectory.hpp"
 #include "Acts/Geometry/GeometryContext.hpp"
 #include "Acts/Utilities/Logger.hpp"
 #include "Acts/Utilities/Result.hpp"
 
 #include <cstddef>
 #include <optional>
+#include <utility>
 
 namespace Acts {
 
@@ -65,7 +67,7 @@ class MbfSmoother {
       // ensure the track state has a smoothed component
       ts.addComponents(TrackStatePropMask::Smoothed);
 
-      InternalTrackState internalTrackState(ts);
+      AnyMutableTrackStateProxy internalTrackState(ts);
 
       // Smoothe the current state
       calculateSmoothed(internalTrackState, bigLambdaHat, smallLambdaHat);
@@ -77,78 +79,31 @@ class MbfSmoother {
 
       // Update the lambdas depending on the type of track state
       if (ts.typeFlags().test(TrackStateFlag::MeasurementFlag)) {
-        visitMeasurement(internalTrackState, bigLambdaHat, smallLambdaHat);
+        visitMeasurement(AnyConstTrackStateProxy{ts}, bigLambdaHat,
+                         smallLambdaHat);
       } else {
-        visitNonMeasurement(internalTrackState, bigLambdaHat, smallLambdaHat);
+        visitNonMeasurement(std::as_const(ts).jacobian(), bigLambdaHat,
+                            smallLambdaHat);
       }
     });
 
     return Result<void>::success();
   }
 
  private:
-  /// Internal track state representation for the smoother.
-  /// @note This allows us to move parts of the implementation into the .cpp
-  struct InternalTrackState final {
-    using Jacobian =
-        typename TrackStateTraits<kMeasurementSizeMax, false>::Covariance;
-    using Parameters =
-        typename TrackStateTraits<kMeasurementSizeMax, false>::Parameters;
-    using Covariance =
-        typename TrackStateTraits<kMeasurementSizeMax, false>::Covariance;
-
-    struct Measurement final {
-      unsigned int calibratedSize{0};
-      // This is used to build a covariance matrix view in the .cpp file
-      const double* calibrated{nullptr};
-      const double* calibratedCovariance{nullptr};
-      BoundSubspaceIndices projector;
-
-      template <typename TrackStateProxy>
-      explicit Measurement(TrackStateProxy ts)
-          : calibratedSize(ts.calibratedSize()),
-            calibrated(ts.effectiveCalibrated().data()),
-            calibratedCovariance(ts.effectiveCalibratedCovariance().data()),
-            projector(ts.projectorSubspaceIndices()) {}
-    };
-
-    Jacobian jacobian;
-
-    Parameters predicted;
-    Covariance predictedCovariance;
-    Parameters filtered;
-    Covariance filteredCovariance;
-    Parameters smoothed;
-    Covariance smoothedCovariance;
-
-    std::optional<Measurement> measurement;
-
-    template <typename TrackStateProxy>
-    explicit InternalTrackState(TrackStateProxy ts)
-        : jacobian(ts.jacobian()),
-          predicted(ts.predicted()),
-          predictedCovariance(ts.predictedCovariance()),
-          filtered(ts.filtered()),
-          filteredCovariance(ts.filteredCovariance()),
-          smoothed(ts.smoothed()),
-          smoothedCovariance(ts.smoothedCovariance()),
-          measurement(ts.typeFlags().test(TrackStateFlag::MeasurementFlag)
-                          ? std::optional<Measurement>(ts)
-                          : std::nullopt) {}
-  };
-
   /// Calculate the smoothed parameters and covariance.
-  void calculateSmoothed(InternalTrackState& ts,
+  void calculateSmoothed(AnyMutableTrackStateProxy& ts,
                          const BoundMatrix& bigLambdaHat,
                          const BoundVector& smallLambdaHat) const;
 
   /// Visit a non-measurement track state and update the lambdas.
-  void visitNonMeasurement(const InternalTrackState& ts,
-                           BoundMatrix& bigLambdaHat,
-                           BoundVector& smallLambdaHat) const;
+  void visitNonMeasurement(
+      const AnyConstTrackStateProxy::ConstCovarianceMap& jacobian,
+      BoundMatrix& bigLambdaHat, BoundVector& smallLambdaHat) const;
 
   /// Visit a measurement track state and update the lambdas.
-  void visitMeasurement(const InternalTrackState& ts, BoundMatrix& bigLambdaHat,
+  void visitMeasurement(const AnyConstTrackStateProxy& ts,
+                        BoundMatrix& bigLambdaHat,
                         BoundVector& smallLambdaHat) const;
 };
 

Core/include/Acts/TrackFitting/detail/GainMatrixUpdaterImpl.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Acts/EventData/AnyTrackStateProxy.hpp"
 #include "Acts/EventData/TrackParameterHelpers.hpp"
 #include "Acts/TrackFitting/GainMatrixUpdater.hpp"
 #include "Acts/TrackFitting/KalmanFitterError.hpp"
@@ -20,25 +21,24 @@ namespace Acts {
 
 template <std::size_t N>
 std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurementImpl(
-    InternalTrackState trackState, const Logger& logger) const {
+    AnyMutableTrackStateProxy trackState, const Logger& logger) const {
   double chi2 = 0;
 
   constexpr std::size_t kMeasurementSize = N;
   using ParametersVector = ActsVector<kMeasurementSize>;
   using CovarianceMatrix = ActsSquareMatrix<kMeasurementSize>;
 
-  typename TrackStateTraits<kMeasurementSize, true>::Calibrated calibrated{
-      trackState.calibrated};
-  typename TrackStateTraits<kMeasurementSize, true>::CalibratedCovariance
-      calibratedCovariance{trackState.calibratedCovariance};
+  auto calibrated = trackState.calibrated<kMeasurementSize>();
+  auto calibratedCovariance =
+      trackState.calibratedCovariance<kMeasurementSize>();
 
   ACTS_VERBOSE("Measurement dimension: " << kMeasurementSize);
   ACTS_VERBOSE("Calibrated measurement: " << calibrated.transpose());
   ACTS_VERBOSE("Calibrated measurement covariance:\n" << calibratedCovariance);
 
-  std::span<const std::uint8_t, kMeasurementSize> validSubspaceIndices(
-      trackState.projector.begin(),
-      trackState.projector.begin() + kMeasurementSize);
+  const auto validSubspaceIndices =
+      trackState.template projectorSubspaceIndices<kMeasurementSize>();
+
   FixedBoundSubspaceHelper<kMeasurementSize> subspaceHelper(
       validSubspaceIndices);
 
@@ -47,11 +47,16 @@ std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurementImpl(
 
   ACTS_VERBOSE("Measurement projector H:\n" << H);
 
-  const auto K = (trackState.predictedCovariance * H.transpose() *
-                  (H * trackState.predictedCovariance * H.transpose() +
-                   calibratedCovariance)
-                      .inverse())
-                     .eval();
+  auto filtered = trackState.filtered();
+  auto filteredCovariance = trackState.filteredCovariance();
+  const auto predicted = trackState.predicted();
+  const auto predictedCovariance = trackState.predictedCovariance();
+
+  const auto K =
+      (predictedCovariance * H.transpose() *
+       (H * predictedCovariance * H.transpose() + calibratedCovariance)
+           .inverse())
+          .eval();
 
   ACTS_VERBOSE("Gain Matrix K:\n" << K);
 
@@ -60,17 +65,16 @@ std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurementImpl(
     return {0, KalmanFitterError::UpdateFailed};
   }
 
-  trackState.filtered =
-      trackState.predicted + K * (calibrated - H * trackState.predicted);
+  filtered = predicted + K * (calibrated - H * predicted);
   // Normalize phi and theta
-  trackState.filtered = normalizeBoundParameters(trackState.filtered);
-  trackState.filteredCovariance =
-      (BoundSquareMatrix::Identity() - K * H) * trackState.predictedCovariance;
-  ACTS_VERBOSE("Filtered parameters: " << trackState.filtered.transpose());
-  ACTS_VERBOSE("Filtered covariance:\n" << trackState.filteredCovariance);
+  filtered = normalizeBoundParameters(filtered);
+  filteredCovariance =
+      (BoundSquareMatrix::Identity() - K * H) * predictedCovariance;
+  ACTS_VERBOSE("Filtered parameters: " << filtered.transpose());
+  ACTS_VERBOSE("Filtered covariance:\n" << filteredCovariance);
 
   ParametersVector residual;
-  residual = calibrated - H * trackState.filtered;
+  residual = calibrated - H * filtered;
   ACTS_VERBOSE("Residual: " << residual.transpose());
 
   CovarianceMatrix m =
@@ -85,10 +89,10 @@ std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurementImpl(
 
 // Ensure thet the compiler does not implicitly instantiate the template
 
-#define _EXTERN(N)                                                          \
-  extern template std::tuple<double, std::error_code>                       \
-  GainMatrixUpdater::visitMeasurementImpl<N>(InternalTrackState trackState, \
-                                             const Logger& logger) const
+#define _EXTERN(N)                                    \
+  extern template std::tuple<double, std::error_code> \
+  GainMatrixUpdater::visitMeasurementImpl<N>(         \
+      AnyMutableTrackStateProxy trackState, const Logger& logger) const
 
 _EXTERN(1);
 _EXTERN(2);