AabbKdTreeHierarchy.h

Go to the documentation of this file.

 
#ifndef PBAT_GEOMETRY_AABBKDTREEHIERARCHY_H
#define PBAT_GEOMETRY_AABBKDTREEHIERARCHY_H
 
#include "KdTree.h"
#include "SpatialSearch.h"
#include "pbat/Aliases.h"
#include "pbat/common/NAryTreeTraversal.h"
#include "pbat/geometry/OverlapQueries.h"
#include "pbat/math/linalg/mini/Eigen.h"
#include "pbat/profiling/Profiling.h"
 
#include <Eigen/Core>
 
namespace pbat::geometry {
template <auto Dims>
class AabbKdTreeHierarchy
{
public:
    static auto constexpr kDims = Dims; 
    using SelfType = AabbKdTreeHierarchy<kDims>; 
 
    AabbKdTreeHierarchy() = default;
    template <class TDerivedL, class TDerivedU>
    AabbKdTreeHierarchy(
        Eigen::DenseBase<TDerivedL> const& L,
        Eigen::DenseBase<TDerivedU> const& U,
        Index maxPointsInLeaf = 8);
    template <class TDerivedL, class TDerivedU>
    void Construct(
        Eigen::DenseBase<TDerivedL> const& L,
        Eigen::DenseBase<TDerivedU> const& U,
        Index maxPointsInLeaf = 8);
    template <class TDerivedL, class TDerivedU>
    void Update(Eigen::DenseBase<TDerivedL> const& L, Eigen::DenseBase<TDerivedU> const& U);
    template <class FNodeOverlaps, class FObjectOverlaps, class FOnOverlap>
    void
    Overlaps(
        FNodeOverlaps fNodeOverlaps,
        FObjectOverlaps fObjectOverlaps,
        FOnOverlap fOnOverlap)
    const;
    template <class FDistanceToNode, class FDistanceToObject, class FOnNearestNeighbour>
    void NearestNeighbours(
        FDistanceToNode fDistanceToNode,
        FDistanceToObject fDistanceToObject,
        FOnNearestNeighbour fOnNearestNeighbour,
        Scalar radius = std::numeric_limits<Scalar>::max(),
        Scalar eps    = Scalar(0)) const;
    template <class FDistanceToNode, class FDistanceToObject, class FOnNearestNeighbour>
    void KNearestNeighbours(
        FDistanceToNode fDistanceToNode,
        FDistanceToObject fDistanceToObject,
        FOnNearestNeighbour fOnNearestNeighbour,
        Index K,
        Scalar radius = std::numeric_limits<Scalar>::max()) const;
    template <class FObjectsOverlap, class FOnSelfOverlap>
    void SelfOverlaps(FObjectsOverlap fObjectsOverlap, FOnSelfOverlap fOnSelfOverlap) const;
    template <class FObjectsOverlap, class FOnOverlap>
    void
    Overlaps(SelfType const& rhs, FObjectsOverlap fObjectsOverlap, FOnOverlap fOnOverlap) const;

    auto InternalNodeBoundingBoxes() const -> Matrix<2 * kDims, Eigen::Dynamic> const&
    {
        return IB;
    }

    auto Tree() const -> KdTree<kDims> const& { return tree; }
    auto Lower(Index node) const { return IB.col(node).template head<kDims>(); }
    auto Upper(Index node) const { return IB.col(node).template tail<kDims>(); }
 
private:
    Matrix<2 * kDims, Eigen::Dynamic>
    IB;
    KdTree<kDims> tree; 
};
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline AabbKdTreeHierarchy<Dims>::AabbKdTreeHierarchy(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U,
    Index maxPointsInLeaf)
{
    Construct(L, U, maxPointsInLeaf);
}
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline void AabbKdTreeHierarchy<Dims>::Construct(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U,
    Index maxPointsInLeaf)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.Construct");
    Matrix<kDims, Eigen::Dynamic> P = 0.5 * (L.derived() + U.derived());
    tree.Construct(P, maxPointsInLeaf);
    auto const nNodes = static_cast<Index>(tree.Nodes().size());
    IB.resize(2 * kDims, nNodes);
    IB.template topRows<kDims>().setConstant(std::numeric_limits<Scalar>::max());
    IB.template bottomRows<kDims>().setConstant(std::numeric_limits<Scalar>::lowest());
}
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline void AabbKdTreeHierarchy<Dims>::Update(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.Update");
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    common::TraverseNAryTreePseudoPostOrder(
        [&](Index n) {
            KdTreeNode const& node = nodes[n];
            if (node.IsLeaf())
            {
                auto inds                        = perm(Eigen::seqN(node.begin, node.n));
                IB.col(n).template head<kDims>() = L(Eigen::placeholders::all, inds).rowwise().
                    minCoeff();
                IB.col(n).template tail<kDims>() = U(Eigen::placeholders::all, inds).rowwise().
                    maxCoeff();
            }
            else
            {
                // Our k-D tree's internal nodes always have both children.
                auto nbox                   = IB.col(n);
                auto lbox                   = IB.col(node.Left());
                auto rbox                   = IB.col(node.Right());
                nbox.template head<kDims>() = lbox.template head<kDims>().cwiseMin(
                    rbox.template head<kDims>());
                nbox.template tail<kDims>() = lbox.template tail<kDims>().cwiseMax(
                    rbox.template tail<kDims>());
            }
        },
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return nodes[n].Left();
            else
                return nodes[n].Right();
        });
}
 
template <auto Dims>
template <class FNodeOverlaps, class FObjectOverlaps, class FOnOverlap>
inline void AabbKdTreeHierarchy<Dims>::Overlaps(
    FNodeOverlaps fNodeOverlaps,
    FObjectOverlaps fObjectOverlaps,
    FOnOverlap fOnOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.Overlaps");
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    geometry::Overlaps(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return nodes[n].Left();
            else
                return nodes[n].Right();
        },
        [&](Index n) { return nodes[n].IsLeaf(); },
        [&](Index n) { return nodes[n].n; },
        [&](Index n, Index i) { return perm(nodes[n].begin + i); },
        [&](Index n) {
            auto L          = IB.col(n).template head<kDims>();
            auto U          = IB.col(n).template tail<kDims>();
            using TDerivedL = decltype(L);
            using TDerivedU = decltype(U);
            return fNodeOverlaps.template operator()<TDerivedL, TDerivedU>(L, U);
        },
        fObjectOverlaps,
        fOnOverlap);
}
 
template <auto Dims>
template <class FDistanceToNode, class FDistanceToObject, class FOnNearestNeighbour>
inline void AabbKdTreeHierarchy<Dims>::NearestNeighbours(
    FDistanceToNode fDistanceToNode,
    FDistanceToObject fDistanceToObject,
    FOnNearestNeighbour fOnNearestNeighbour,
    Scalar radius,
    Scalar eps) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.NearestNeighbours");
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    geometry::DfsAllNearestNeighbours(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return nodes[n].Left();
            else
                return nodes[n].Right();
        },
        [&](Index n) { return nodes[n].IsLeaf(); },
        [&](Index n) { return nodes[n].n; },
        [&](Index n, Index i) { return perm(nodes[n].begin + i); },
        [&](Index n) {
            auto L          = IB.col(n).template head<kDims>();
            auto U          = IB.col(n).template tail<kDims>();
            using TDerivedL = decltype(L);
            using TDerivedU = decltype(U);
            return fDistanceToNode.template operator()<TDerivedL, TDerivedU>(L, U);
        },
        fDistanceToObject,
        fOnNearestNeighbour,
        false /*bUseBestFirstSearch*/,
        radius,
        eps);
}
 
template <auto Dims>
template <class FDistanceToNode, class FDistanceToObject, class FOnNearestNeighbour>
inline void AabbKdTreeHierarchy<Dims>::KNearestNeighbours(
    FDistanceToNode fDistanceToNode,
    FDistanceToObject fDistanceToObject,
    FOnNearestNeighbour fOnNearestNeighbour,
    Index K,
    Scalar radius) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.KNearestNeighbours");
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    geometry::KNearestNeighbours(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return nodes[n].Left();
            else
                return nodes[n].Right();
        },
        [&](Index n) { return nodes[n].IsLeaf(); },
        [&](Index n) { return nodes[n].n; },
        [&](Index n, Index i) { return perm(nodes[n].begin + i); },
        [&](Index n) {
            auto L          = IB.col(n).template head<kDims>();
            auto U          = IB.col(n).template tail<kDims>();
            using TDerivedL = decltype(L);
            using TDerivedU = decltype(U);
            return fDistanceToNode.template operator()<TDerivedL, TDerivedU>(L, U);
        },
        fDistanceToObject,
        fOnNearestNeighbour,
        K,
        radius);
}
 
template <auto Dims>
template <class FObjectsOverlap, class FOnSelfOverlap>
inline void AabbKdTreeHierarchy<Dims>::SelfOverlaps(
    FObjectsOverlap fObjectsOverlap,
    FOnSelfOverlap fOnSelfOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.SelfOverlaps");
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    geometry::SelfOverlaps(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return nodes[n].Left();
            else
                return nodes[n].Right();
        },
        [&](Index n) { return nodes[n].IsLeaf(); },
        [&](Index n) { return nodes[n].n; },
        [&](Index n, Index i) { return perm(nodes[n].begin + i); },
        [&](Index n1, Index n2) {
            using math::linalg::mini::FromEigen;
            auto L1 = IB.col(n1).template head<kDims>();
            auto U1 = IB.col(n1).template tail<kDims>();
            auto L2 = IB.col(n2).template head<kDims>();
            auto U2 = IB.col(n2).template tail<kDims>();
            return geometry::OverlapQueries::AxisAlignedBoundingBoxes(
                FromEigen(L1),
                FromEigen(U1),
                FromEigen(L2),
                FromEigen(U2));
        },
        fObjectsOverlap,
        fOnSelfOverlap);
}
 
template <auto Dims>
template <class FObjectsOverlap, class FOnOverlap>
inline void AabbKdTreeHierarchy<Dims>::Overlaps(
    SelfType const& rhs,
    FObjectsOverlap fObjectsOverlap,
    FOnOverlap fOnOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbKdTreeHierarchy.Overlaps");
    // This tree will be the left-hand side tree
    KdTreeNode const* nodes  = tree.Nodes().data();
    IndexVectorX const& perm = tree.Permutation();
    auto const fChildLhs     = [&]<auto c>(Index n) -> Index {
        if constexpr (c == 0)
            return nodes[n].Left();
        else
            return nodes[n].Right();
    };
    auto const fIsLeafLhs = [&](Index n) {
        return nodes[n].IsLeaf();
    };
    auto const fLeafSizeLhs = [&]([[maybe_unused]] Index n) {
        return nodes[n].n;
    };
    auto const fLeafObjectLhs = [&](Index n, [[maybe_unused]] Index i) {
        return perm(nodes[n].begin + i);
    };
    // This tree will be the right-hand side tree
    KdTreeNode const* nodesRhs  = rhs.tree.Nodes().data();
    IndexVectorX const& permRhs = rhs.tree.Permutation();
    auto const fChildRhs        = [&]<auto c>(Index n) -> Index {
        if constexpr (c == 0)
            return nodesRhs[n].Left();
        else
            return nodesRhs[n].Right();
    };
    auto const fIsLeafRhs = [&](Index n) {
        return nodesRhs[n].IsLeaf();
    };
    auto const fLeafSizeRhs = [&]([[maybe_unused]] Index n) {
        return nodesRhs[n].n;
    };
    auto const fLeafObjectRhs = [&](Index n, [[maybe_unused]] Index i) {
        return permRhs(nodesRhs[n].begin + i);
    };
    // Register overlaps
    geometry::Overlaps(
        fChildLhs,
        fIsLeafLhs,
        fLeafSizeLhs,
        fLeafObjectLhs,
        fChildRhs,
        fIsLeafRhs,
        fLeafSizeRhs,
        fLeafObjectRhs,
        [&](Index n1, Index n2) {
            using math::linalg::mini::FromEigen;
            auto L1 = IB.col(n1).template head<kDims>();
            auto U1 = IB.col(n1).template tail<kDims>();
            auto L2 = rhs.IB.col(n2).template head<kDims>();
            auto U2 = rhs.IB.col(n2).template tail<kDims>();
            return geometry::OverlapQueries::AxisAlignedBoundingBoxes(
                FromEigen(L1),
                FromEigen(U1),
                FromEigen(L2),
                FromEigen(U2));
        },
        fObjectsOverlap,
        fOnOverlap);
}
} // namespace pbat::geometry
 
#endif // PBAT_GEOMETRY_AABBKDTREEHIERARCHY_H