feat: CKF brem recovery

Core/include/Acts/TrackFinding/CombinatorialKalmanFilterExtensions.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Acts/TrackFitting/GsfComponent.hpp"
 #include "Acts/TrackFitting/KalmanFitter.hpp"
 #include "Acts/Utilities/Result.hpp"
 
@@ -96,6 +97,15 @@ struct CombinatorialKalmanFilterExtensions {
   ///   which makes uses of @ref MeasurementSelector and SourceLinkAccessor
   TrackStateCreator createTrackStates;
 
+  // The following options are only relevant if a multi stepper is used
+
+  /// Type alias for component reducer delegate function
+  using ComponentReducer =
+      Delegate<void(std::vector<GsfComponent>&, std::size_t, const Surface&)>;
+
+  /// Takes a vector of components and reduces its number
+  ComponentReducer mixtureReducer;
+
  private:
   /// Default branch stopper which will never stop
   /// @return false

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

@@ -281,11 +281,25 @@
       std::vector<GsfComponent>& componentCache = result.componentCache;
       componentCache.clear();
 
-      convoluteComponents(
-          state, stepper, tmpStates, *m_cfg.bethe_heitler_approx,
-          result.betheHeitlerCache, m_cfg.weightCutoff, componentCache,
-          result.nInvalidBetheHeitler.tmp(), result.maxPathXOverX0.tmp(),
-          result.sumPathXOverX0.tmp(), logger());
+      double pathXOverX0 = 0.0;
+      for (const TrackIndexType idx : tmpStates.tips) {
+        auto proxy = tmpStates.traj.getTrackState(idx);
+
+        const BoundTrackParameters bound(
+            surface.getSharedPtr(), proxy.filtered(),
+            proxy.filteredCovariance(),
+            stepper.particleHypothesis(state.stepping));
+
+        pathXOverX0 += applyBetheHeitler(
+            state.options.geoContext, surface, state.options.direction, bound,
+            tmpStates.weights.at(idx), *m_cfg.bethe_heitler_approx,
+            result.betheHeitlerCache, m_cfg.weightCutoff, componentCache,
+            result.nInvalidBetheHeitler.tmp(), result.maxPathXOverX0.tmp(),
+            logger());
+      }
+      // Store average material seen by the components
+      // Should not be too broadly distributed
+      result.sumPathXOverX0.tmp() += pathXOverX0 / tmpStates.tips.size();
 
       if (componentCache.empty()) {
         ACTS_WARNING(

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

@@ -351,39 +351,6 @@ double applyBetheHeitler(
     std::size_t &nInvalidBetheHeitler, double &maxPathXOverX0,
     const Logger &logger);
 
-template <typename traj_t, typename propagator_state_t, typename stepper_t>
-void convoluteComponents(
-    propagator_state_t &state, const stepper_t &stepper,
-    const TemporaryStates<traj_t> &tmpStates,
-    const BetheHeitlerApprox &betheHeitlerApprox,
-    std::vector<BetheHeitlerApprox::Component> &betheHeitlerCache,
-    double weightCutoff, std::vector<GsfComponent> &componentCache,
-    std::size_t &nInvalidBetheHeitler, double &maxPathXOverX0,
-    double &sumPathXOverX0, const Logger &logger) {
-  const GeometryContext &geoContext = state.options.geoContext;
-  const Direction direction = state.options.direction;
-
-  double pathXOverX0 = 0.0;
-  auto cmps = stepper.componentIterable(state.stepping);
-  for (auto [idx, cmp] : zip(tmpStates.tips, cmps)) {
-    auto proxy = tmpStates.traj.getTrackState(idx);
-    const Surface &surface = proxy.referenceSurface();
-
-    BoundTrackParameters bound(surface.getSharedPtr(), proxy.filtered(),
-                               proxy.filteredCovariance(),
-                               stepper.particleHypothesis(state.stepping));
-
-    pathXOverX0 += applyBetheHeitler(
-        geoContext, surface, direction, bound, tmpStates.weights.at(idx),
-        betheHeitlerApprox, betheHeitlerCache, weightCutoff, componentCache,
-        nInvalidBetheHeitler, maxPathXOverX0, logger);
-  }
-
-  // Store average material seen by the components
-  // Should not be too broadly distributed
-  sumPathXOverX0 += pathXOverX0 / tmpStates.tips.size();
-}
-
 /// Apply the multiple scattering to the state
 template <typename propagator_state_t, typename stepper_t>
 void applyMultipleScattering(propagator_state_t &state,

Core/src/TrackFitting/GsfUtils.cpp

@@ -113,7 +113,7 @@ double detail::Gsf::applyBetheHeitler(
       continue;
     }
 
-    if (gaussian.mean < 1.e-8) {
+    if (gaussian.mean < 1e-8) {
       ACTS_WARNING("Skip component with gaussian " << gaussian.mean << " +- "
                                                    << gaussian.var);
       continue;
@@ -122,24 +122,24 @@ double detail::Gsf::applyBetheHeitler(
     // compute delta p from mixture and update parameters
     BoundVector newPars = initialParameters.parameters();
 
-    const auto delta_p = [&]() {
+    const double deltaP = [&]() {
       if (direction == Direction::Forward()) {
-        return initialMomentum * (gaussian.mean - 1.);
+        return initialMomentum * (gaussian.mean - 1);
       } else {
-        return initialMomentum * (1. / gaussian.mean - 1.);
+        return initialMomentum * (1 / gaussian.mean - 1);
       }
     }();
 
-    assert(initialMomentum + delta_p > 0. && "new momentum must be > 0");
+    assert(initialMomentum + deltaP > 0 && "new momentum must be > 0");
     newPars[eBoundQOverP] = particleHypothesis.qOverP(
-        initialMomentum + delta_p, initialParameters.charge());
+        initialMomentum + deltaP, initialParameters.charge());
 
     // compute inverse variance of p from mixture and update covariance
     BoundMatrix newCov = initialParameters.covariance().value();
 
-    const auto varInvP = [&]() {
+    const double varInvP = [&]() {
       if (direction == Direction::Forward()) {
-        const double f = 1. / (initialMomentum * gaussian.mean);
+        const double f = 1 / (initialMomentum * gaussian.mean);
         return f * f * gaussian.var;
       } else {
         return gaussian.var / (initialMomentum * initialMomentum);

Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp

@@ -62,14 +62,17 @@ class TrackFindingAlgorithm final : public IAlgorithm {
   };
 
   /// Create the track finder function implementation.
-  ///
-  /// The magnetic field is intentionally given by-value since the variant
-  /// contains shared_ptr anyway.
   static std::shared_ptr<TrackFinderFunction> makeTrackFinderFunction(
       std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
       std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
       const Acts::Logger& logger);
 
+  /// Create the bremstrahlung track finder function implementation.
+  static std::shared_ptr<TrackFinderFunction> makeBremTrackFinderFunction(
+      std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
+      std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
+      const Acts::Logger& logger);
+
   struct Config {
     /// Input measurements collection.
     std::string inputMeasurements;
@@ -88,6 +91,8 @@ class TrackFindingAlgorithm final : public IAlgorithm {
 
     /// Type erased track finder function.
     std::shared_ptr<TrackFinderFunction> findTracks;
+    /// Type erased track finder with brem recovery function.
+    std::shared_ptr<TrackFinderFunction> findTracksBrem;
     /// CKF measurement selector config
     Acts::MeasurementSelector::Config measurementSelectorCfg;
     /// Track selector config

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp

@@ -12,6 +12,7 @@
 #include "Acts/Definitions/Direction.hpp"
 #include "Acts/Definitions/TrackParametrization.hpp"
 #include "Acts/EventData/MultiTrajectory.hpp"
+#include "Acts/EventData/ParticleHypothesis.hpp"
 #include "Acts/EventData/ProxyAccessor.hpp"
 #include "Acts/EventData/SourceLink.hpp"
 #include "Acts/EventData/TrackContainer.hpp"
@@ -26,7 +27,9 @@
 #include "Acts/Surfaces/PerigeeSurface.hpp"
 #include "Acts/Surfaces/Surface.hpp"
 #include "Acts/TrackFinding/TrackStateCreator.hpp"
+#include "Acts/TrackFitting/BetheHeitlerApprox.hpp"
 #include "Acts/TrackFitting/GainMatrixUpdater.hpp"
+#include "Acts/TrackFitting/GsfMixtureReduction.hpp"
 #include "Acts/Utilities/Enumerate.hpp"
 #include "Acts/Utilities/Logger.hpp"
 #include "Acts/Utilities/TrackHelpers.hpp"
@@ -345,6 +348,7 @@ ProcessCode TrackFindingAlgorithm::execute(const AlgorithmContext& ctx) const {
   extensions.createTrackStates
       .template connect<&TrackStateCreatorType ::createTrackStates>(
           &trackStateCreator);
+  extensions.mixtureReducer.connect<&Acts::reduceMixtureLargestWeights>();
 
   Acts::PropagatorPlainOptions firstPropOptions(ctx.geoContext,
                                                 ctx.magFieldContext);
@@ -367,12 +371,16 @@ ProcessCode TrackFindingAlgorithm::execute(const AlgorithmContext& ctx) const {
                                   firstPropOptions);
 
   firstOptions.targetSurface = m_cfg.reverseSearch ? pSurface.get() : nullptr;
+  firstOptions.betheHeitlerApprox =
+      std::make_shared<Acts::AtlasBetheHeitlerApprox>(
+          Acts::makeDefaultBetheHeitlerApprox());
 
   TrackFinderOptions secondOptions(ctx.geoContext, ctx.magFieldContext,
                                    ctx.calibContext, extensions,
                                    secondPropOptions);
   secondOptions.targetSurface = m_cfg.reverseSearch ? nullptr : pSurface.get();
   secondOptions.skipPrePropagationUpdate = true;
+  secondOptions.betheHeitlerApprox = firstOptions.betheHeitlerApprox;
 
   using Extrapolator = Acts::Propagator<Acts::SympyStepper, Acts::Navigator>;
   using ExtrapolatorOptions = Extrapolator::template Options<
@@ -481,9 +489,16 @@ ProcessCode TrackFindingAlgorithm::execute(const AlgorithmContext& ctx) const {
     ACTS_VERBOSE("Processing seed " << iSeed << " with initial parameters "
                                     << firstInitialParameters);
 
+    const TrackFinderFunction& findTracks =
+        (m_cfg.findTracksBrem == nullptr ||
+         firstInitialParameters.particleHypothesis() !=
+             Acts::ParticleHypothesis::electron())
+            ? *m_cfg.findTracks
+            : *m_cfg.findTracksBrem;
+
     auto firstRootBranch = tracksTemp.makeTrack();
-    auto firstResult = (*m_cfg.findTracks)(firstInitialParameters, firstOptions,
-                                           tracksTemp, firstRootBranch);
+    auto firstResult = findTracks(firstInitialParameters, firstOptions,
+                                  tracksTemp, firstRootBranch);
     nSeed++;
 
     if (!firstResult.ok()) {
@@ -551,9 +566,8 @@ ProcessCode TrackFindingAlgorithm::execute(const AlgorithmContext& ctx) const {
 
           auto secondRootBranch = tracksTemp.makeTrack();
           secondRootBranch.copyFromWithoutStates(trackCandidate);
-          auto secondResult =
-              (*m_cfg.findTracks)(secondInitialParameters, secondOptions,
-                                  tracksTemp, secondRootBranch);
+          auto secondResult = findTracks(secondInitialParameters, secondOptions,
+                                         tracksTemp, secondRootBranch);
 
           if (!secondResult.ok()) {
             ACTS_WARNING("Second track finding failed for seed "

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithmFunction.cpp

@@ -6,7 +6,7 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
-#include "Acts/EventData/TrackContainer.hpp"
+#include "Acts/Propagator/MultiStepperLoop.hpp"
 #include "Acts/Propagator/Navigator.hpp"
 #include "Acts/Propagator/Propagator.hpp"
 #include "Acts/Propagator/SympyStepper.hpp"
@@ -18,34 +18,39 @@
 #include <memory>
 #include <utility>
 
+namespace ActsExamples {
+
 namespace {
 
 using Stepper = Acts::SympyStepper;
 using Navigator = Acts::Navigator;
 using Propagator = Acts::Propagator<Stepper, Navigator>;
-using CKF =
-    Acts::CombinatorialKalmanFilter<Propagator, ActsExamples::TrackContainer>;
+using CKF = Acts::CombinatorialKalmanFilter<Propagator, TrackContainer>;
+
+using BremStepper = Acts::MultiStepperLoop<Stepper>;
+using BremPropagator = Acts::Propagator<BremStepper, Navigator>;
+using BremCKF = Acts::CombinatorialKalmanFilter<BremPropagator, TrackContainer>;
 
+template <typename CKFType>
 struct TrackFinderFunctionImpl
-    : public ActsExamples::TrackFindingAlgorithm::TrackFinderFunction {
-  CKF trackFinder;
+    : public TrackFindingAlgorithm::TrackFinderFunction {
+  CKFType trackFinder;
 
-  explicit TrackFinderFunctionImpl(CKF&& f) : trackFinder(std::move(f)) {}
+  explicit TrackFinderFunctionImpl(CKFType&& f) : trackFinder(std::move(f)) {}
 
-  ActsExamples::TrackFindingAlgorithm::TrackFinderResult operator()(
-      const ActsExamples::TrackParameters& initialParameters,
-      const ActsExamples::TrackFindingAlgorithm::TrackFinderOptions& options,
-      ActsExamples::TrackContainer& tracks,
-      ActsExamples::TrackProxy rootBranch) const override {
+  TrackFindingAlgorithm::TrackFinderResult operator()(
+      const TrackParameters& initialParameters,
+      const TrackFindingAlgorithm::TrackFinderOptions& options,
+      TrackContainer& tracks, TrackProxy rootBranch) const override {
     return trackFinder.findTracks(initialParameters, options, tracks,
                                   rootBranch);
   }
 };
 
 }  // namespace
 
-std::shared_ptr<ActsExamples::TrackFindingAlgorithm::TrackFinderFunction>
-ActsExamples::TrackFindingAlgorithm::makeTrackFinderFunction(
+std::shared_ptr<TrackFindingAlgorithm::TrackFinderFunction>
+TrackFindingAlgorithm::makeTrackFinderFunction(
     std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
     std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
     const Acts::Logger& logger) {
@@ -60,5 +65,28 @@ ActsExamples::TrackFindingAlgorithm::makeTrackFinderFunction(
   CKF trackFinder(std::move(propagator), logger.cloneWithSuffix("Finder"));
 
   // build the track finder functions. owns the track finder object.
-  return std::make_shared<TrackFinderFunctionImpl>(std::move(trackFinder));
+  return std::make_shared<TrackFinderFunctionImpl<CKF>>(std::move(trackFinder));
 }
+
+std::shared_ptr<TrackFindingAlgorithm::TrackFinderFunction>
+TrackFindingAlgorithm::makeBremTrackFinderFunction(
+    std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
+    std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
+    const Acts::Logger& logger) {
+  BremStepper stepper(std::move(magneticField));
+  Navigator::Config cfg{std::move(trackingGeometry)};
+  cfg.resolvePassive = false;
+  cfg.resolveMaterial = true;
+  cfg.resolveSensitive = true;
+  Navigator navigator(cfg, logger.cloneWithSuffix("Navigator"));
+  BremPropagator propagator(std::move(stepper), std::move(navigator),
+                            logger.cloneWithSuffix("BremPropagator"));
+  BremCKF trackFinder(std::move(propagator),
+                      logger.cloneWithSuffix("BremFinder"));
+
+  // build the track finder functions. owns the track finder object.
+  return std::make_shared<TrackFinderFunctionImpl<BremCKF>>(
+      std::move(trackFinder));
+}
+
+}  // namespace ActsExamples

Examples/Io/Root/src/RootTrackFitterPerformanceWriter.cpp

@@ -12,7 +12,6 @@
 #include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/EventData/Track.hpp"
 #include "ActsExamples/Framework/AlgorithmContext.hpp"
-#include "ActsExamples/Validation/TrackClassification.hpp"
 #include "ActsFatras/EventData/Barcode.hpp"
 #include "ActsPlugins/Root/HistogramConverter.hpp"
 

Python/Examples/python/reconstruction.py

@@ -1694,6 +1694,9 @@ def addCKFTracks(
         findTracks=acts.examples.TrackFindingAlgorithm.makeTrackFinderFunction(
             trackingGeometry, field, customLogLevel()
         ),
+        findTracksBrem=acts.examples.TrackFindingAlgorithm.makeBremTrackFinderFunction(
+            trackingGeometry, field, customLogLevel()
+        ),
         **acts.examples.defaultKWArgs(
             trackingGeometry=trackingGeometry,
             magneticField=field,

Python/Examples/src/TrackFinding.cpp

@@ -132,17 +132,27 @@ void addTrackFinding(py::module& mex) {
               *Acts::getDefaultLogger("TrackFinding", level));
         });
 
+    alg.def_static(
+        "makeBremTrackFinderFunction",
+        [](std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
+           std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
+           Acts::Logging::Level level) {
+          return Alg::makeBremTrackFinderFunction(
+              std::move(trackingGeometry), std::move(magneticField),
+              *Acts::getDefaultLogger("TrackFindingBrem", level));
+        });
+
     py::class_<Alg::TrackFinderFunction,
                std::shared_ptr<Alg::TrackFinderFunction>>(
         alg, "TrackFinderFunction");
 
-    ACTS_PYTHON_STRUCT(c, inputMeasurements, inputInitialTrackParameters,
-                       inputSeeds, outputTracks, trackingGeometry,
-                       magneticField, findTracks, measurementSelectorCfg,
-                       trackSelectorCfg, maxSteps, twoWay, reverseSearch,
-                       seedDeduplication, stayOnSeed, pixelVolumeIds,
-                       stripVolumeIds, maxPixelHoles, maxStripHoles, trimTracks,
-                       constrainToVolumeIds, endOfWorldVolumeIds);
+    ACTS_PYTHON_STRUCT(
+        c, inputMeasurements, inputInitialTrackParameters, inputSeeds,
+        outputTracks, trackingGeometry, magneticField, findTracks,
+        findTracksBrem, measurementSelectorCfg, trackSelectorCfg, maxSteps,
+        twoWay, reverseSearch, seedDeduplication, stayOnSeed, pixelVolumeIds,
+        stripVolumeIds, maxPixelHoles, maxStripHoles, trimTracks,
+        constrainToVolumeIds, endOfWorldVolumeIds);
   }
 }