KdTree.h

Go to the documentation of this file.

 
#ifndef PBAT_GEOMETRY_KDTREE_H
#define PBAT_GEOMETRY_KDTREE_H
 
#include "AxisAlignedBoundingBox.h"
#include "pbat/common/NAryTreeTraversal.h"
#include "pbat/common/Stack.h"
 
#include <algorithm>
#include <pbat/Aliases.h>
#include <pbat/profiling/Profiling.h>
#include <ranges>
#include <vector>
 
namespace pbat {
namespace geometry {

struct KdTreeNode
{
    enum { kLeafNodeLeftChild = -2 };
    Index begin; 
    Index n;     
    Index c{kLeafNodeLeftChild}; 

    [[maybe_unused]] bool IsLeaf() const { return c == -2; }
    [[maybe_unused]] bool IsInternal() const { return not IsLeaf(); }
    [[maybe_unused]] auto Left() const { return c; }
    [[maybe_unused]] auto Right() const { return c + 1; }
};

template <int Dims>
class KdTree
{
  public:
    KdTree() = default;

    template <class TDerivedP>
    KdTree(Eigen::DenseBase<TDerivedP> const& P, Index maxPointsInLeaf = 8);
    template <class TDerivedP>
    void Construct(Eigen::DenseBase<TDerivedP> const& P, Index maxPointsInLeaf);
    template <class FVisit>
    void DepthFirstSearch(FVisit visit, Index root = 0) const;

    std::vector<KdTreeNode> const& Nodes() const { return mNodes; }
    IndexVectorX const& Permutation() const { return mPermutation; }
    auto PointsInNode(Index const nodeIdx) const;
    auto PointsInNode(KdTreeNode const& node) const;
    Index constexpr Root() const { return 0; }
 
  private:
    IndexVectorX mPermutation;      
    std::vector<KdTreeNode> mNodes; 
};
 
template <int Dims>
template <class TDerivedP>
inline KdTree<Dims>::KdTree(Eigen::DenseBase<TDerivedP> const& P, Index maxPointsInLeaf)
{
    Construct(P, maxPointsInLeaf);
}
 
template <int Dims>
template <class TDerivedP>
inline void KdTree<Dims>::Construct(Eigen::DenseBase<TDerivedP> const& P, Index maxPointsInLeaf)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.KdTree.Construct");
    // Our k-D tree is a full binary tree (i.e. # nodes = 2*leaves - 1). We try to estimate
    // the number of leaf nodes to reserve memory up-front and prevent any reallocation, but without
    // excessively using up memory. We estimate the number of nodes per leaf to be 80% of the
    // maximum number of points in a leaf node.
    Index const nPoints                      = P.cols();
    Scalar constexpr kEstimatedLeafOccupancy = 0.8;
    auto const nEstimatedNodesPerLeaf =
        (maxPointsInLeaf > 1) ?
            static_cast<Index>(
                std::ceil(kEstimatedLeafOccupancy * static_cast<Scalar>(maxPointsInLeaf))) :
            1;
    auto const nEstimatedLeafNodes =
        (nPoints / nEstimatedNodesPerLeaf) + (nPoints % nEstimatedNodesPerLeaf != 0);
    mNodes.clear();
#include "pbat/warning/Push.h"
#include "pbat/warning/SignConversion.h"
    mNodes.reserve(2 * nEstimatedLeafNodes - 1);
    auto iota = std::views::iota(Index{0}, nPoints);
    mPermutation.resize(nPoints);
    std::copy(iota.begin(), iota.end(), mPermutation.data());
 
    // Top-down construction of the k-D tree
    struct StackFrame
    {
        geometry::AxisAlignedBoundingBox<Dims> aabb;
        Index nodeIdx;
        Index begin;
        Index n;
    };
    common::Stack<StackFrame, 128> stack{};
    Index constexpr root = 0;
    stack.Push({geometry::AxisAlignedBoundingBox<Dims>{P}, root, 0, nPoints});
    mNodes.emplace_back(0, nPoints, KdTreeNode::kLeafNodeLeftChild);
    while (not stack.IsEmpty())
    {
        auto const [aabb, nodeIdx, begin, n] = stack.Pop();
        if (n <= maxPointsInLeaf)
            continue;
 
        // Find the dimension with the largest extent
        Eigen::Index dimension{};
        (aabb.max() - aabb.min()).maxCoeff(&dimension);
        // Partition the points along the dimension on left/right sides of median
        Index const halfn = n / 2;
        std::nth_element(
            mPermutation.data() + begin,
            mPermutation.data() + begin + halfn,
            mPermutation.data() + begin + n,
            [&](Index const lhs, Index const rhs) {
                return P(dimension, lhs) < P(dimension, rhs);
            });
        // Set the left (and implicitly right) child index of the current node
        mNodes[nodeIdx].c = static_cast<Index>(mNodes.size());
        // Split bounding box into child boxes
        Scalar const split           = P(dimension, mPermutation[begin + halfn]);
        AxisAlignedBoundingBox laabb = aabb;
        laabb.max()(dimension)       = split;
        AxisAlignedBoundingBox raabb = aabb;
        raabb.min()(dimension)       = split;
        // Store left/right node contiguously in memory
        mNodes.emplace_back(begin, halfn, KdTreeNode::kLeafNodeLeftChild);
        mNodes.emplace_back(begin + halfn, n - halfn, KdTreeNode::kLeafNodeLeftChild);
        // Schedule left and right sub-tree constructions
        stack.Push({laabb, mNodes[nodeIdx].Left(), begin, halfn});
        stack.Push({raabb, mNodes[nodeIdx].Right(), begin + halfn, n - halfn});
    }
#include "pbat/warning/Pop.h"
}
 
template <int Dims>
template <class FVisit>
inline void KdTree<Dims>::DepthFirstSearch(FVisit fVisit, Index root) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.KdTree.DepthFirstSearch");
#include "pbat/warning/Push.h"
#include "pbat/warning/SignConversion.h"
    common::TraverseNAryTreePseudoPreOrder(
        [&](Index node) { return fVisit(node, mNodes[node]); },
        [&]<auto c>(Index node) {
            if constexpr (c == 0)
                return mNodes[node].Left();
            else
                return mNodes[node].Right();
        },
        root);
#include "pbat/warning/Pop.h"
}
 
template <int Dims>
inline auto KdTree<Dims>::PointsInNode(Index const nodeIdx) const
{
    auto const& node = mNodes[static_cast<std::size_t>(nodeIdx)];
    return PointsInNode(node);
}
 
template <int Dims>
inline auto KdTree<Dims>::PointsInNode(KdTreeNode const& node) const
{
    namespace vi = std::views;
    auto indrng  = mPermutation | vi::drop(node.begin) | vi::take(node.n);
    return indrng;
}
 
} // namespace geometry
} // namespace pbat
 
#endif // PBAT_GEOMETRY_KDTREE_H