EBGeometry_BVH.hpp

Go to the documentation of this file.

/* EBGeometry
 * Copyright © 2022 Robert Marskar
 * Please refer to Copyright.txt and LICENSE in the EBGeometry root directory.
 */
 
#ifndef EBGeometry_BVH
#define EBGeometry_BVH
 
// Std includes
#include <memory>
#include <vector>
#include <array>
#include <functional>
#include <queue>
 
// Our includes
#include "EBGeometry_Vec.hpp"
#include "EBGeometry_SFC.hpp"
#include "EBGeometry_NamespaceHeader.hpp"
 
namespace BVH {
 
  enum class Build
  {
    TopDown,
    Morton,
    Nested
  };
 
  template <class T, class P, class BV, size_t K>
  class NodeT;
 
  template <class T, class P, class BV, size_t K>
  class LinearNodeT;
 
  template <class T, class P, class BV, size_t K>
  class LinearBVH;
 
  template <class P>
  using PrimitiveListT = std::vector<std::shared_ptr<const P>>;
 
  template <class P, class BV>
  using PrimAndBV = std::pair<std::shared_ptr<const P>, BV>;
 
  template <class P, class BV>
  using PrimAndBVListT = std::vector<PrimAndBV<P, BV>>;
 
  template <class P, class BV, size_t K>
  using PartitionerT = std::function<std::array<PrimAndBVListT<P, BV>, K>(const PrimAndBVListT<P, BV>& a_primsAndBVs)>;
 
  template <class T, class P, class BV, size_t K>
  using StopFunctionT = std::function<bool(const NodeT<T, P, BV, K>& a_node)>;
 
  template <class P>
  using Updater = std::function<void(const PrimitiveListT<P>& a_primitives)>;
 
  template <class NodeType, class Meta>
  using Visiter = std::function<bool(const NodeType& a_node, const Meta& a_meta)>;
 
  template <class NodeType, class Meta, size_t K>
  using Sorter = std::function<void(std::array<std::pair<std::shared_ptr<const NodeType>, Meta>, K>& a_children)>;
 
  template <class NodeType, class Meta>
  using MetaUpdater = std::function<Meta(const NodeType& a_node)>;
 
  template <class X, size_t K>
  auto equalCounts = [](const std::vector<X>& a_primitives) noexcept -> std::array<std::vector<X>, K> {
    int length = a_primitives.size() / K;
    int remain = a_primitives.size() % K;
 
    int begin = 0;
    int end   = 0;
 
    std::array<std::vector<X>, K> chunks;
 
    for (size_t k = 0; k < K; k++) {
      end += (remain > 0) ? length + 1 : length;
      remain--;
 
      chunks[k] = std::vector<X>(a_primitives.begin() + begin, a_primitives.begin() + end);
 
      begin = end;
    }
 
    return chunks;
  };
 
  template <class T, class P, class BV, size_t K>
  auto PrimitiveCentroidPartitioner =
    [](const PrimAndBVListT<P, BV>& a_primsAndBVs) noexcept -> std::array<PrimAndBVListT<P, BV>, K> {
    Vec3T<T> lo = Vec3T<T>::max();
    Vec3T<T> hi = -Vec3T<T>::max();
 
    for (const auto& pbv : a_primsAndBVs) {
      lo = min(lo, pbv.first->getCentroid());
      hi = max(hi, pbv.first->getCentroid());
    }
 
    const size_t splitDir = (hi - lo).maxDir(true);
 
    // Sort the primitives along the above coordinate direction.
    PrimAndBVListT<P, BV> sortedPrimsAndBVs(a_primsAndBVs);
 
    std::sort(sortedPrimsAndBVs.begin(),
              sortedPrimsAndBVs.end(),
              [splitDir](const PrimAndBV<P, BV>& pbv1, const PrimAndBV<P, BV>& pbv2) -> bool {
                return pbv1.first->getCentroid(splitDir) < pbv2.first->getCentroid(splitDir);
              });
 
    return BVH::equalCounts<PrimAndBV<P, BV>, K>(sortedPrimsAndBVs);
  };
 
  template <class T, class P, class BV, size_t K>
  auto BVCentroidPartitioner = [](const PrimAndBVListT<P, BV>& a_primsAndBVs) -> std::array<PrimAndBVListT<P, BV>, K> {
    Vec3T<T> lo = Vec3T<T>::max();
    Vec3T<T> hi = -Vec3T<T>::max();
 
    for (const auto& pbv : a_primsAndBVs) {
      lo = min(lo, pbv.second.getCentroid());
      hi = max(hi, pbv.second.getCentroid());
    }
 
    const size_t splitDir = (hi - lo).maxDir(true);
 
    // Sort the primitives along the above coordinate direction.
    PrimAndBVListT<P, BV> sortedPrimsAndBVs(a_primsAndBVs);
 
    std::sort(sortedPrimsAndBVs.begin(),
              sortedPrimsAndBVs.end(),
              [splitDir](const PrimAndBV<P, BV>& pbv1, const PrimAndBV<P, BV>& pbv2) -> bool {
                return pbv1.second.getCentroid()[splitDir] < pbv2.second.getCentroid()[splitDir];
              });
 
    return BVH::equalCounts<PrimAndBV<P, BV>, K>(sortedPrimsAndBVs);
  };
 
  template <class T, class P, class BV, size_t K>
  auto DefaultStopFunction =
    [](const BVH::NodeT<T, P, BV, K>& a_node) noexcept -> bool { return (a_node.getPrimitives()).size() < K; };
 
  template <class T, class P, class BV, size_t K>
  class NodeT : public std::enable_shared_from_this<NodeT<T, P, BV, K>>
  {
  public:
    using PrimitiveList = PrimitiveListT<P>;
 
    using Vec3 = Vec3T<T>;
 
    using Node = NodeT<T, P, BV, K>;
 
    using NodePtr = std::shared_ptr<Node>;
 
    using Partitioner = PartitionerT<P, BV, K>;
 
    using StopFunction = StopFunctionT<T, P, BV, K>;
 
    NodeT() noexcept;
 
    NodeT(const std::vector<PrimAndBV<P, BV>>& a_primsAndBVs) noexcept;
 
    virtual ~NodeT() noexcept;
 
    inline void
    topDownSortAndPartition(const Partitioner&  a_partitioner = BVCentroidPartitioner<T, P, BV, K>,
                            const StopFunction& a_stopCrit    = DefaultStopFunction<T, P, BV, K>) noexcept;
 
    template <typename S>
    inline void
    bottomUpSortAndPartition() noexcept;
 
    inline bool
    isLeaf() const noexcept;
 
    inline bool
    isPartitioned() const noexcept;
 
    inline const BV&
    getBoundingVolume() const noexcept;
 
    inline const PrimitiveList&
    getPrimitives() const noexcept;
 
    inline const std::vector<BV>&
    getBoundingVolumes() const noexcept;
 
    inline T
    getDistanceToBoundingVolume(const Vec3& a_point) const noexcept;
 
    inline const std::array<std::shared_ptr<NodeT<T, P, BV, K>>, K>&
    getChildren() const noexcept;
 
    template <class Meta>
    inline void
    traverse(const BVH::Updater<P>&              a_updater,
             const BVH::Visiter<Node, Meta>&     a_visiter,
             const BVH::Sorter<Node, Meta, K>&   a_sorter,
             const BVH::MetaUpdater<Node, Meta>& a_metaUpdater) const noexcept;
 
    inline std::shared_ptr<LinearBVH<T, P, BV, K>>
    flattenTree() const noexcept;
 
  protected:
    BV m_boundingVolume;
 
    bool m_partitioned;
 
    std::vector<std::shared_ptr<const P>> m_primitives;
 
    std::vector<BV> m_boundingVolumes;
 
    std::array<std::shared_ptr<NodeT<T, P, BV, K>>, K> m_children;
 
    inline PrimitiveList&
    getPrimitives() noexcept;
 
    inline std::vector<BV>&
    getBoundingVolumes() noexcept;
 
    inline void
    setChildren(const std::array<std::shared_ptr<NodeT<T, P, BV, K>>, K>& a_children) noexcept;
 
    inline size_t
    flattenTree(std::vector<std::shared_ptr<LinearNodeT<T, P, BV, K>>>& a_linearNodes,
                std::vector<std::shared_ptr<const P>>&                  a_sortedPrimitives,
                size_t&                                                 a_offset) const noexcept;
  };
 
  template <class T, class P, class BV, size_t K>
  class LinearNodeT
  {
  public:
    using Vec3 = Vec3T<T>;
 
    inline LinearNodeT() noexcept;
 
    inline virtual ~LinearNodeT();
 
    inline void
    setBoundingVolume(const BV& a_boundingVolume) noexcept;
 
    inline void
    setPrimitivesOffset(const size_t a_primitivesOffset) noexcept;
 
    inline void
    setNumPrimitives(const size_t a_numPrimitives) noexcept;
 
    inline void
    setChildOffset(const size_t a_childOffset, const size_t a_whichChild) noexcept;
 
    inline const BV&
    getBoundingVolume() const noexcept;
 
    inline const size_t&
    getPrimitivesOffset() const noexcept;
 
    inline const size_t&
    getNumPrimitives() const noexcept;
 
    inline const std::array<size_t, K>&
    getChildOffsets() const noexcept;
 
    inline bool
    isLeaf() const noexcept;
 
    inline bool
    isPartitioned() const noexcept;
 
    inline T
    getDistanceToBoundingVolume(const Vec3& a_point) const noexcept;
 
    inline std::vector<T>
    getDistances(const Vec3& a_point, const std::vector<std::shared_ptr<const P>>& a_primitives) const noexcept;
 
  protected:
    BV m_boundingVolume;
 
    bool m_partitioned;
 
    size_t m_primitivesOffset;
 
    size_t m_numPrimitives;
 
    std::array<size_t, K> m_childOffsets;
  };
 
  template <class T, class P, class BV, size_t K>
  class LinearBVH
  {
  public:
    using Vec3 = Vec3T<T>;
 
    using LinearNode = LinearNodeT<T, P, BV, K>;
 
    using PrimitiveList = std::vector<std::shared_ptr<const P>>;
 
    inline LinearBVH() = default;
 
    inline LinearBVH(const LinearBVH&) = default;
 
    inline LinearBVH(const std::vector<std::shared_ptr<const LinearNodeT<T, P, BV, K>>>& a_linearNodes,
                     const std::vector<std::shared_ptr<const P>>&                        a_primitives);
 
    inline LinearBVH(const std::vector<std::shared_ptr<LinearNodeT<T, P, BV, K>>>& a_linearNodes,
                     const std::vector<std::shared_ptr<const P>>&                  a_primitives);
 
    inline virtual ~LinearBVH();
 
    inline const BV&
    getBoundingVolume() const noexcept;
 
    template <class Meta>
    inline void
    traverse(const BVH::Updater<P>&                    a_updater,
             const BVH::Visiter<LinearNode, Meta>&     a_visiter,
             const BVH::Sorter<LinearNode, Meta, K>&   a_sorter,
             const BVH::MetaUpdater<LinearNode, Meta>& a_metaUpdater) const noexcept;
 
  protected:
    std::vector<std::shared_ptr<const LinearNodeT<T, P, BV, K>>> m_linearNodes;
 
    std::vector<std::shared_ptr<const P>> m_primitives;
  };
} // namespace BVH
 
#include "EBGeometry_NamespaceFooter.hpp"
 
#include "EBGeometry_BVHImplem.hpp"
 
#endif