HierarchicalHashGrid.h

#ifndef PBAT_GEOMETRY_HIERARCHICALHASHGRID_H
#define PBAT_GEOMETRY_HIERARCHICALHASHGRID_H
 
#include "pbat/Aliases.h"
#include "pbat/common/BruteSet.h"
#include "pbat/common/Concepts.h"
#include "pbat/common/Eigen.h"
#include "pbat/common/Modulo.h"
#include "pbat/geometry/OverlapQueries.h"
#include "pbat/math/linalg/mini/Eigen.h"
#include "pbat/profiling/Profiling.h"
 
#include <Eigen/Core>
#include <algorithm>
#include <cassert>
#include <limits>
#include <span>
 
namespace pbat::geometry {
template <int Dims, common::CFloatingPoint TScalar = Scalar, common::CIndex TIndex = Index>
class HierarchicalHashGrid
{
public:
    using ScalarType           = TScalar; 
    using IndexType            = TIndex;  
    static constexpr int kDims = Dims;    
    static_assert(Dims >= 2 and Dims <= 3, "HierarchicalHashGrid only supports 2D and 3D");
    static constexpr int kMaxCellsPerPrimitive =
        Dims == 2 ? 4 : 8; 

    HierarchicalHashGrid() = default;
    HierarchicalHashGrid(IndexType nPrimitives, IndexType nBuckets = 0);
    void Configure(IndexType nPrimitives, IndexType nBuckets = 0);
    template <class TDerivedL, class TDerivedU>
    void Construct(Eigen::DenseBase<TDerivedL> const& L, Eigen::DenseBase<TDerivedU> const& U);
    template <class FOnPair, class TDerivedX>
    void BroadPhase(Eigen::DenseBase<TDerivedX> const& X, FOnPair fOnPair) const;
    template <class FOnPair, class TDerivedL, class TDerivedU>
    void BroadPhase(
        Eigen::DenseBase<TDerivedL> const& L,
        Eigen::DenseBase<TDerivedU> const& U,
        FOnPair fOnPair) const;
    template <class FOnPair, class TDerivedLP, class TDerivedUP, class TDerivedX>
    void BroadPhase(
        Eigen::DenseBase<TDerivedLP> const& LP,
        Eigen::DenseBase<TDerivedUP> const& UP,
        Eigen::DenseBase<TDerivedX> const& XQ,
        FOnPair fOnPair) const;
    template <class FOnPair, class TDerivedLP, class TDerivedUP, class TDerivedLQ, class TDerivedUQ>
    void BroadPhase(
        Eigen::DenseBase<TDerivedLP> const& LP,
        Eigen::DenseBase<TDerivedUP> const& UP,
        Eigen::DenseBase<TDerivedLQ> const& LQ,
        Eigen::DenseBase<TDerivedUQ> const& UQ,
        FOnPair fOnPair) const;
    IndexType NumberOfBuckets() const { return static_cast<IndexType>(mPrefix.size()) - 1; }
    IndexType NumberOfLevels() const { return static_cast<IndexType>(mSetOfLevels.Size()); }
    auto Levels() const -> std::span<std::int16_t const>
    {
        return std::span<std::int16_t const>(mSetOfLevels.begin(), mSetOfLevels.Size());
    }

    template <class TDerivedX>
    auto ToIntegerCoordinates(
        Eigen::DenseBase<TDerivedX> const& X,
        ScalarType const cellSize) const
        -> Eigen::Vector<IndexType, kDims>;
    template <class TDerivedX>
    auto Hash(Eigen::DenseBase<TDerivedX> const& X, std::int16_t l) const;
 
protected:
    void ClearHashTable();
 
private:
    Eigen::Vector<std::uint8_t, Eigen::Dynamic>
    mCellCounts; 
    Eigen::Vector<IndexType, Eigen::Dynamic>
    mBucketIds;
    Eigen::Vector<IndexType, Eigen::Dynamic>
    mPrefix;
    Eigen::Vector<IndexType, Eigen::Dynamic>
    mPrimitiveIds;
    common::BruteSet<std::int16_t> mSetOfLevels; 
};
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
HierarchicalHashGrid<Dims, TScalar, TIndex>::HierarchicalHashGrid(
    IndexType nPrimitives,
    IndexType nBuckets)
    : mCellCounts(),
      mBucketIds(),
      mPrefix(),
      mPrimitiveIds(),
      mSetOfLevels()
{
    Configure(nPrimitives, nBuckets);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
inline void
HierarchicalHashGrid<Dims, TScalar, TIndex>::Configure(IndexType nPrimitives, IndexType nBuckets)
{
    nBuckets = std::max(nPrimitives * kMaxCellsPerPrimitive, nBuckets);
    mCellCounts.resize(nPrimitives);
    mBucketIds.resize(nBuckets);
    mPrefix.resize(nBuckets + 1);
    mPrimitiveIds.resize(nBuckets);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class TDerivedL, class TDerivedU>
void HierarchicalHashGrid<Dims, TScalar, TIndex>::Construct(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.HierarchicalHashGrid.Construct");
    assert(L.rows() == kDims and U.rows() == kDims); // L and U must have |# dims| rows
    assert(L.cols() == U.cols());                    // L and U must have the same number of columns
    auto const nPrimitives = static_cast<IndexType>(L.cols());
    auto const nBuckets    = NumberOfBuckets();
    assert(nPrimitives <= nBuckets);
    assert(mPrefix.size() == nBuckets + 1); // Prefix size must be nBuckets + 1
    // Build hash table from scratch
    ClearHashTable();
    // Map primitives to cells
    IndexType k = 0; // Index for mBucketIds and mPrimitiveIds
    for (IndexType i = 0; i < nPrimitives; ++i)
    {
        ScalarType const maxExtent               = (U.col(i) - L.col(i)).maxCoeff();
        ScalarType const l                       = std::ceil(std::log2(maxExtent));
        ScalarType const cellSize                = std::pow(ScalarType(2), l);
        Eigen::Vector<IndexType, kDims> const ib =
            ToIntegerCoordinates(L.col(i).template segment<kDims>(0), cellSize);
        Eigen::Vector<IndexType, kDims> const ie =
            ToIntegerCoordinates(U.col(i).template segment<kDims>(0), cellSize);
        auto const level = static_cast<std::int16_t>(l);
        Eigen::Vector<IndexType, kDims> ix;
        // Add primitive i to every cell at level l overlapping with i's bounding box
        for (ix(0) = ib(0); ix(0) <= ie(0); ++ix(0))
        {
            for (ix(1) = ib(1); ix(1) <= ie(1); ++ix(1))
            {
                if constexpr (kDims == 2)
                {
                    IndexType const hash = Hash(ix, level);
                    mBucketIds(k++)      = common::Modulo(hash, nBuckets);
                    ++mCellCounts(i);
                }
                else
                {
                    for (ix(2) = ib(2); ix(2) <= ie(2); ++ix(2))
                    {
                        IndexType const hash = Hash(ix, level);
                        mBucketIds(k++)      = common::Modulo(hash, nBuckets);
                        ++mCellCounts(i);
                    }
                }
            }
        }
        mSetOfLevels.Insert(level);
    }
    // Compute id counts in the prefix sum's memory
    mPrefix(mBucketIds.segment(0, k)).array() += 1;
    // Compute the shifted prefix sum
    std::inclusive_scan(mPrefix.begin(), mPrefix.end(), mPrefix.begin());
    // Construct primitive IDs while unshifting prefix sum
    k = 0;
    for (IndexType i = 0; i < nPrimitives; ++i)
    {
        for (std::uint8_t j = 0; j < mCellCounts(i); ++j, ++k)
        {
            auto bucketId                      = mBucketIds(k);
            mPrimitiveIds(--mPrefix(bucketId)) = i;
        }
    }
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FOnPair, class TDerivedX>
void HierarchicalHashGrid<Dims, TScalar, TIndex>::BroadPhase(
    Eigen::DenseBase<TDerivedX> const& X,
    FOnPair fOnPair) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.HierarchicalHashGrid.BroadPhase");
    assert(X.rows() == kDims); // X must have |# dims| rows
    auto const nQueries = static_cast<IndexType>(X.cols());
    auto const nBuckets = NumberOfBuckets();
    for (IndexType q = 0; q < nQueries; ++q)
    {
        for (std::int16_t l : mSetOfLevels)
        {
            auto cellSize = std::pow(ScalarType(2), static_cast<ScalarType>(l));
            Eigen::Vector<IndexType, kDims> iq = ToIntegerCoordinates(X.col(q), cellSize);
            auto bucketId = (kDims == 2) ?
                                common::Modulo(Hash(iq, l), nBuckets) :
                                common::Modulo(Hash(iq, l), nBuckets);
            auto beginBucket = mPrefix(bucketId);
            auto endBucket   = mPrefix(bucketId + 1);
            for (auto k = beginBucket; k < endBucket; ++k)
            {
                auto const i = mPrimitiveIds(k);
                fOnPair(q, i);
            }
        }
    }
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FOnPair, class TDerivedL, class TDerivedU>
inline void HierarchicalHashGrid<Dims, TScalar, TIndex>::BroadPhase(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U,
    FOnPair fOnPair) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.HierarchicalHashGrid.BroadPhase");
    assert(L.rows() == kDims);    // L must have |# dims| rows
    assert(U.rows() == kDims);    // U must have |# dims| rows
    assert(L.cols() == U.cols()); // L and U must have the same number of columns
    auto const nQueries = static_cast<IndexType>(L.cols());
    auto const nBuckets = NumberOfBuckets();
    for (IndexType q = 0; q < nQueries; ++q)
    {
        for (std::int16_t l : mSetOfLevels)
        {
            auto cellSize = std::pow(ScalarType(2), static_cast<ScalarType>(l));
            Eigen::Vector<IndexType, kDims> iqb = ToIntegerCoordinates(L.col(q), cellSize);
            Eigen::Vector<IndexType, kDims> iqe = ToIntegerCoordinates(U.col(q), cellSize);
            Eigen::Vector<IndexType, kDims> iq;
            for (iq(0) = iqb(0); iq(0) <= iqe(0); ++iq(0))
            {
                for (iq(1) = iqb(1); iq(1) <= iqe(1); ++iq(1))
                {
                    if constexpr (kDims == 2)
                    {
                        auto bucketId    = common::Modulo(Hash(iq, l), nBuckets);
                        auto beginBucket = mPrefix(bucketId);
                        auto endBucket   = mPrefix(bucketId + 1);
                        for (auto k = beginBucket; k < endBucket; ++k)
                        {
                            auto const i = mPrimitiveIds(k);
                            fOnPair(q, i);
                        }
                    }
                    else
                    {
                        for (iq(2) = iqb(2); iq(2) <= iqe(2); ++iq(2))
                        {
                            auto bucketId    = common::Modulo(Hash(iq, l), nBuckets);
                            auto beginBucket = mPrefix(bucketId);
                            auto endBucket   = mPrefix(bucketId + 1);
                            for (auto k = beginBucket; k < endBucket; ++k)
                            {
                                auto const i = mPrimitiveIds(k);
                                fOnPair(q, i);
                            }
                        }
                    }
                }
            }
        }
    }
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FOnPair, class TDerivedLP, class TDerivedUP, class TDerivedX>
inline void HierarchicalHashGrid<Dims, TScalar, TIndex>::BroadPhase(
    Eigen::DenseBase<TDerivedLP> const& LP,
    Eigen::DenseBase<TDerivedUP> const& UP,
    Eigen::DenseBase<TDerivedX> const& XQ,
    FOnPair fOnPair) const
{
    BroadPhase(
        XQ,
        [&](Index q, Index p) {
            using pbat::math::linalg::mini::FromEigen;
            if (OverlapQueries::PointAxisAlignedBoundingBox(
                FromEigen(XQ.col(q).template head<kDims>()),
                FromEigen(LP.col(p).template head<kDims>()),
                FromEigen(UP.col(p).template head<kDims>())))
            {
                fOnPair(q, p);
            }
        });
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FOnPair, class TDerivedLP, class TDerivedUP, class TDerivedLQ, class TDerivedUQ>
inline void HierarchicalHashGrid<Dims, TScalar, TIndex>::BroadPhase(
    Eigen::DenseBase<TDerivedLP> const& LP,
    Eigen::DenseBase<TDerivedUP> const& UP,
    Eigen::DenseBase<TDerivedLQ> const& LQ,
    Eigen::DenseBase<TDerivedUQ> const& UQ,
    FOnPair fOnPair) const
{
    using pbat::math::linalg::mini::FromEigen;
    BroadPhase(
        LQ,
        UQ,
        [&](Index q, Index p) {
            if (OverlapQueries::AxisAlignedBoundingBoxes(
                FromEigen(LP.col(p).template head<kDims>()),
                FromEigen(UP.col(p).template head<kDims>()),
                FromEigen(LQ.col(q).template head<kDims>()),
                FromEigen(UQ.col(q).template head<kDims>())))
            {
                fOnPair(q, p);
            }
        });
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class TDerivedX>
inline auto HierarchicalHashGrid<Dims, TScalar, TIndex>::ToIntegerCoordinates(
    Eigen::DenseBase<TDerivedX> const& X,
    ScalarType const cellSize) const -> Eigen::Vector<IndexType, kDims>
{
    return Eigen::Vector<ScalarType, kDims>(X.derived().array() / static_cast<ScalarType>(cellSize))
           .array()
           .floor()
           .template cast<IndexType>();
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class TDerivedX>
inline auto HierarchicalHashGrid<Dims, TScalar, TIndex>::Hash(
    Eigen::DenseBase<TDerivedX> const& X,
    std::int16_t l) const
{
    if constexpr (kDims == 2)
        return (X(0) * 73856093) ^ (X(1) * 19349663) ^ (l * 67867979);
    else
        return (X(0) * 73856093) ^ (X(1) * 19349663) ^ (X(2) * 834927911) ^ (l * 67867979);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
void HierarchicalHashGrid<Dims, TScalar, TIndex>::ClearHashTable()
{
    mCellCounts.setZero();
    mPrimitiveIds.setConstant(IndexType(-1));
    mPrefix.setZero();
    mSetOfLevels.Clear();
}
} // namespace pbat::geometry
 
#endif // PBAT_GEOMETRY_HIERARCHICALHASHGRID_H