TrackFitting/detail/SymmetricKlDistanceMatrix.hpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2023 CERN for the benefit of the Acts project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/EventData/MultiComponentTrackParameters.hpp"
0012 #include "Acts/EventData/MultiTrajectory.hpp"
0013 #include "Acts/EventData/TrackParameters.hpp"
0014 #include "Acts/TrackFitting/detail/GsfComponentMerging.hpp"
0015 #include "Acts/TrackFitting/detail/GsfUtils.hpp"
0016
0017 namespace Acts::detail {
0018
0019 /// Computes the Kullback-Leibler distance between two components as shown in
0020 /// https://arxiv.org/abs/2001.00727v1 but ignoring the weights
0021 template <typename component_t, typename component_projector_t>
0022 auto computeSymmetricKlDivergence(const component_t &a, const component_t &b,
0023                                   const component_projector_t &proj) {
0024   using namespace Acts;
0025   const auto parsA = proj(a).boundPars[eBoundQOverP];
0026   const auto parsB = proj(b).boundPars[eBoundQOverP];
0027   const auto covA = proj(a).boundCov(eBoundQOverP, eBoundQOverP);
0028   const auto covB = proj(b).boundCov(eBoundQOverP, eBoundQOverP);
0029
0030   assert(covA != 0.0);
0031   assert(std::isfinite(covA));
0032   assert(covB != 0.0);
0033   assert(std::isfinite(covB));
0034
0035   const auto kl = covA * (1 / covB) + covB * (1 / covA) +
0036                   (parsA - parsB) * (1 / covA + 1 / covB) * (parsA - parsB);
0037
0038   assert(kl >= 0.0 && "kl-divergence must be non-negative");
0039
0040   return kl;
0041 }
0042
0043 template <typename component_t, typename component_projector_t,
0044           typename angle_desc_t>
0045 auto mergeComponents(const component_t &a, const component_t &b,
0046                      const component_projector_t &proj,
0047                      const angle_desc_t &angle_desc) {
0048   assert(proj(a).weight >= 0.0 && proj(b).weight >= 0.0 &&
0049          "non-positive weight");
0050
0051   std::array range = {std::ref(proj(a)), std::ref(proj(b))};
0052   const auto refProj = [](auto &c) {
0053     return std::tie(c.get().weight, c.get().boundPars, c.get().boundCov);
0054   };
0055
0056   auto [mergedPars, mergedCov] =
0057       gaussianMixtureMeanCov(range, refProj, angle_desc);
0058
0059   component_t ret = a;
0060   proj(ret).boundPars = mergedPars;
0061   proj(ret).boundCov = mergedCov;
0062   proj(ret).weight = proj(a).weight + proj(b).weight;
0063
0064   return ret;
0065 }
0066
0067 /// @brief Class representing a symmetric distance matrix
0068 class SymmetricKLDistanceMatrix {
0069   using Array = Eigen::Array<Acts::ActsScalar, Eigen::Dynamic, 1>;
0070   using Mask = Eigen::Array<bool, Eigen::Dynamic, 1>;
0071
0072   Array m_distances;
0073   Mask m_mask;
0074   std::vector<std::pair<std::size_t, std::size_t>> m_mapToPair;
0075   std::size_t m_numberComponents;
0076
0077   template <typename array_t, typename setter_t>
0078   void setAssociated(std::size_t n, array_t &array, setter_t &&setter) {
0079     const auto indexConst = (n - 1) * n / 2;
0080
0081     // Rows
0082     for (auto i = 0ul; i < n; ++i) {
0083       array[indexConst + i] = setter(n, i);
0084     }
0085
0086     // Columns
0087     for (auto i = n + 1; i < m_numberComponents; ++i) {
0088       array[(i - 1) * i / 2 + n] = setter(n, i);
0089     }
0090   }
0091
0092  public:
0093   template <typename component_t, typename projector_t>
0094   SymmetricKLDistanceMatrix(const std::vector<component_t> &cmps,
0095                             const projector_t &proj)
0096       : m_distances(Array::Zero(cmps.size() * (cmps.size() - 1) / 2)),
0097         m_mask(Mask::Ones(cmps.size() * (cmps.size() - 1) / 2)),
0098         m_mapToPair(m_distances.size()),
0099         m_numberComponents(cmps.size()) {
0100     for (auto i = 1ul; i < m_numberComponents; ++i) {
0101       const auto indexConst = (i - 1) * i / 2;
0102       for (auto j = 0ul; j < i; ++j) {
0103         m_mapToPair.at(indexConst + j) = {i, j};
0104         m_distances[indexConst + j] =
0105             computeSymmetricKlDivergence(cmps[i], cmps[j], proj);
0106       }
0107     }
0108   }
0109
0110   auto at(std::size_t i, std::size_t j) const {
0111     return m_distances[i * (i - 1) / 2 + j];
0112   }
0113
0114   template <typename component_t, typename projector_t>
0115   void recomputeAssociatedDistances(std::size_t n,
0116                                     const std::vector<component_t> &cmps,
0117                                     const projector_t &proj) {
0118     assert(cmps.size() == m_numberComponents && "size mismatch");
0119
0120     setAssociated(n, m_distances, [&](std::size_t i, std::size_t j) {
0121       return computeSymmetricKlDivergence(cmps[i], cmps[j], proj);
0122     });
0123   }
0124
0125   void maskAssociatedDistances(std::size_t n) {
0126     setAssociated(n, m_mask, [&](std::size_t, std::size_t) { return false; });
0127   }
0128
0129   auto minDistancePair() const {
0130     auto min = std::numeric_limits<Acts::ActsScalar>::max();
0131     std::size_t idx = 0;
0132
0133     for (auto i = 0l; i < m_distances.size(); ++i) {
0134       if (auto new_min = std::min(min, m_distances[i]);
0135           m_mask[i] && new_min < min) {
0136         min = new_min;
0137         idx = i;
0138       }
0139     }
0140
0141     return m_mapToPair.at(idx);
0142   }
0143
0144   friend std::ostream &operator<<(std::ostream &os,
0145                                   const SymmetricKLDistanceMatrix &m) {
0146     const auto prev_precision = os.precision();
0147     const int width = 8;
0148     const int prec = 2;
0149
0150     os << "\n";
0151     os << std::string(width, ' ') << " | ";
0152     for (auto j = 0ul; j < m.m_numberComponents - 1; ++j) {
0153       os << std::setw(width) << j << "  ";
0154     }
0155     os << "\n";
0156     os << std::string((width + 3) + (width + 2) * (m.m_numberComponents - 1),
0157                       '-');
0158     os << "\n";
0159
0160     for (auto i = 1ul; i < m.m_numberComponents; ++i) {
0161       const auto indexConst = (i - 1) * i / 2;
0162       os << std::setw(width) << i << " | ";
0163       for (auto j = 0ul; j < i; ++j) {
0164         os << std::setw(width) << std::setprecision(prec)
0165            << m.m_distances[indexConst + j] << "  ";
0166       }
0167       os << "\n";
0168     }
0169     os << std::setprecision(prev_precision);
0170     return os;
0171   }
0172 };
0173
0174 }  // namespace Acts::detail