ACTS kalman filter updates

Author: DraTeots

source/tdis/CMakeLists.txt

@@ -53,6 +53,11 @@ add_executable(tdis
         tracking/KalmanFittingFactory.cpp
         # tracking/CKFTracking.h
         # tracking/CKFTracking.cc
+        tracking/TrackFitterFunction.hpp
+        tracking/KalmanFitterFunction.cpp
+        tracking/RefittingCalibrator.hpp
+        tracking/RefittingCalibrator.cpp
+
 )
 
 # ---------- FIND REQUIRED PACKAGES -------------
@@ -66,6 +71,7 @@ find_package(Acts REQUIRED COMPONENTS Core PluginTGeo PluginJson)
 target_include_directories(tdis PUBLIC ${CMAKE_CURRENT_LIST_DIR} "${CMAKE_CURRENT_LIST_DIR}/..")
 target_link_libraries(tdis
         ${JANA_LIB}
+        libc++
         ROOT::RIO ROOT::Core
         podio::podio podio::podioRootIO podio_model_lib podio_model_dict
         spdlog::spdlog

source/tdis/tracking/DD4hepBField.cc

@@ -0,0 +1,194 @@
+// This file is part of the ACTS project.
+//
+// Copyright (C) 2016 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 https://mozilla.org/MPL/2.0/.
+
+#include "Acts/Definitions/Direction.hpp"
+#include "Acts/Definitions/TrackParametrization.hpp"
+#include "Acts/EventData/MultiTrajectory.hpp"
+#include "Acts/EventData/TrackContainer.hpp"
+#include "Acts/EventData/TrackStatePropMask.hpp"
+#include "Acts/EventData/VectorMultiTrajectory.hpp"
+#include "Acts/EventData/VectorTrackContainer.hpp"
+#include "Acts/EventData/detail/CorrectedTransformationFreeToBound.hpp"
+#include "Acts/Geometry/GeometryIdentifier.hpp"
+#include "Acts/Propagator/DirectNavigator.hpp"
+#include "Acts/Propagator/Navigator.hpp"
+#include "Acts/Propagator/Propagator.hpp"
+#include "Acts/Propagator/SympyStepper.hpp"
+#include "Acts/TrackFitting/GainMatrixSmoother.hpp"
+#include "Acts/TrackFitting/GainMatrixUpdater.hpp"
+#include "Acts/TrackFitting/KalmanFitter.hpp"
+#include "Acts/Utilities/Delegate.hpp"
+#include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/EventData/IndexSourceLink.hpp"
+#include "ActsExamples/EventData/MeasurementCalibration.hpp"
+#include "ActsExamples/EventData/Track.hpp"
+#include "ActsExamples/TrackFitting/RefittingCalibrator.hpp"
+#include "TrackFitterFunction.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <functional>
+#include <memory>
+#include <utility>
+#include <vector>
+
+namespace Acts {
+class MagneticFieldProvider;
+class SourceLink;
+class Surface;
+class TrackingGeometry;
+}  // namespace Acts
+
+namespace {
+
+using Stepper = Acts::SympyStepper;
+using Propagator = Acts::Propagator<Stepper, Acts::Navigator>;
+using Fitter = Acts::KalmanFitter<Propagator, Acts::VectorMultiTrajectory>;
+using DirectPropagator = Acts::Propagator<Stepper, Acts::DirectNavigator>;
+using DirectFitter =
+    Acts::KalmanFitter<DirectPropagator, Acts::VectorMultiTrajectory>;
+
+using TrackContainer =
+    Acts::TrackContainer<Acts::VectorTrackContainer,
+                         Acts::VectorMultiTrajectory, std::shared_ptr>;
+
+struct SimpleReverseFilteringLogic {
+  double momentumThreshold = 0;
+
+  bool doBackwardFiltering(
+      Acts::VectorMultiTrajectory::ConstTrackStateProxy trackState) const {
+    auto momentum = std::abs(1 / trackState.filtered()[Acts::eBoundQOverP]);
+    return (momentum <= momentumThreshold);
+  }
+};
+
+using namespace ActsExamples;
+
+struct KalmanFitterFunctionImpl final : public TrackFitterFunction {
+  Fitter fitter;
+  DirectFitter directFitter;
+
+  Acts::GainMatrixUpdater kfUpdater;
+  Acts::GainMatrixSmoother kfSmoother;
+  SimpleReverseFilteringLogic reverseFilteringLogic;
+
+  bool multipleScattering = false;
+  bool energyLoss = false;
+  Acts::FreeToBoundCorrection freeToBoundCorrection;
+
+  IndexSourceLink::SurfaceAccessor slSurfaceAccessor;
+
+  KalmanFitterFunctionImpl(Fitter&& f, DirectFitter&& df,
+                           const Acts::TrackingGeometry& trkGeo)
+      : fitter(std::move(f)),
+        directFitter(std::move(df)),
+        slSurfaceAccessor{trkGeo} {}
+
+  template <typename calibrator_t>
+  auto makeKfOptions(const GeneralFitterOptions& options,
+                     const calibrator_t& calibrator) const {
+    Acts::KalmanFitterExtensions<Acts::VectorMultiTrajectory> extensions;
+    extensions.updater.connect<
+        &Acts::GainMatrixUpdater::operator()<Acts::VectorMultiTrajectory>>(
+        &kfUpdater);
+    extensions.smoother.connect<
+        &Acts::GainMatrixSmoother::operator()<Acts::VectorMultiTrajectory>>(
+        &kfSmoother);
+    extensions.reverseFilteringLogic
+        .connect<&SimpleReverseFilteringLogic::doBackwardFiltering>(
+            &reverseFilteringLogic);
+
+    Acts::KalmanFitterOptions<Acts::VectorMultiTrajectory> kfOptions(
+        options.geoContext, options.magFieldContext, options.calibrationContext,
+        extensions, options.propOptions, &(*options.referenceSurface));
+
+    kfOptions.referenceSurfaceStrategy =
+        Acts::KalmanFitterTargetSurfaceStrategy::first;
+    kfOptions.multipleScattering = multipleScattering;
+    kfOptions.energyLoss = energyLoss;
+    kfOptions.freeToBoundCorrection = freeToBoundCorrection;
+    kfOptions.extensions.calibrator.connect<&calibrator_t::calibrate>(
+        &calibrator);
+
+    if (options.doRefit) {
+      kfOptions.extensions.surfaceAccessor.connect<&RefittingCalibrator::accessSurface>();
+    } else {
+      kfOptions.extensions.surfaceAccessor.connect<&IndexSourceLink::SurfaceAccessor::operator()>(
+              &slSurfaceAccessor);
+    }
+
+    return kfOptions;
+  }
+
+  TrackFitterResult operator()(const std::vector<Acts::SourceLink>& sourceLinks,
+                               const TrackParameters& initialParameters,
+                               const GeneralFitterOptions& options,
+                               const MeasurementCalibratorAdapter& calibrator,
+                               TrackContainer& tracks) const override {
+    const auto kfOptions = makeKfOptions(options, calibrator);
+    return fitter.fit(sourceLinks.begin(), sourceLinks.end(), initialParameters,
+                      kfOptions, tracks);
+  }
+
+  TrackFitterResult operator()(
+      const std::vector<Acts::SourceLink>& sourceLinks,
+      const TrackParameters& initialParameters,
+      const GeneralFitterOptions& options,
+      const RefittingCalibrator& calibrator,
+      const std::vector<const Acts::Surface*>& surfaceSequence,
+      TrackContainer& tracks) const override {
+    const auto kfOptions = makeKfOptions(options, calibrator);
+    return directFitter.fit(sourceLinks.begin(), sourceLinks.end(),
+                            initialParameters, kfOptions, surfaceSequence,
+                            tracks);
+  }
+};
+
+}  // namespace
+
+std::shared_ptr<ActsExamples::TrackFitterFunction>
+ActsExamples::makeKalmanFitterFunction(
+    std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry,
+    std::shared_ptr<const Acts::MagneticFieldProvider> magneticField,
+    bool multipleScattering, bool energyLoss,
+    double reverseFilteringMomThreshold,
+    Acts::FreeToBoundCorrection freeToBoundCorrection,
+    const Acts::Logger& logger) {
+  // Stepper should be copied into the fitters
+  const Stepper stepper(std::move(magneticField));
+
+  // Standard fitter
+  const auto& geo = *trackingGeometry;
+  Acts::Navigator::Config cfg{std::move(trackingGeometry)};
+  cfg.resolvePassive = false;
+  cfg.resolveMaterial = true;
+  cfg.resolveSensitive = true;
+  Acts::Navigator navigator(cfg, logger.cloneWithSuffix("Navigator"));
+  Propagator propagator(stepper, std::move(navigator),
+                        logger.cloneWithSuffix("Propagator"));
+  Fitter trackFitter(std::move(propagator), logger.cloneWithSuffix("Fitter"));
+
+  // Direct fitter
+  Acts::DirectNavigator directNavigator{
+      logger.cloneWithSuffix("DirectNavigator")};
+  DirectPropagator directPropagator(stepper, std::move(directNavigator),
+                                    logger.cloneWithSuffix("DirectPropagator"));
+  DirectFitter directTrackFitter(std::move(directPropagator),
+                                 logger.cloneWithSuffix("DirectFitter"));
+
+  // build the fitter function. owns the fitter object.
+  auto fitterFunction = std::make_shared<KalmanFitterFunctionImpl>(
+      std::move(trackFitter), std::move(directTrackFitter), geo);
+  fitterFunction->multipleScattering = multipleScattering;
+  fitterFunction->energyLoss = energyLoss;
+  fitterFunction->reverseFilteringLogic.momentumThreshold =
+      reverseFilteringMomThreshold;
+  fitterFunction->freeToBoundCorrection = freeToBoundCorrection;
+
+  return fitterFunction;
+}

source/tdis/tracking/DD4hepBField.h

@@ -31,6 +31,7 @@
 #include <ActsExamples/EventData/IndexSourceLink.hpp>
 #include <ActsExamples/EventData/Measurement.hpp>
 #include <ActsExamples/EventData/MeasurementCalibration.hpp>
+
 #include <ActsExamples/EventData/Track.hpp>
 //#include <edm4eic/Cov3f.h>
 //#include <edm4eic/Cov6f.h>
@@ -78,6 +79,8 @@ void tdis::tracking::KalmanFittingFactory::Execute(int32_t runNumber, uint64_t e
 
     m_log->debug("{}::Execute", this->GetTypeName());
 
+    std::shared_ptr<ActsExamples::TrackFitterFunction> fit;
+
 }
 
 template <typename stepper_t>