KdtSurfacesProvider.hpp

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

// 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/.

#pragma once

#include "Acts/Definitions/Algebra.hpp"
#include "Acts/Detector/detail/ReferenceGenerators.hpp"
#include "Acts/Detector/interface/ISurfacesProvider.hpp"
#include "Acts/Surfaces/Surface.hpp"
#include "Acts/Utilities/BinningData.hpp"
#include "Acts/Utilities/KDTree.hpp"

#include <array>
#include <stdexcept>
#include <tuple>
#include <vector>

namespace Acts {

namespace Experimental {

template <size_t kDIM = 2u, size_t bSize = 100u,
          typename reference_generator =
              detail::PolyhedronReferenceGenerator<1u, false>>
class KdtSurfaces {
 public:
  using KDTS =
      KDTree<kDIM, std::shared_ptr<Surface>, ActsScalar, std::array, bSize>;

  using Query = std::array<ActsScalar, kDIM>;

  using Entry = std::pair<Query, std::shared_ptr<Surface>>;

  KdtSurfaces(const GeometryContext& gctx,
              const std::vector<std::shared_ptr<Surface>>& surfaces,
              const std::array<BinningValue, kDIM>& casts,
              const reference_generator& rgen =
                  detail::PolyhedronReferenceGenerator<1u, false>{})
      : m_kdt(nullptr), m_casts(casts), m_rGenerator(rgen) {
    // Simple check if the dimension is correct
    if (kDIM == 0u) {
      throw std::invalid_argument(
          "KdtSurfaces: dimension and/or cast rules are incorrect.");
    }
    // Fill the tree from surfaces
    std::vector<Entry> kdtEntries;
    kdtEntries.reserve(surfaces.size());
    for (auto& s : surfaces) {
      // Generate the references and the center of gravity from it
      const auto references = m_rGenerator.references(gctx, *s);
      std::vector<Query> castedReferences;
      castedReferences.reserve(references.size());
      for (const auto& r : references) {
        //  Now cast into the correct fill position
        Query rc = {};
        fillCasts(r, rc, std::make_integer_sequence<std::size_t, kDIM>{});
        castedReferences.push_back(rc);
      }
      // Calculate the center of gravity in casted frame
      kdtEntries.push_back({cog(castedReferences), s});
    }
    // Create the KDTree
    m_kdt = std::make_unique<KDTS>(std::move(kdtEntries));
  }

  std::vector<std::shared_ptr<Surface>> surfaces(
      const RangeXD<kDIM, ActsScalar>& range) const {
    // Strip the surfaces
    std::vector<std::shared_ptr<Surface>> surfacePtrs;
    auto surfaceQuery = m_kdt->rangeSearchWithKey(range);
    std::for_each(surfaceQuery.begin(), surfaceQuery.end(),
                  [&](auto& s) { surfacePtrs.push_back(s.second); });
    return surfacePtrs;
  }

  std::vector<std::shared_ptr<Surface>> surfaces(const Extent& extent) const {
    RangeXD<kDIM, ActsScalar> qRange;
    for (auto [ibv, v] : enumerate(m_casts)) {
      qRange[ibv] = extent.range(v);
    }
    return surfaces(qRange);
  }

 private:
  std::unique_ptr<KDTS> m_kdt = nullptr;

  std::array<BinningValue, kDIM> m_casts = {};

  reference_generator m_rGenerator;

  template <typename Array, std::size_t... idx>
  void fillCasts(const Vector3& position, Array& a,
                 std::index_sequence<idx...> /*indices*/) const {
    ((a[idx] = VectorHelpers::cast(position, m_casts[idx])), ...);
  }

  Query cog(const std::vector<Query>& cQueries) const {
    // If there is only one position, return it
    if (cQueries.size() == 1) {
      return cQueries.front();
    }
    // Build the center of gravity of the n positions
    Query c{};
    float weight = 1. / cQueries.size();
    for (auto& q : cQueries) {
      std::transform(c.begin(), c.end(), q.begin(), c.begin(),
                     std::plus<ActsScalar>());
    }
    std::for_each(c.begin(), c.end(), [&](auto& v) { v *= weight; });
    return c;
  }
};

template <size_t kDIM = 2u, size_t bSize = 100u,
          typename reference_generator =
              detail::PolyhedronReferenceGenerator<1u, false>>
class KdtSurfacesProvider : public ISurfacesProvider {
 public:
  KdtSurfacesProvider(
      std::shared_ptr<KdtSurfaces<kDIM, bSize, reference_generator>> kdts,
      const Extent& kregion)
      : m_kdt(std::move(kdts)), m_region(kregion) {
    if (m_kdt == nullptr) {
      throw std::invalid_argument(
          "KdtSurfacesProvider: no KDTree structure provided.");
    }
  }

  std::vector<std::shared_ptr<Surface>> surfaces(
      [[maybe_unused]] const GeometryContext& gctx) const final {
    return m_kdt->surfaces(m_region);
  }

 private:
  std::shared_ptr<KdtSurfaces<kDIM, bSize, reference_generator>> m_kdt =
      nullptr;
  Extent m_region;
};

}  // namespace Experimental

}  // namespace Acts