refactor: Allow to configure GSF mixture reducer as delegate (#2466)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

Core/include/Acts/TrackFitting/GsfMixtureReduction.hpp

@@ -0,0 +1,20 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2021 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include "Acts/Surfaces/Surface.hpp"
+#include "Acts/TrackFitting/GsfOptions.hpp"
+
+namespace Acts {
+
+void reduceMixtureWithKLDistance(
+    std::vector<Acts::Experimental::GsfComponent> &cmpCache,
+    std::size_t maxCmpsAfterMerge, const Surface &surface);
+
+}

Core/include/Acts/TrackFitting/GsfOptions.hpp

@@ -22,6 +22,12 @@ namespace Acts {
 
 namespace Experimental {
 
+struct GsfComponent {
+  ActsScalar weight = 0;
+  BoundVector boundPars = BoundVector::Zero();
+  BoundSquareMatrix boundCov = BoundSquareMatrix::Identity();
+};
+
 namespace GsfConstants {
 constexpr std::string_view kFinalMultiComponentStateColumn =
     "gsf-final-multi-component-state";
@@ -36,14 +42,17 @@ struct GsfExtensions {
   using ConstTrackStateProxy = typename traj_t::ConstTrackStateProxy;
 
   using Calibrator =
-      Delegate<void(const GeometryContext&, const CalibrationContext&,
-                    const SourceLink&, TrackStateProxy)>;
+      Delegate<void(const GeometryContext &, const CalibrationContext &,
+                    const SourceLink &, TrackStateProxy)>;
 
-  using Updater = Delegate<Result<void>(const GeometryContext&, TrackStateProxy,
-                                        Direction, const Logger&)>;
+  using Updater = Delegate<Result<void>(
+      const GeometryContext &, TrackStateProxy, Direction, const Logger &)>;
 
   using OutlierFinder = Delegate<bool(ConstTrackStateProxy)>;
 
+  using ComponentReducer =
+      Delegate<void(std::vector<GsfComponent> &, std::size_t, const Surface &)>;
+
   /// The Calibrator is a dedicated calibration algorithm that allows
   /// to calibrate measurements using track information, this could be
   /// e.g. sagging for wires, module deformations, etc.
@@ -59,6 +68,9 @@ struct GsfExtensions {
   /// Retrieves the associated surface from a source link
   SourceLinkSurfaceAccessor surfaceAccessor;
 
+  /// Takes a vector of components and reduces its number
+  ComponentReducer mixtureReducer;
+
   /// Default constructor which connects the default void components
   GsfExtensions() {
     calibrator.template connect<&detail::voidFitterCalibrator<traj_t>>();
@@ -77,7 +89,7 @@ struct GsfOptions {
 
   PropagatorPlainOptions propagatorPlainOptions;
 
-  const Surface* referenceSurface = nullptr;
+  const Surface *referenceSurface = nullptr;
 
   std::size_t maxComponents = 4;
 

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

@@ -21,7 +21,6 @@
 #include "Acts/TrackFitting/GsfOptions.hpp"
 #include "Acts/TrackFitting/KalmanFitter.hpp"
 #include "Acts/TrackFitting/detail/GsfUtils.hpp"
-#include "Acts/TrackFitting/detail/KLMixtureReduction.hpp"
 #include "Acts/TrackFitting/detail/KalmanUpdateHelpers.hpp"
 #include "Acts/Utilities/Zip.hpp"
 
@@ -32,7 +31,7 @@
 namespace Acts {
 namespace detail {
 
-template <typename traj_t, typename component_cache_t>
+template <typename traj_t>
 struct GsfResult {
   /// The multi-trajectory which stores the graph of components
   traj_t* fittedStates{nullptr};
@@ -62,7 +61,7 @@ struct GsfResult {
   Result<void> result{Result<void>::success()};
 
   // Internal: component cache to avoid reallocation
-  std::vector<component_cache_t> componentCache;
+  std::vector<Acts::Experimental::GsfComponent> componentCache;
 };
 
 /// The actor carrying out the GSF algorithm
@@ -71,15 +70,10 @@ struct GsfActor {
   /// Enforce default construction
   GsfActor() = default;
 
-  /// Stores parameters of a gaussian component
-  struct ComponentCache {
-    ActsScalar weight = 0;
-    BoundVector boundPars = BoundVector::Zero();
-    BoundSquareMatrix boundCov = BoundSquareMatrix::Identity();
-  };
+  using ComponentCache = Acts::Experimental::GsfComponent;
 
   /// Broadcast the result_type
-  using result_type = GsfResult<traj_t, ComponentCache>;
+  using result_type = GsfResult<traj_t>;
 
   // Actor configuration
   struct Config {
@@ -303,7 +297,10 @@ struct GsfActor {
         return;
       }
 
-      reduceComponents(stepper, surface, componentCache);
+      // reduce component number
+      const auto finalCmpNumber = std::min(
+          static_cast<std::size_t>(stepper.maxComponents), m_cfg.maxComponents);
+      m_cfg.extensions.mixtureReducer(componentCache, finalCmpNumber, surface);
 
       removeLowWeightComponents(componentCache);
 
@@ -427,22 +424,6 @@ struct GsfActor {
     }
   }
 
-  template <typename stepper_t>
-  void reduceComponents(const stepper_t& stepper, const Surface& surface,
-                        std::vector<ComponentCache>& cmps) const {
-    // Final component number
-    const auto final_cmp_number = std::min(
-        static_cast<std::size_t>(stepper.maxComponents), m_cfg.maxComponents);
-
-    auto proj = [](auto& a) -> decltype(auto) { return a; };
-
-    // We must differ between surface types, since there can be different
-    // local coordinates
-    detail::angleDescriptionSwitch(surface, [&](const auto& desc) {
-      detail::reduceWithKLDistance(cmps, final_cmp_number, proj, desc);
-    });
-  }
-
   /// Remove components with low weights and renormalize from the component
   /// cache
   /// TODO This function does not expect normalized components, but this

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

@@ -1,6 +1,6 @@
 // This file is part of the Acts project.
 //
-// Copyright (C) 2021 CERN for the benefit of the Acts project
+// Copyright (C) 2023 CERN for the benefit of the Acts project
 //
 // This Source Code Form is subject to the terms of the Mozilla Public
 // License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -14,15 +14,14 @@
 #include "Acts/TrackFitting/detail/GsfUtils.hpp"
 #include "Acts/Utilities/GaussianMixtureReduction.hpp"
 
-namespace Acts {
-
-namespace detail {
+namespace Acts::detail {
 
 /// Computes the Kullback-Leibler distance between two components as shown in
 /// https://arxiv.org/abs/2001.00727v1 but ignoring the weights
 template <typename component_t, typename component_projector_t>
 auto computeSymmetricKlDivergence(const component_t &a, const component_t &b,
                                   const component_projector_t &proj) {
+  using namespace Acts;
   const auto parsA = proj(a).boundPars[eBoundQOverP];
   const auto parsB = proj(b).boundPars[eBoundQOverP];
   const auto covA = proj(a).boundCov(eBoundQOverP, eBoundQOverP);
@@ -67,7 +66,7 @@ auto mergeComponents(const component_t &a, const component_t &b,
 
 /// @brief Class representing a symmetric distance matrix
 class SymmetricKLDistanceMatrix {
-  using Array = Eigen::Array<ActsScalar, Eigen::Dynamic, 1>;
+  using Array = Eigen::Array<Acts::ActsScalar, Eigen::Dynamic, 1>;
   using Mask = Eigen::Array<bool, Eigen::Dynamic, 1>;
 
   Array m_distances;
@@ -128,7 +127,7 @@ class SymmetricKLDistanceMatrix {
   }
 
   auto minDistancePair() const {
-    ActsScalar min = std::numeric_limits<ActsScalar>::max();
+    auto min = std::numeric_limits<Acts::ActsScalar>::max();
     std::size_t idx = 0;
 
     for (auto i = 0l; i < m_distances.size(); ++i) {
@@ -172,49 +171,4 @@ class SymmetricKLDistanceMatrix {
   }
 };
 
-template <typename component_t, typename component_projector_t,
-          typename angle_desc_t>
-void reduceWithKLDistance(std::vector<component_t> &cmpCache,
-                          std::size_t maxCmpsAfterMerge,
-                          const component_projector_t &proj,
-                          const angle_desc_t &angle_desc) {
-  if (cmpCache.size() <= maxCmpsAfterMerge) {
-    return;
-  }
-
-  SymmetricKLDistanceMatrix distances(cmpCache, proj);
-
-  auto remainingComponents = cmpCache.size();
-
-  while (remainingComponents > maxCmpsAfterMerge) {
-    const auto [minI, minJ] = distances.minDistancePair();
-
-    // Set one component and compute associated distances
-    cmpCache[minI] =
-        mergeComponents(cmpCache[minI], cmpCache[minJ], proj, angle_desc);
-    distances.recomputeAssociatedDistances(minI, cmpCache, proj);
-
-    // Set weight of the other component to -1 so we can remove it later and
-    // mask its distances
-    proj(cmpCache[minJ]).weight = -1.0;
-    distances.maskAssociatedDistances(minJ);
-
-    remainingComponents--;
-  }
-
-  // Remove all components which are labeled with weight -1
-  std::sort(cmpCache.begin(), cmpCache.end(),
-            [&](const auto &a, const auto &b) {
-              return proj(a).weight < proj(b).weight;
-            });
-  cmpCache.erase(
-      std::remove_if(cmpCache.begin(), cmpCache.end(),
-                     [&](const auto &a) { return proj(a).weight == -1.0; }),
-      cmpCache.end());
-
-  assert(cmpCache.size() == maxCmpsAfterMerge && "size mismatch");
-}
-
-}  // namespace detail
-
-}  // namespace Acts
+}  // namespace Acts::detail

Core/src/TrackFitting/CMakeLists.txt

@@ -8,4 +8,5 @@ target_sources(
     GsfError.cpp
     GsfUtils.cpp
     BetheHeitlerApprox.cpp
+    GsfMixtureReduction.cpp
 )

Core/src/TrackFitting/GsfMixtureReduction.cpp

@@ -0,0 +1,69 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2023 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include "Acts/TrackFitting/GsfMixtureReduction.hpp"
+
+#include "Acts/TrackFitting/detail/SymmetricKlDistanceMatrix.hpp"
+
+template <typename proj_t, typename angle_desc_t>
+void reduceWithKLDistanceImpl(
+    std::vector<Acts::Experimental::GsfComponent> &cmpCache,
+    std::size_t maxCmpsAfterMerge, const proj_t &proj,
+    const angle_desc_t &desc) {
+  Acts::detail::SymmetricKLDistanceMatrix distances(cmpCache, proj);
+
+  auto remainingComponents = cmpCache.size();
+
+  while (remainingComponents > maxCmpsAfterMerge) {
+    const auto [minI, minJ] = distances.minDistancePair();
+
+    // Set one component and compute associated distances
+    cmpCache[minI] =
+        mergeComponents(cmpCache[minI], cmpCache[minJ], proj, desc);
+    distances.recomputeAssociatedDistances(minI, cmpCache, proj);
+
+    // Set weight of the other component to -1 so we can remove it later and
+    // mask its distances
+    proj(cmpCache[minJ]).weight = -1.0;
+    distances.maskAssociatedDistances(minJ);
+
+    remainingComponents--;
+  }
+
+  // Remove all components which are labeled with weight -1
+  std::sort(cmpCache.begin(), cmpCache.end(),
+            [&](const auto &a, const auto &b) {
+              return proj(a).weight < proj(b).weight;
+            });
+  cmpCache.erase(
+      std::remove_if(cmpCache.begin(), cmpCache.end(),
+                     [&](const auto &a) { return proj(a).weight == -1.0; }),
+      cmpCache.end());
+
+  assert(cmpCache.size() == maxCmpsAfterMerge && "size mismatch");
+}
+
+namespace Acts {
+
+void reduceMixtureWithKLDistance(
+    std::vector<Acts::Experimental::GsfComponent> &cmpCache,
+    std::size_t maxCmpsAfterMerge, const Surface &surface) {
+  if (cmpCache.size() <= maxCmpsAfterMerge) {
+    return;
+  }
+
+  auto proj = [](auto &a) -> decltype(auto) { return a; };
+
+  // We must differ between surface types, since there can be different
+  // local coordinates
+  detail::angleDescriptionSwitch(surface, [&](const auto &desc) {
+    reduceWithKLDistanceImpl(cmpCache, maxCmpsAfterMerge, proj, desc);
+  });
+}
+
+}  // namespace Acts

Examples/Algorithms/TrackFitting/src/GsfFitterFunction.cpp

@@ -21,6 +21,7 @@
 #include "Acts/Propagator/Propagator.hpp"
 #include "Acts/TrackFitting/GainMatrixUpdater.hpp"
 #include "Acts/TrackFitting/GaussianSumFitter.hpp"
+#include "Acts/TrackFitting/GsfMixtureReduction.hpp"
 #include "Acts/TrackFitting/GsfOptions.hpp"
 #include "Acts/Utilities/Delegate.hpp"
 #include "Acts/Utilities/GaussianMixtureReduction.hpp"
@@ -117,6 +118,8 @@ struct GsfFitterFunctionImpl final : public ActsExamples::TrackFitterFunction {
     gsfOptions.extensions.surfaceAccessor
         .connect<&IndexSourceLink::SurfaceAccessor::operator()>(
             &m_slSurfaceAccessor);
+    gsfOptions.extensions.mixtureReducer
+        .connect<&Acts::reduceMixtureWithKLDistance>();
 
     return gsfOptions;
   }