AdaptiveMultiVertexFitter.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file AdaptiveMultiVertexFitter.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2019-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/.

#pragma once

#include "Acts/Definitions/Algebra.hpp"
#include "Acts/EventData/TrackParameters.hpp"
#include "Acts/Utilities/AnnealingUtility.hpp"
#include "Acts/Utilities/Logger.hpp"
#include "Acts/Utilities/Result.hpp"
#include "Acts/Vertexing/AMVFInfo.hpp"
#include "Acts/Vertexing/ImpactPointEstimator.hpp"
#include "Acts/Vertexing/LinearizerConcept.hpp"
#include "Acts/Vertexing/TrackAtVertex.hpp"
#include "Acts/Vertexing/Vertex.hpp"
#include "Acts/Vertexing/VertexingOptions.hpp"

#include <functional>

namespace Acts {

template <typename input_track_t, typename linearizer_t>
class AdaptiveMultiVertexFitter {
  static_assert(LinearizerConcept<linearizer_t>,
                "Linearizer does not fulfill linearizer concept.");

 public:
  using InputTrack_t = input_track_t;
  using Propagator_t = typename linearizer_t::Propagator_t;
  using Linearizer_t = linearizer_t;

 private:
  using IPEstimator = ImpactPointEstimator<InputTrack_t, Propagator_t>;

 public:
  struct State {
    State(const MagneticFieldProvider& field,
          const Acts::MagneticFieldContext& magContext)
        : ipState(field.makeCache(magContext)),
          linearizerState(field.makeCache(magContext)) {}
    // Vertex collection to be fitted
    std::vector<Vertex<InputTrack_t>*> vertexCollection;

    // Annealing state
    AnnealingUtility::State annealingState;

    // IPEstimator state
    typename IPEstimator::State ipState;

    // Linearizer state
    typename Linearizer_t::State linearizerState;

    // Map to store vertices information
    std::map<Vertex<InputTrack_t>*, VertexInfo<InputTrack_t>> vtxInfoMap;

    std::multimap<const InputTrack_t*, Vertex<InputTrack_t>*>
        trackToVerticesMultiMap;

    std::map<std::pair<const InputTrack_t*, Vertex<InputTrack_t>*>,
             TrackAtVertex<InputTrack_t>>
        tracksAtVerticesMap;

    State() = default;

    // Adds a vertex to trackToVerticesMultiMap
    void addVertexToMultiMap(Vertex<InputTrack_t>& vtx) {
      for (auto trk : vtxInfoMap[&vtx].trackLinks) {
        trackToVerticesMultiMap.emplace(trk, &vtx);
      }
    }

    // Removes a vertex from trackToVerticesMultiMap
    void removeVertexFromMultiMap(Vertex<InputTrack_t>& vtx) {
      for (auto iter = trackToVerticesMultiMap.begin();
           iter != trackToVerticesMultiMap.end();) {
        if (iter->second == &vtx) {
          iter = trackToVerticesMultiMap.erase(iter);
        } else {
          ++iter;
        }
      }
    }
  };

  struct Config {
    Config(const IPEstimator& est) : ipEst(est) {}

    // ImpactPointEstimator
    IPEstimator ipEst;

    AnnealingUtility annealingTool;

    // Number of max iterations
    unsigned int maxIterations{30};

    // Max distance to linearization point allowed
    // without relinearization
    double maxDistToLinPoint{0.5};

    // Minimum track weight needed for track to be considered
    double minWeight{0.0001};

    // Max relative shift of vertex during one iteration
    double maxRelativeShift{0.01};

    // Do smoothing after multivertex fit
    bool doSmoothing{false};

    // Use time information when calculating the vertex compatibility
    bool useTime{false};
  };

  template <
      typename T = InputTrack_t,
      std::enable_if_t<std::is_same<T, BoundTrackParameters>::value, int> = 0>
  AdaptiveMultiVertexFitter(Config& cfg,
                            std::unique_ptr<const Logger> logger =
                                getDefaultLogger("AdaptiveMultiVertexFitter",
                                                 Logging::INFO))
      : m_cfg(std::move(cfg)),
        m_extractParameters([](T params) { return params; }),
        m_logger(std::move(logger)) {}

  AdaptiveMultiVertexFitter(
      Config& cfg, std::function<BoundTrackParameters(InputTrack_t)> func,
      std::unique_ptr<const Logger> logger =
          getDefaultLogger("AdaptiveMultiVertexFitter", Logging::INFO))
      : m_cfg(std::move(cfg)),
        m_extractParameters(func),
        m_logger(std::move(logger)) {}

  Result<void> fit(
      State& state, const std::vector<Vertex<InputTrack_t>*>& verticesToFit,
      const Linearizer_t& linearizer,
      const VertexingOptions<InputTrack_t>& vertexingOptions) const;

  Result<void> addVtxToFit(
      State& state, Vertex<InputTrack_t>& newVertex,
      const Linearizer_t& linearizer,
      const VertexingOptions<InputTrack_t>& vertexingOptions) const;

 private:
  const Config m_cfg;

  std::function<BoundTrackParameters(InputTrack_t)> m_extractParameters;

  std::unique_ptr<const Logger> m_logger;

  const Logger& logger() const { return *m_logger; }

  Result<void> fitImpl(
      State& state, const Linearizer_t& linearizer,
      const VertexingOptions<InputTrack_t>& vertexingOptions) const;

  bool isAlreadyInList(
      Vertex<InputTrack_t>* vtx,
      const std::vector<Vertex<InputTrack_t>*>& verticesVec) const;

  Result<void> prepareVertexForFit(
      State& state, Vertex<InputTrack_t>* vtx,
      const VertexingOptions<InputTrack_t>& vertexingOptions) const;

  Result<void> setAllVertexCompatibilities(
      State& state, Vertex<InputTrack_t>* currentVtx,
      const VertexingOptions<input_track_t>& vertexingOptions) const;

  Result<void> setWeightsAndUpdate(
      State& state, const Linearizer_t& linearizer,
      const VertexingOptions<input_track_t>& vertexingOptions) const;

  std::vector<double> collectTrackToVertexCompatibilities(
      State& state, const InputTrack_t* trk) const;

  bool checkSmallShift(State& state) const;

  void doVertexSmoothing(State& state) const;
};

}  // namespace Acts

#include "Acts/Vertexing/AdaptiveMultiVertexFitter.ipp"