HashGrid.h

Go to the documentation of this file.

 
#ifndef PBAT_GEOMETRY_HASHGRID_H
#define PBAT_GEOMETRY_HASHGRID_H
 
#include "pbat/common/Concepts.h"
#include "pbat/common/Modulo.h"
#include "pbat/profiling/Profiling.h"
 
#include <Eigen/Core>
#include <algorithm>
#include <cassert>
 
namespace pbat::geometry {

template <common::CIndex THashed>
[[maybe_unused]] inline auto HashByXorOfPrimeMultiples()
{
    return []<class TDerivedX>(Eigen::DenseBase<TDerivedX> const& X) {
        static_assert(
            TDerivedX::SizeAtCompileTime == 2 or TDerivedX::SizeAtCompileTime == 3,
            "Only eigen vectors of compile-time size 2 or 3 are supported.");
        if constexpr (TDerivedX::SizeAtCompileTime == 2)
        {
            return static_cast<THashed>(X(0) * 73856093) ^ static_cast<THashed>(X(1) * 19349663);
        }
        if constexpr (TDerivedX::SizeAtCompileTime == 3)
        {
            return static_cast<THashed>(X(0) * 73856093) ^ static_cast<THashed>(X(1) * 19349663) ^
                   static_cast<THashed>(X(2) * 83492791);
        }
    };
}

template <int Dims, common::CFloatingPoint TScalar = Scalar, common::CIndex TIndex = Index>
class HashGrid
{
  public:
    using ScalarType           = TScalar; 
    using IndexType            = TIndex;  
    static constexpr int kDims = Dims;    
    static_assert(Dims >= 2 and Dims <= 3, "HashGrid only supports 2D and 3D");

    HashGrid() = default;
    HashGrid(ScalarType cellSize, IndexType nBuckets);
    void Configure(ScalarType cellSize, IndexType nBuckets);
    void Reserve(IndexType nPrimitives);
    template <class FHash, class TDerivedL, class TDerivedU>
    void Construct(
        Eigen::DenseBase<TDerivedL> const& L,
        Eigen::DenseBase<TDerivedU> const& U,
        FHash fHash);
    template <class FHash, class TDerivedX>
    void Construct(Eigen::DenseBase<TDerivedX> const& X, FHash fHash);
    template <class FOnPair, class FHash, class TDerivedX>
    void BroadPhase(Eigen::DenseBase<TDerivedX> const& X, FOnPair fOnPair, FHash fHash) const;
    IndexType NumberOfBuckets() const { return static_cast<IndexType>(mPrefix.size()) - 1; }
    template <class TDerivedX>
    auto Cell(Eigen::DenseBase<TDerivedX> const& X) const -> Eigen::Matrix<ScalarType, kDims, 2>;
    template <class TDerivedX>
    auto ToIntegerCoordinates(Eigen::DenseBase<TDerivedX> const& X) const
        -> Eigen::Vector<IndexType, kDims>;
 
  private:
    ScalarType mCellSize; 
    Eigen::Vector<IndexType, Eigen::Dynamic>
        mBucketIds; 
    Eigen::Vector<IndexType, Eigen::Dynamic>
        mPrefix; 
    Eigen::Vector<IndexType, Eigen::Dynamic>
        mPrimitiveIds; 
};
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
inline HashGrid<Dims, TScalar, TIndex>::HashGrid(ScalarType cellSize, IndexType nBuckets)
    : HashGrid<Dims, TScalar, TIndex>()
{
    Configure(cellSize, nBuckets);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
inline void HashGrid<Dims, TScalar, TIndex>::Configure(ScalarType cellSize, IndexType nBuckets)
{
    mCellSize = cellSize;
    mPrefix.resize(nBuckets + 1);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
inline void HashGrid<Dims, TScalar, TIndex>::Reserve(IndexType nPrimitives)
{
    mBucketIds.resize(nPrimitives);
    mPrimitiveIds.resize(nPrimitives);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FHash, class TDerivedL, class TDerivedU>
void HashGrid<Dims, TScalar, TIndex>::Construct(
    Eigen::DenseBase<TDerivedL> const& L,
    Eigen::DenseBase<TDerivedU> const& U,
    FHash fHash)
{
    assert(L.rows() == U.rows());
    assert(L.cols() == U.cols());
    Construct(ScalarType(0.5) * (L.derived() + U.derived()), fHash);
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FHash, class TDerivedX>
inline void
HashGrid<Dims, TScalar, TIndex>::Construct(Eigen::DenseBase<TDerivedX> const& X, FHash fHash)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.HashGrid.Construct");
    auto const nPrimitives = static_cast<IndexType>(X.cols());
    auto const nBuckets    = NumberOfBuckets();
    assert(nBuckets >= nPrimitives);
    assert(nPrimitives == static_cast<IndexType>(X.cols()));
    assert(X.rows() == kDims);
    Reserve(nPrimitives);
    // Map primitives to hash IDs
    for (IndexType i = 0; i < nPrimitives; ++i)
    {
        auto ixyz     = ToIntegerCoordinates(X.col(i));
        mBucketIds[i] = common::Modulo(fHash(ixyz), nBuckets);
    }
    // Compute id counts in the prefix sum's memory
    mPrefix.setZero();
    mPrefix(mBucketIds).array() += 1;
    // Compute the shifted prefix sum
    std::inclusive_scan(mPrefix.begin(), mPrefix.end(), mPrefix.begin());
    // Construct primitive IDs while unshifting prefix sum
    for (IndexType i = 0; i < nPrimitives; ++i)
    {
        auto bucketId                      = mBucketIds[i];
        mPrimitiveIds[--mPrefix[bucketId]] = i;
    }
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class FOnPair, class FHash, class TDerivedX>
inline void HashGrid<Dims, TScalar, TIndex>::BroadPhase(
    Eigen::DenseBase<TDerivedX> const& X,
    FOnPair fOnPair,
    FHash fHash) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.geometry.HashGrid.BroadPhase");
    assert(X.rows() == kDims);
    auto const nQueries   = static_cast<IndexType>(X.cols());
    auto const nBuckets   = NumberOfBuckets();
    auto const fVisitCell = [&](IndexType q, Eigen::Vector<IndexType, kDims> const& ixyz) {
        auto bucketId = common::Modulo(fHash(ixyz), nBuckets);
        auto begin    = mPrefix[bucketId];
        auto end      = mPrefix[bucketId + 1];
        for (auto j = begin; j < end; ++j)
            fOnPair(q, mPrimitiveIds[j]);
    };
    for (IndexType q = 0; q < nQueries; ++q)
    {
        Eigen::Vector<IndexType, kDims> ixyz = ToIntegerCoordinates(X.col(q));
        if constexpr (kDims == 2)
        {
            for (auto cx = -1; cx <= 1; ++cx)
            {
                ixyz.x() += cx;
                for (auto cy = -1; cy <= 1; ++cy)
                {
                    ixyz.y() += cy;
                    fVisitCell(q, ixyz);
                    ixyz.y() -= cy;
                }
                ixyz.x() -= cx;
            }
        }
        if constexpr (kDims == 3)
        {
            for (auto cx = -1; cx <= 1; ++cx)
            {
                ixyz.x() += cx;
                for (auto cy = -1; cy <= 1; ++cy)
                {
                    ixyz.y() += cy;
                    for (auto cz = -1; cz <= 1; ++cz)
                    {
                        ixyz.z() += cz;
                        fVisitCell(q, ixyz);
                        ixyz.z() -= cz;
                    }
                    ixyz.y() -= cy;
                }
                ixyz.x() -= cx;
            }
        }
    }
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class TDerivedX>
inline auto HashGrid<Dims, TScalar, TIndex>::Cell(Eigen::DenseBase<TDerivedX> const& X) const
    -> Eigen::Matrix<ScalarType, kDims, 2>
{
    Eigen::Matrix<ScalarType, kDims, 2> cell{};
    auto const ixyz = ToIntegerCoordinates(X);
    cell.col(0)     = ixyz.template cast<ScalarType>() * mCellSize;
    cell.col(1)     = cell.col(0).array() + static_cast<ScalarType>(mCellSize);
    return cell;
}
 
template <int Dims, common::CFloatingPoint TScalar, common::CIndex TIndex>
template <class TDerivedX>
inline auto
HashGrid<Dims, TScalar, TIndex>::ToIntegerCoordinates(Eigen::DenseBase<TDerivedX> const& X) const
    -> Eigen::Vector<IndexType, kDims>
{
    return Eigen::Vector<ScalarType, kDims>(
               X.derived().array() / static_cast<ScalarType>(mCellSize))
        .array()
        .floor()
        .template cast<IndexType>();
}
 
} // namespace pbat::geometry
 
#endif // PBAT_GEOMETRY_HASHGRID_H