KDTree.hpp

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

// This file is part of the Acts project.
//
// Copyright (C) 2021 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/Utilities/RangeXD.hpp"

#include <algorithm>
#include <array>
#include <cmath>
#include <functional>
#include <memory>
#include <string>
#include <vector>

namespace Acts {
template <std::size_t Dims, typename Type, typename Scalar = double,
          template <typename, std::size_t> typename Vector = std::array,
          std::size_t LeafSize = 4>
class KDTree {
 public:
  using value_t = Type;

  using range_t = RangeXD<Dims, Scalar, Vector>;

  using coordinate_t = Vector<Scalar, Dims>;

  using pair_t = std::pair<coordinate_t, Type>;

  using vector_t = std::vector<pair_t>;

  using iterator_t = typename vector_t::iterator;

  using const_iterator_t = typename vector_t::const_iterator;

  // We do not need an empty constructor - this is never useful.
  KDTree() = delete;

  KDTree(vector_t &&d) : m_elems(d) {
    // To start out, we need to check whether we need to construct a leaf node
    // or an internal node. We create a leaf only if we have at most as many
    // elements as the number of elements that can fit into a leaf node.
    // Hopefully most invocations of this constructor will have more than a few
    // elements!
    //
    // One interesting thing to note is that all of the nodes in the k-d tree
    // have a range in the element vector of the outermost node. They simply
    // make in-place changes to this array, and they hold no memory of their
    // own.
    m_root = std::make_unique<KDTreeNode>(m_elems.begin(), m_elems.end(),
                                          m_elems.size() > LeafSize
                                              ? KDTreeNode::NodeType::Internal
                                              : KDTreeNode::NodeType::Leaf,
                                          0UL);
  }

  std::vector<Type> rangeSearch(const range_t &r) const {
    std::vector<Type> out;

    rangeSearch(r, out);

    return out;
  }

  std::vector<pair_t> rangeSearchWithKey(const range_t &r) const {
    std::vector<pair_t> out;

    rangeSearchWithKey(r, out);

    return out;
  }

  void rangeSearch(const range_t &r, std::vector<Type> &v) const {
    rangeSearchInserter(r, std::back_inserter(v));
  }

  void rangeSearchWithKey(const range_t &r, std::vector<pair_t> &v) const {
    rangeSearchInserterWithKey(r, std::back_inserter(v));
  }

  template <typename OutputIt>
  void rangeSearchInserter(const range_t &r, OutputIt i) const {
    rangeSearchMapDiscard(
        r, [i](const coordinate_t &, const Type &v) mutable { i = v; });
  }

  template <typename OutputIt>
  void rangeSearchInserterWithKey(const range_t &r, OutputIt i) const {
    rangeSearchMapDiscard(r, [i](const coordinate_t &c, const Type &v) mutable {
      i = {c, v};
    });
  }

  template <typename Result>
  std::vector<Result> rangeSearchMap(
      const range_t &r,
      std::function<Result(const coordinate_t &, const Type &)> f) const {
    std::vector<Result> out;

    rangeSearchMapInserter(r, f, std::back_inserter(out));

    return out;
  }

  template <typename Result, typename OutputIt>
  void rangeSearchMapInserter(
      const range_t &r,
      std::function<Result(const coordinate_t &, const Type &)> f,
      OutputIt i) const {
    rangeSearchMapDiscard(r, [i, f](const coordinate_t &c,
                                    const Type &v) mutable { i = f(c, v); });
  }

  template <typename Callable>
  void rangeSearchMapDiscard(const range_t &r, Callable &&f) const {
    m_root->rangeSearchMapDiscard(r, std::forward<Callable>(f));
  }

  std::size_t size(void) const { return m_root->size(); }

  const_iterator_t begin(void) const { return m_elems.begin(); }

  const_iterator_t end(void) const { return m_elems.end(); }

 private:
  static Scalar nextRepresentable(Scalar v) {
    // I'm not super happy with this bit of code, but since 1D ranges are
    // semi-open, we can't simply incorporate values by setting the maximum to
    // them. Instead, what we need to do is get the next representable value.
    // For integer values, this means adding one. For floating point types, we
    // rely on the nextafter method to get the smallest possible value that is
    // larger than the one we requested.
    if constexpr (std::is_integral_v<Scalar>) {
      return v + 1;
    } else if constexpr (std::is_floating_point_v<Scalar>) {
      return std::nextafter(v, std::numeric_limits<Scalar>::max());
    }
  }

  static range_t boundingBox(iterator_t b, iterator_t e) {
    // Firstly, we find the minimum and maximum value in each dimension to
    // construct a bounding box around this node's values.
    std::array<Scalar, Dims> min_v{}, max_v{};

    for (std::size_t i = 0; i < Dims; ++i) {
      min_v[i] = std::numeric_limits<Scalar>::max();
      max_v[i] = std::numeric_limits<Scalar>::lowest();
    }

    for (iterator_t i = b; i != e; ++i) {
      for (std::size_t j = 0; j < Dims; ++j) {
        min_v[j] = std::min(min_v[j], i->first[j]);
        max_v[j] = std::max(max_v[j], i->first[j]);
      }
    }

    // Then, we construct a k-dimensional range from the given minima and
    // maxima, which again is just a bounding box.
    range_t r;

    for (std::size_t j = 0; j < Dims; ++j) {
      r[j] = {min_v[j], nextRepresentable(max_v[j])};
    }

    return r;
  }

  class KDTreeNode {
   public:
    enum class NodeType { Internal, Leaf };

    KDTreeNode(iterator_t _b, iterator_t _e, NodeType _t, std::size_t _d)
        : m_type(_t),
          m_begin_it(_b),
          m_end_it(_e),
          m_range(boundingBox(m_begin_it, m_end_it)) {
      if (m_type == NodeType::Internal) {
        // This constant determines the maximum number of elements where we
        // still
        // calculate the exact median of the values for the purposes of
        // splitting. In general, the closer the pivot value is to the true
        // median, the more balanced the tree will be. However, calculating the
        // median exactly is an O(n log n) operation, while approximating it is
        // an O(1) time.
        constexpr std::size_t max_exact_median = 128;

        iterator_t pivot;

        // Next, we need to determine the pivot point of this node, that is to
        // say the point in the selected pivot dimension along which point we
        // will split the range. To do this, we check how large the set of
        // elements is. If it is sufficiently small, we use the median.
        // Otherwise we use the mean.
        if (size() > max_exact_median) {
          // In this case, we have a lot of elements, and sorting the range to
          // find the true median might be too expensive. Therefore, we will
          // just use the middle value between the minimum and maximum. This is
          // not nearly as accurate as using the median, but it's a nice cheat.
          Scalar mid = static_cast<Scalar>(0.5) *
                       (m_range[_d].max() + m_range[_d].min());

          pivot = std::partition(m_begin_it, m_end_it, [=](const pair_t &i) {
            return i.first[_d] < mid;
          });
        } else {
          // If the number of elements is fairly small, we will just calculate
          // the median exactly. We do this by finding the values in the
          // dimension, sorting it, and then taking the middle one.
          std::sort(m_begin_it, m_end_it,
                    [_d](const typename iterator_t::value_type &a,
                         const typename iterator_t::value_type &b) {
                      return a.first[_d] < b.first[_d];
                    });

          pivot = m_begin_it + (std::distance(m_begin_it, m_end_it) / 2);
        }

        // This should never really happen, but in very select cases where there
        // are a lot of equal values in the range, the pivot can end up all the
        // way at the end of the array and we end up in an infinite loop. We
        // check for pivot points which would not split the range, and fix them
        // if they occur.
        if (pivot == m_begin_it || pivot == std::prev(m_end_it)) {
          pivot = std::next(m_begin_it, LeafSize);
        }

        // Calculate the number of elements on the left-hand side, as well as
        // the right-hand side. We do this by calculating the difference from
        // the begin and end of the array to the pivot point.
        std::size_t lhs_size = std::distance(m_begin_it, pivot);
        std::size_t rhs_size = std::distance(pivot, m_end_it);

        // Next, we check whether the left-hand node should be another internal
        // node or a leaf node, and we construct the node recursively.
        m_lhs = std::make_unique<KDTreeNode>(
            m_begin_it, pivot,
            lhs_size > LeafSize ? NodeType::Internal : NodeType::Leaf,
            (_d + 1) % Dims);

        // Same on the right hand side.
        m_rhs = std::make_unique<KDTreeNode>(
            pivot, m_end_it,
            rhs_size > LeafSize ? NodeType::Internal : NodeType::Leaf,
            (_d + 1) % Dims);
      }
    }

    template <typename Callable>
    void rangeSearchMapDiscard(const range_t &r, Callable &&f) const {
      // Determine whether the range completely covers the bounding box of
      // this leaf node. If it is, we can copy all values without having to
      // check for them being inside the range again.
      bool contained = r >= m_range;

      if (m_type == NodeType::Internal) {
        // Firstly, we can check if the range completely contains the bounding
        // box of this node. If that is the case, we know for certain that any
        // value contained below this node should end up in the output, and we
        // can stop recursively looking for them.
        if (contained) {
          // We can also pre-allocate space for the number of elements, since we
          // are inserting all of them anyway.
          for (iterator_t i = m_begin_it; i != m_end_it; ++i) {
            f(i->first, i->second);
          }

          return;
        }

        assert(m_lhs && m_rhs && "Did not find lhs and rhs");

        // If we have a left-hand node (which we should!), then we check if
        // there is any overlap between the target range and the bounding box of
        // the left-hand node. If there is, we recursively search in that node.
        if (m_lhs->range() && r) {
          m_lhs->rangeSearchMapDiscard(r, std::forward<Callable>(f));
        }

        // Then, we perform exactly the same procedure for the right hand side.
        if (m_rhs->range() && r) {
          m_rhs->rangeSearchMapDiscard(r, std::forward<Callable>(f));
        }
      } else {
        // Iterate over all the elements in this leaf node. This should be a
        // relatively small number (the LeafSize template parameter).
        for (iterator_t i = m_begin_it; i != m_end_it; ++i) {
          // We need to check whether the element is actually inside the range.
          // In case this node's bounding box is fully contained within the
          // range, we don't actually need to check this.
          if (contained || r.contains(i->first)) {
            f(i->first, i->second);
          }
        }
      }
    }

    std::size_t size() const { return std::distance(m_begin_it, m_end_it); }

    const range_t &range() const { return m_range; }

   protected:
    NodeType m_type;

    const iterator_t m_begin_it, m_end_it;

    const range_t m_range;

    std::unique_ptr<KDTreeNode> m_lhs;
    std::unique_ptr<KDTreeNode> m_rhs;
  };

  vector_t m_elems;

  std::unique_ptr<KDTreeNode> m_root;
};
}  // namespace Acts