AabbRadixTreeHierarchy.h

#ifndef PBAT_GEOMETRY_AABBRADIXTREEHIERARCHY_H
#define PBAT_GEOMETRY_AABBRADIXTREEHIERARCHY_H
 
#include "AxisAlignedBoundingBox.h"
#include "SpatialSearch.h"
#include "pbat/Aliases.h"
#include "pbat/common/BinaryRadixTree.h"
#include "pbat/common/ConstexprFor.h"
#include "pbat/common/NAryTreeTraversal.h"
#include "pbat/common/Permute.h"
#include "pbat/geometry/Morton.h"
#include "pbat/geometry/OverlapQueries.h"
#include "pbat/math/linalg/mini/Eigen.h"
#include "pbat/profiling/Profiling.h"
 
#include <cpp-sort/sorters/ska_sorter.h>
#include <numeric>
 
namespace pbat::geometry {
template <auto Dims>
class AabbRadixTreeHierarchy
{
public:
    using IndexType = Index; 
    static auto constexpr kDims = Dims; 
    using SelfType = AabbRadixTreeHierarchy<kDims>; 
 
    AabbRadixTreeHierarchy() = default;
    template <class TDerivedL, class TDerivedU>
    AabbRadixTreeHierarchy(
        Eigen::DenseBase<TDerivedL> const& L,
        Eigen::DenseBase<TDerivedU> const& U);
    template <class TDerivedL, class TDerivedU>
    void Construct(Eigen::DenseBase<TDerivedL> const& L, Eigen::DenseBase<TDerivedU> const& U);
    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 { return IB; }
    auto Tree() const -> common::BinaryRadixTree<IndexType> const& { return tree; }
 
protected:
    template <class TDerivedL, class TDerivedU>
    void
    ComputeMortonCodes(Eigen::DenseBase<TDerivedL> const& L, Eigen::DenseBase<TDerivedU> const& U);
    void SortMortonCodes();
 
private:
    Eigen::Vector<MortonCodeType, Eigen::Dynamic> codes; 
    Eigen::Vector<IndexType, Eigen::Dynamic> inds;       
    Matrix<2 * kDims, Eigen::Dynamic>
    IB;
    common::BinaryRadixTree<IndexType> tree; 
};
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline AabbRadixTreeHierarchy<Dims>::AabbRadixTreeHierarchy(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    Construct(L.derived(), U.derived());
}
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline void AabbRadixTreeHierarchy<Dims>::Construct(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.Construct");
    auto const n = L.cols();
    codes.resize(n);
    inds.resize(n);
    ComputeMortonCodes(L.derived(), U.derived());
    SortMortonCodes();
    tree.Construct(codes);
    IB.resize(2 * kDims, tree.InternalNodeCount());
}
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline void AabbRadixTreeHierarchy<Dims>::Update(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.Update");
    IB.template topRows<kDims>().setConstant(std::numeric_limits<Scalar>::max());
    IB.template bottomRows<kDims>().setConstant(std::numeric_limits<Scalar>::lowest());
    common::TraverseNAryTreePseudoPostOrder(
        [&](Index n) {
            // n is guaranteed to be internal node, due to our fChild functor
            auto const lc = tree.Left(n);
            auto const rc = tree.Right(n);
            Vector<kDims> LL, LU, RL, RU;
            if (tree.IsLeaf(lc))
            {
                auto i = inds(tree.CodeIndex(lc));
                LL     = L.col(i).template head<kDims>();
                LU     = U.col(i).template head<kDims>();
            }
            else
            {
                LL = IB.col(lc).template head<kDims>();
                LU = IB.col(lc).template tail<kDims>();
            }
            if (tree.IsLeaf(rc))
            {
                auto i = inds(tree.CodeIndex(rc));
                RL     = L.col(i).template head<kDims>();
                RU     = U.col(i).template head<kDims>();
            }
            else
            {
                RL = IB.col(rc).template head<kDims>();
                RU = IB.col(rc).template tail<kDims>();
            }
            IB.col(n).template head<kDims>() = LL.cwiseMin(RL);
            IB.col(n).template tail<kDims>() = LU.cwiseMax(RU);
        },
        // fChild functor that only returns non-leaf children
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return tree.IsLeaf(tree.Left(n)) ? -1 : tree.Left(n);
            else
                return tree.IsLeaf(tree.Right(n)) ? -1 : tree.Right(n);
        });
}
 
template <auto Dims>
template <class FNodeOverlaps, class FObjectOverlaps, class FOnOverlap>
inline void AabbRadixTreeHierarchy<Dims>::Overlaps(
    FNodeOverlaps fNodeOverlaps,
    FObjectOverlaps fObjectOverlaps,
    FOnOverlap fOnOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.Overlaps");
    geometry::Overlaps(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return tree.IsLeaf(n) ? -1 : tree.Left(n);
            else
                return tree.IsLeaf(n) ? -1 : tree.Right(n);
        },
        [&](Index n) { return tree.IsLeaf(n); },
        []([[maybe_unused]] Index n) { return 1; },
        [&](Index n, [[maybe_unused]] Index i) { return inds(tree.CodeIndex(n)); },
        [&](Index n) {
            if (tree.IsLeaf(n))
                return true; // Radix tree leaf nodes correspond to individual objects
            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 AabbRadixTreeHierarchy<Dims>::NearestNeighbours(
    FDistanceToNode fDistanceToNode,
    FDistanceToObject fDistanceToObject,
    FOnNearestNeighbour fOnNearestNeighbour,
    Scalar radius,
    Scalar eps) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.NearestNeighbours");
    geometry::DfsAllNearestNeighbours(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return tree.IsLeaf(n) ? -1 : tree.Left(n);
            else
                return tree.IsLeaf(n) ? -1 : tree.Right(n);
        },
        [&](Index n) { return tree.IsLeaf(n); },
        []([[maybe_unused]] Index n) { return 1; },
        [&](Index n, [[maybe_unused]] Index i) { return inds(tree.CodeIndex(n)); },
        [&](Index n) {
            using TDerivedL  = decltype(IB.col(n).template head<kDims>());
            using TDerivedU  = decltype(IB.col(n).template tail<kDims>());
            using ScalarType = std::invoke_result_t<FDistanceToNode, TDerivedL, TDerivedU>;
            if (tree.IsLeaf(n))
                return ScalarType(0); // Radix tree leaf nodes correspond to individual objects
            auto L = IB.col(n).template head<kDims>();
            auto U = IB.col(n).template tail<kDims>();
            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 AabbRadixTreeHierarchy<Dims>::KNearestNeighbours(
    FDistanceToNode fDistanceToNode,
    FDistanceToObject fDistanceToObject,
    FOnNearestNeighbour fOnNearestNeighbour,
    Index K,
    Scalar radius) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.KNearestNeighbours");
    geometry::KNearestNeighbours(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return tree.IsLeaf(n) ? -1 : tree.Left(n);
            else
                return tree.IsLeaf(n) ? -1 : tree.Right(n);
        },
        [&](Index n) { return tree.IsLeaf(n); },
        []([[maybe_unused]] Index n) { return 1; },
        [&](Index n, [[maybe_unused]] Index i) { return inds(tree.CodeIndex(n)); },
        [&](Index n) {
            using TDerivedL  = decltype(IB.col(n).template head<kDims>());
            using TDerivedU  = decltype(IB.col(n).template tail<kDims>());
            using ScalarType = std::invoke_result_t<FDistanceToNode, TDerivedL, TDerivedU>;
            if (tree.IsLeaf(n))
                return ScalarType(0); // Radix tree leaf nodes correspond to individual objects
            auto L = IB.col(n).template head<kDims>();
            auto U = IB.col(n).template tail<kDims>();
            return fDistanceToNode.template operator()<TDerivedL, TDerivedU>(L, U);
        },
        fDistanceToObject,
        fOnNearestNeighbour,
        K,
        radius);
}
 
template <auto Dims>
template <class FObjectsOverlap, class FOnSelfOverlap>
inline void AabbRadixTreeHierarchy<Dims>::SelfOverlaps(
    FObjectsOverlap fObjectsOverlap,
    FOnSelfOverlap fOnSelfOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.SelfOverlaps");
    geometry::SelfOverlaps(
        [&]<auto c>(Index n) -> Index {
            if constexpr (c == 0)
                return tree.IsLeaf(n) ? -1 : tree.Left(n);
            else
                return tree.IsLeaf(n) ? -1 : tree.Right(n);
        },
        [&](Index n) { return tree.IsLeaf(n); },
        []([[maybe_unused]] Index n) { return 1; },
        [&](Index n, [[maybe_unused]] Index i) { return inds(tree.CodeIndex(n)); },
        [&](Index n1, Index n2) {
            if (tree.IsLeaf(n1) or tree.IsLeaf(n2))
                return true; // Radix tree leaf nodes correspond to individual objects
            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 AabbRadixTreeHierarchy<Dims>::Overlaps(
    SelfType const& rhs,
    FObjectsOverlap fObjectsOverlap,
    FOnOverlap fOnOverlap) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.Overlaps");
    // This tree will be the left-hand side tree
    auto const fChildLhs = [&]<auto c>(Index n) -> Index {
        if constexpr (c == 0)
            return tree.IsLeaf(n) ? -1 : tree.Left(n);
        else
            return tree.IsLeaf(n) ? -1 : tree.Right(n);
    };
    auto const fIsLeafLhs = [&](Index n) {
        return tree.IsLeaf(n);
    };
    auto const fLeafSizeLhs = []([[maybe_unused]] Index n) {
        return 1;
    };
    auto const fLeafObjectLhs = [&](Index n, [[maybe_unused]] Index i) {
        return inds(tree.CodeIndex(n));
    };
    // This tree will be the right-hand side tree
    auto const fChildRhs = [&]<auto c>(Index n) -> Index {
        if constexpr (c == 0)
            return rhs.tree.IsLeaf(n) ? -1 : rhs.tree.Left(n);
        else
            return rhs.tree.IsLeaf(n) ? -1 : rhs.tree.Right(n);
    };
    auto const fIsLeafRhs = [&](Index n) {
        return rhs.tree.IsLeaf(n);
    };
    auto const fLeafSizeRhs = []([[maybe_unused]] Index n) {
        return 1;
    };
    auto const fLeafObjectRhs = [&](Index n, [[maybe_unused]] Index i) {
        return rhs.inds(rhs.tree.CodeIndex(n));
    };
    // Register overlaps
    geometry::Overlaps(
        fChildLhs,
        fIsLeafLhs,
        fLeafSizeLhs,
        fLeafObjectLhs,
        fChildRhs,
        fIsLeafRhs,
        fLeafSizeRhs,
        fLeafObjectRhs,
        [&](Index n1, Index n2) {
            if (tree.IsLeaf(n1) or rhs.tree.IsLeaf(n2))
                return true; // Radix tree leaf nodes correspond to individual objects
            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);
}
 
template <auto Dims>
template <class TDerivedL, class TDerivedU>
inline void AabbRadixTreeHierarchy<Dims>::ComputeMortonCodes(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.ComputeMortonCodes");
    Vector<kDims> min     = L.rowwise().minCoeff();
    Vector<kDims> max     = U.rowwise().maxCoeff();
    Vector<kDims> extents = max - min;
    auto C                = 0.5 * (L.derived() + U.derived());
    for (Index i = 0; i < L.cols(); ++i)
    {
        Vector<kDims> center;
        common::ForRange<0, kDims>([&]<auto d>() { center(d) = (C(d, i) - min(d)) / extents(d); });
        codes(i) = Morton3D(center);
    }
}
 
template <auto Dims>
inline void AabbRadixTreeHierarchy<Dims>::SortMortonCodes()
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.AabbRadixTreeHierarchy.SortMortonCodes");
    // NOTE:
    // It would potentially be more efficient not to re-sort from scratch every time, due to
    // this std::iota call. If we kept the inds ordering between calls, we would benefit from
    // significant performance speedups of certain sorting algorithms on nearly sorted input.
    // However, this would mean that in the call to ComputeMortonCodes, we would have to
    // recompute the Morton codes in the same order as the previous call, which will not
    // exploit cache locality of iterating through B's columns. This is a trade-off that
    // should be considered if performance is critical.
    std::iota(inds.begin(), inds.end(), 0);
    cppsort::ska_sort(inds.begin(), inds.end(), [&](IndexType i) { return codes(i); });
    common::Permute(codes.begin(), codes.end(), inds.begin());
}
} // namespace pbat::geometry
 
#endif // PBAT_GEOMETRY_AABBRADIXTREEHIERARCHY_H