feat: max-weight reduction for GSF trackstates, fill smoothed states (#…

…2115)

This PR

- refactors the component merging code a bit
- allows to configure, how the multi-component states on measurement-surfaces are reduced
- fills the `MultiTrajectory` in a more standard way, similar to the KF / CKF:
  - `forward pass predicted` -> `predicted`
  - `forward pass filtered` -> `filtered`
  - `reverse pass filtered` -> `smoothed`

Core/include/Acts/Propagator/MultiEigenStepperLoop.hpp

@@ -20,9 +20,9 @@
 #include "Acts/Propagator/EigenStepper.hpp"
 #include "Acts/Propagator/EigenStepperError.hpp"
 #include "Acts/Propagator/Propagator.hpp"
+#include "Acts/Utilities/GaussianMixtureReduction.hpp"
 #include "Acts/Utilities/Intersection.hpp"
 #include "Acts/Utilities/Result.hpp"
-#include "Acts/Utilities/detail/gaussian_mixture_helpers.hpp"
 
 #include <algorithm>
 #include <cmath>
@@ -159,12 +159,6 @@ struct MaxMomentumReducerLoop {
   }
 };
 
-/// @enum FinalReductionMethod
-///
-/// Available reduction methods for the reduction in the boundState and
-/// curvilinearState member functions of the MultiEigenStepperLoop.
-enum class FinalReductionMethod { eMean, eMaxWeight };
-
 /// @brief Stepper based on the EigenStepper, but handles Multi-Component Tracks
 /// (e.g., for the GSF). Internally, this only manages a vector of
 /// EigenStepper::States. This simplifies implementation, but has several
@@ -189,7 +183,7 @@ class MultiEigenStepperLoop
 
   /// How to extract a single component state when calling .boundState() or
   /// .curvilinearState()
-  FinalReductionMethod m_finalReductionMethod = FinalReductionMethod::eMean;
+  MixtureReductionMethod m_finalReductionMethod = MixtureReductionMethod::eMean;
 
   /// The logger (used if no logger is provided by caller of methods)
   std::unique_ptr<const Acts::Logger> m_logger;
@@ -219,11 +213,11 @@ class MultiEigenStepperLoop
   static constexpr int maxComponents = std::numeric_limits<int>::max();
 
   /// Constructor from a magnetic field and a optionally provided Logger
-  MultiEigenStepperLoop(
-      std::shared_ptr<const MagneticFieldProvider> bField,
-      FinalReductionMethod finalReductionMethod = FinalReductionMethod::eMean,
-      std::unique_ptr<const Logger> logger = getDefaultLogger("GSF",
-                                                              Logging::INFO))
+  MultiEigenStepperLoop(std::shared_ptr<const MagneticFieldProvider> bField,
+                        MixtureReductionMethod finalReductionMethod =
+                            MixtureReductionMethod::eMean,
+                        std::unique_ptr<const Logger> logger =
+                            getDefaultLogger("GSF", Logging::INFO))
       : EigenStepper<extensionlist_t, auctioneer_t>(std::move(bField)),
         m_finalReductionMethod(finalReductionMethod),
         m_logger(std::move(logger)) {}

Core/include/Acts/Propagator/MultiEigenStepperLoop.ipp

@@ -58,19 +58,8 @@ auto MultiEigenStepperLoop<E, R, A>::boundState(
     return MultiStepperError::AllComponentsConversionToBoundFailed;
   }
 
-  const auto [mean, cov] =
-      detail::angleDescriptionSwitch(surface, [&](const auto& desc) {
-        return detail::combineGaussianMixture(states, Acts::Identity{}, desc);
-      });
-
-  const auto finalPars =
-      (m_finalReductionMethod == FinalReductionMethod::eMaxWeight)
-          ? std::get<BoundVector>(*std::max_element(
-                states.begin(), states.end(),
-                [](const auto& a, const auto& b) {
-                  return std::get<double>(a) < std::get<double>(b);
-                }))
-          : mean;
+  const auto [finalPars, cov] =
+      Acts::reduceGaussianMixture(states, surface, m_finalReductionMethod);
 
   std::optional<BoundSymMatrix> finalCov = std::nullopt;
   if (cov != BoundSymMatrix::Zero()) {
@@ -91,7 +80,7 @@ auto MultiEigenStepperLoop<E, R, A>::curvilinearState(State& state,
   if (numberComponents(state) == 1) {
     return SingleStepper::curvilinearState(state.components.front().state,
                                            transportCov);
-  } else if (m_finalReductionMethod == FinalReductionMethod::eMaxWeight) {
+  } else if (m_finalReductionMethod == MixtureReductionMethod::eMaxWeight) {
     auto cmpIt = std::max_element(
         state.components.begin(), state.components.end(),
         [](const auto& a, const auto& b) { return a.weight < b.weight; });

Core/include/Acts/TrackFitting/GaussianSumFitter.hpp

@@ -366,8 +366,8 @@ struct GaussianSumFitter {
 
       auto proxy =
           r.fittedStates->getTrackState(fwdGsfResult.lastMeasurementTip);
-      proxy.filtered() = proxy.predicted();
-      proxy.filteredCovariance() = proxy.predictedCovariance();
+      proxy.shareFrom(TrackStatePropMask::Filtered,
+                      TrackStatePropMask::Smoothed);
 
       r.currentTip = fwdGsfResult.lastMeasurementTip;
       r.visitedSurfaces.push_back(&proxy.referenceSurface());
@@ -426,6 +426,7 @@ struct GaussianSumFitter {
       if (not found && state.typeFlags().test(MeasurementFlag)) {
         state.typeFlags().set(OutlierFlag);
         state.typeFlags().reset(MeasurementFlag);
+        state.unset(TrackStatePropMask::Smoothed);
       }
 
       measurementStatesFinal +=

Core/include/Acts/TrackFitting/GsfOptions.hpp

@@ -11,6 +11,7 @@
 #include "Acts/EventData/MultiTrajectory.hpp"
 #include "Acts/Geometry/GeometryContext.hpp"
 #include "Acts/MagneticField/MagneticFieldContext.hpp"
+#include "Acts/Propagator/MultiEigenStepperLoop.hpp"
 #include "Acts/Propagator/Propagator.hpp"
 #include "Acts/TrackFitting/detail/VoidKalmanComponents.hpp"
 #include "Acts/Utilities/CalibrationContext.hpp"
@@ -84,6 +85,9 @@ struct GsfOptions {
 
   std::string_view finalMultiComponentStateColumn = "";
 
+  MixtureReductionMethod stateReductionMethod =
+      MixtureReductionMethod::eMaxWeight;
+
 #if __cplusplus < 202002L
   GsfOptions() = delete;
 #endif

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

@@ -107,6 +107,9 @@ struct GsfActor {
     /// be started backwards in the first pass
     bool inReversePass = false;
 
+    /// How to reduce the states that are stored in the multi trajectory
+    MixtureReductionMethod reductionMethod = MixtureReductionMethod::eMaxWeight;
+
     const Logger* logger{nullptr};
   } m_cfg;
 
@@ -720,56 +723,67 @@ struct GsfActor {
 
   void addCombinedState(result_type& result, const TemporaryStates& tmpStates,
                         const Surface& surface) const {
-    using PredProjector =
+    using PrtProjector =
         MultiTrajectoryProjector<StatesType::ePredicted, traj_t>;
-    using FiltProjector =
+    using FltProjector =
         MultiTrajectoryProjector<StatesType::eFiltered, traj_t>;
 
-    // We do not need smoothed and jacobian for now
-    const auto mask = TrackStatePropMask::Calibrated |
-                      TrackStatePropMask::Predicted |
-                      TrackStatePropMask::Filtered;
-
     if (not m_cfg.inReversePass) {
-      // The predicted state is the forward pass
-      const auto [filtMean, filtCov] =
-          angleDescriptionSwitch(surface, [&](const auto& desc) {
-            return combineGaussianMixture(
-                tmpStates.tips,
-                FiltProjector{tmpStates.traj, tmpStates.weights}, desc);
-          });
+      const auto firstCmpProxy =
+          tmpStates.traj.getTrackState(tmpStates.tips.front());
+      const auto isMeasurement =
+          firstCmpProxy.typeFlags().test(MeasurementFlag);
+
+      const auto mask =
+          isMeasurement
+              ? TrackStatePropMask::Calibrated | TrackStatePropMask::Predicted |
+                    TrackStatePropMask::Filtered | TrackStatePropMask::Smoothed
+              : TrackStatePropMask::Calibrated | TrackStatePropMask::Predicted;
 
       result.currentTip =
           result.fittedStates->addTrackState(mask, result.currentTip);
       auto proxy = result.fittedStates->getTrackState(result.currentTip);
-      auto firstCmpProxy = tmpStates.traj.getTrackState(tmpStates.tips.front());
 
       proxy.setReferenceSurface(surface.getSharedPtr());
       proxy.copyFrom(firstCmpProxy, mask);
 
-      // We set predicted & filtered the same so that the fields are not
-      // uninitialized when not finding this state in the reverse pass.
-      proxy.predicted() = filtMean;
-      proxy.predictedCovariance() = filtCov;
-      proxy.filtered() = filtMean;
-      proxy.filteredCovariance() = filtCov;
+      auto [prtMean, prtCov] = reduceGaussianMixture(
+          tmpStates.tips, surface, m_cfg.reductionMethod,
+          PrtProjector{tmpStates.traj, tmpStates.weights});
+      proxy.predicted() = prtMean;
+      proxy.predictedCovariance() = prtCov;
+
+      if (isMeasurement) {
+        auto [fltMean, fltCov] = reduceGaussianMixture(
+            tmpStates.tips, surface, m_cfg.reductionMethod,
+            FltProjector{tmpStates.traj, tmpStates.weights});
+        proxy.filtered() = fltMean;
+        proxy.filteredCovariance() = fltCov;
+        proxy.smoothed() = BoundVector::Constant(-2);
+        proxy.smoothedCovariance() = BoundSymMatrix::Constant(-2);
+      } else {
+        proxy.shareFrom(TrackStatePropMask::Predicted,
+                        TrackStatePropMask::Filtered);
+      }
+
     } else {
       assert((result.currentTip != MultiTrajectoryTraits::kInvalid &&
               "tip not valid"));
+
       result.fittedStates->applyBackwards(
           result.currentTip, [&](auto trackState) {
             auto fSurface = &trackState.referenceSurface();
             if (fSurface == &surface) {
-              const auto [filtMean, filtCov] =
-                  angleDescriptionSwitch(surface, [&](const auto& desc) {
-                    return combineGaussianMixture(
-                        tmpStates.tips,
-                        FiltProjector{tmpStates.traj, tmpStates.weights}, desc);
-                  });
-
-              trackState.filtered() = filtMean;
-              trackState.filteredCovariance() = filtCov;
               result.surfacesVisitedBwdAgain.push_back(&surface);
+
+              if (trackState.hasSmoothed()) {
+                const auto [smtMean, smtCov] = reduceGaussianMixture(
+                    tmpStates.tips, surface, m_cfg.reductionMethod,
+                    FltProjector{tmpStates.traj, tmpStates.weights});
+
+                trackState.smoothed() = smtMean;
+                trackState.smoothedCovariance() = smtCov;
+              }
               return false;
             }
             return true;
@@ -785,6 +799,7 @@ struct GsfActor {
     m_cfg.abortOnError = options.abortOnError;
     m_cfg.disableAllMaterialHandling = options.disableAllMaterialHandling;
     m_cfg.weightCutoff = options.weightCutoff;
+    m_cfg.reductionMethod = options.stateReductionMethod;
   }
 };
 

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

@@ -11,9 +11,8 @@
 #include "Acts/EventData/MultiComponentBoundTrackParameters.hpp"
 #include "Acts/EventData/MultiTrajectory.hpp"
 #include "Acts/EventData/TrackParameters.hpp"
-#include "Acts/Utilities/detail/gaussian_mixture_helpers.hpp"
-
-#include "GsfUtils.hpp"
+#include "Acts/TrackFitting/detail/GsfUtils.hpp"
+#include "Acts/Utilities/GaussianMixtureReduction.hpp"
 
 namespace Acts {
 
@@ -56,7 +55,7 @@ auto mergeComponents(const component_t &a, const component_t &b,
   };
 
   auto [mergedPars, mergedCov] =
-      combineGaussianMixture(range, refProj, angle_desc);
+      gaussianMixtureMeanCov(range, refProj, angle_desc);
 
   component_t ret = a;
   proj(ret).boundPars = mergedPars;

Core/include/Acts/Utilities/detail/gaussian_mixture_helpers.hpp → Core/include/Acts/Utilities/GaussianMixtureReduction.hpp

@@ -76,9 +76,10 @@ auto angleDescriptionSwitch(const Surface &surface, Callable &&callable) {
 
 template <int D, typename components_t, typename projector_t,
           typename angle_desc_t>
-auto combineCov(const components_t components, const ActsVector<D> &mean,
-                double sumOfWeights, projector_t &&projector,
-                const angle_desc_t &angleDesc) {
+auto gaussianMixtureCov(const components_t components,
+                        const ActsVector<D> &mean, double sumOfWeights,
+                        projector_t &&projector,
+                        const angle_desc_t &angleDesc) {
   ActsSymMatrix<D> cov = ActsSymMatrix<D>::Zero();
 
   for (const auto &cmp : components) {
@@ -124,7 +125,7 @@ auto combineCov(const components_t components, const ActsVector<D> &mean,
 /// angles in the bound parameters
 template <typename components_t, typename projector_t = Identity,
           typename angle_desc_t = AngleDescription<Surface::Plane>::Desc>
-auto combineGaussianMixture(const components_t components,
+auto gaussianMixtureMeanCov(const components_t components,
                             projector_t &&projector = projector_t{},
                             const angle_desc_t &angleDesc = angle_desc_t{}) {
   // Extract the first component
@@ -195,10 +196,48 @@ auto combineGaussianMixture(const components_t components,
   std::apply([&](auto... dsc) { (wrap(dsc), ...); }, angleDesc);
 
   const auto cov =
-      combineCov(components, mean, sumOfWeights, projector, angleDesc);
+      gaussianMixtureCov(components, mean, sumOfWeights, projector, angleDesc);
 
   return RetType{mean, cov};
 }
 
 }  // namespace detail
+
+/// @enum MixtureReductionMethod
+///
+/// Available reduction methods for the reduction of a Gaussian mixture
+enum class MixtureReductionMethod { eMean, eMaxWeight };
+
+/// Reduce Gaussian mixture
+///
+/// @param mixture The mixture iterable
+/// @param surface The surface, on which the bound state is
+/// @param method How to reduce the mixture
+/// @param projector How to project a element of the iterable to something
+/// like a std::tuple< double, BoundVector, BoundMatrix >
+///
+/// @return parameters and covariance as std::tuple< BoundVector, BoundMatrix >
+template <typename mixture_t, typename projector_t = Acts::Identity>
+auto reduceGaussianMixture(const mixture_t &mixture, const Surface &surface,
+                           MixtureReductionMethod method,
+                           projector_t &&projector = projector_t{}) {
+  using R = std::tuple<Acts::BoundVector, Acts::BoundSymMatrix>;
+  const auto [mean, cov] =
+      detail::angleDescriptionSwitch(surface, [&](const auto &desc) {
+        return detail::gaussianMixtureMeanCov(mixture, projector, desc);
+      });
+
+  if (method == MixtureReductionMethod::eMean) {
+    return R{mean, cov};
+  } else {
+    const auto maxWeightIt = std::max_element(
+        mixture.begin(), mixture.end(), [&](const auto &a, const auto &b) {
+          return std::get<0>(projector(a)) < std::get<0>(projector(b));
+        });
+    const BoundVector meanMaxWeight = std::get<1>(projector(*maxWeightIt));
+
+    return R{meanMaxWeight, cov};
+  }
+}
+
 }  // namespace Acts

Examples/Algorithms/TrackFitting/include/ActsExamples/TrackFitting/TrackFitterFunction.hpp

@@ -88,7 +88,7 @@ std::shared_ptr<TrackFitterFunction> makeGsfFitterFunction(
     std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
     std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
     BetheHeitlerApprox betheHeitlerApprox, std::size_t maxComponents,
-    double weightCutoff, Acts::FinalReductionMethod finalReductionMethod,
+    double weightCutoff, Acts::MixtureReductionMethod finalReductionMethod,
     bool abortOnError, bool disableAllMaterialHandling,
     const Acts::Logger& logger);
 

Examples/Algorithms/TrackFitting/src/GsfFitterFunction.cpp

@@ -54,6 +54,8 @@ struct GsfFitterFunctionImpl final : public ActsExamples::TrackFitterFunction {
   double weightCutoff = 0;
   bool abortOnError = false;
   bool disableAllMaterialHandling = false;
+  Acts::MixtureReductionMethod reductionMethod =
+      Acts::MixtureReductionMethod::eMaxWeight;
 
   GsfFitterFunctionImpl(Fitter&& f, DirectFitter&& df)
       : fitter(std::move(f)), directFitter(std::move(df)) {}
@@ -130,7 +132,7 @@ std::shared_ptr<TrackFitterFunction> ActsExamples::makeGsfFitterFunction(
     std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
     std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
     BetheHeitlerApprox betheHeitlerApprox, std::size_t maxComponents,
-    double weightCutoff, Acts::FinalReductionMethod finalReductionMethod,
+    double weightCutoff, Acts::MixtureReductionMethod finalReductionMethod,
     bool abortOnError, bool disableAllMaterialHandling,
     const Acts::Logger& logger) {
   // Standard fitter
@@ -166,6 +168,7 @@ std::shared_ptr<TrackFitterFunction> ActsExamples::makeGsfFitterFunction(
   fitterFunction->weightCutoff = weightCutoff;
   fitterFunction->abortOnError = abortOnError;
   fitterFunction->disableAllMaterialHandling = disableAllMaterialHandling;
+  fitterFunction->reductionMethod = finalReductionMethod;
 
   return fitterFunction;
 }

Examples/Python/src/TrackFitting.cpp

@@ -75,9 +75,9 @@ void addTrackFitting(Context& ctx) {
               return std::make_shared<PassThroughCalibrator>();
             });
 
-    py::enum_<Acts::FinalReductionMethod>(mex, "FinalReductionMethod")
-        .value("mean", Acts::FinalReductionMethod::eMean)
-        .value("maxWeight", Acts::FinalReductionMethod::eMaxWeight);
+    py::enum_<Acts::MixtureReductionMethod>(mex, "FinalReductionMethod")
+        .value("mean", Acts::MixtureReductionMethod::eMean)
+        .value("maxWeight", Acts::MixtureReductionMethod::eMaxWeight);
 
     py::class_<ActsExamples::BetheHeitlerApprox>(mex, "AtlasBetheHeitlerApprox")
         .def_static("loadFromFiles",
@@ -92,9 +92,9 @@ void addTrackFitting(Context& ctx) {
         [](std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
            std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
            BetheHeitlerApprox betheHeitlerApprox, std::size_t maxComponents,
-           double weightCutoff, Acts::FinalReductionMethod finalReductionMethod,
-           bool abortOnError, bool disableAllMaterialHandling,
-           Logging::Level level) {
+           double weightCutoff,
+           Acts::MixtureReductionMethod finalReductionMethod, bool abortOnError,
+           bool disableAllMaterialHandling, Logging::Level level) {
           return ActsExamples::makeGsfFitterFunction(
               trackingGeometry, magneticField, betheHeitlerApprox,
               maxComponents, weightCutoff, finalReductionMethod, abortOnError,