KdTree.h

Go to the documentation of this file.

#pragma once
 
 
#include "io/Logger.h"
#include "objects/finders/NeighbourFinder.h"
#include "objects/geometry/Box.h"
#include "objects/utility/IteratorAdapters.h"
#include "objects/wrappers/Finally.h"
#include "objects/wrappers/Function.h"
#include "objects/wrappers/Outcome.h"
#include "thread/ThreadLocal.h"
#include <set>
#include <shared_mutex>
 
NAMESPACE_SPH_BEGIN
 
struct KdNode : public Noncopyable {
    enum class Type : Size { X, Y, Z, LEAF };
    Type type;
 
    Box box;
 
    KdNode(const Type& type)
        : type(type) {}
 
    INLINE bool isLeaf() const {
        return type == Type::LEAF;
    }
};
 
template <typename TBase>
struct InnerNode : public TBase {
    float splitPosition;
 
    Size left;
 
    Size right;
 
    InnerNode()
        : TBase(KdNode::Type(-1)) {}
 
    InnerNode(const KdNode::Type& type)
        : TBase(type) {}
};
 
template <typename TBase>
struct LeafNode : public TBase {
    Size from;
 
    Size to;
 
    Size padding;
 
    LeafNode(const KdNode::Type& type)
        : TBase(type) {}
 
    INLINE Size size() const {
        return to - from;
    }
};
 
static_assert(sizeof(Size) == sizeof(float), "Sizes must match to keep this layout");
 
class LeafIndexIterator : public IndexIterator {
private:
    ArrayView<const Size> mapping;
 
public:
    INLINE LeafIndexIterator(const Size idx, ArrayView<const Size> mapping)
        : IndexIterator(idx)
        , mapping(mapping) {}
 
    INLINE Size operator*() const {
        return mapping[idx];
    }
};
 
class LeafIndexSequence : public IndexSequence {
private:
    ArrayView<const Size> mapping;
 
public:
    INLINE LeafIndexSequence(const Size from, const Size to, ArrayView<const Size> mapping)
        : IndexSequence(from, to)
        , mapping(mapping) {
        SPH_ASSERT(to <= mapping.size());
    }
 
    INLINE LeafIndexIterator begin() const {
        return LeafIndexIterator(from, mapping);
    }
 
    INLINE LeafIndexIterator end() const {
        return LeafIndexIterator(to, mapping);
    }
};
 
struct EuclideanMetric {
    INLINE Float operator()(const Vector& v) const {
        return getSqrLength(v);
    }
};
 
enum class KdChild;
 
template <typename TNode, typename TMetric = EuclideanMetric>
class KdTree : public FinderTemplate<KdTree<TNode, TMetric>> {
private:
    struct {
        Size leafSize;
 
        Size maxParallelDepth;
    } config;
 
    Box entireBox;
 
    Array<Size> idxs;
 
 
    Array<InnerNode<TNode>> nodes;
    std::atomic_int nodeCounter;
    std::shared_timed_mutex nodesMutex;
 
    static constexpr Size ROOT_PARENT_NODE = -1;
 
public:
    explicit KdTree(const Size leafSize = 25, const Size maxParallelDepth = 50) {
        SPH_ASSERT(leafSize >= 1);
        config.leafSize = leafSize;
        config.maxParallelDepth = maxParallelDepth;
    }
 
    template <bool FindAll>
    Size find(const Vector& pos, const Size index, const Float radius, Array<NeighbourRecord>& neighs) const;
 
    INLINE TNode& getNode(const Size nodeIdx) {
        return nodes[nodeIdx];
    }
 
    INLINE const TNode& getNode(const Size nodeIdx) const {
        return nodes[nodeIdx];
    }
 
    INLINE Size getNodeCnt() const {
        return nodes.size();
    }
 
    INLINE LeafIndexSequence getLeafIndices(const LeafNode<TNode>& leaf) const {
        return LeafIndexSequence(leaf.from, leaf.to, idxs);
    }
 
    Outcome sanityCheck() const;
 
protected:
    virtual void buildImpl(IScheduler& scheduler, ArrayView<const Vector> points) override;
 
private:
    void init();
 
    void buildTree(IScheduler& scheduler,
        const Size parent,
        const KdChild child,
        const Size from,
        const Size to,
        const Box& box,
        const Size slidingCnt,
        const Size depth);
 
    void addLeaf(const Size parent, const KdChild child, const Size from, const Size to);
 
    Size addInner(const Size parent, const KdChild child, const Float splitPosition, const Size splitIdx);
 
    bool isSingular(const Size from, const Size to, const Size splitIdx) const;
 
    bool checkBoxes(const Size from, const Size to, const Size mid, const Box& box1, const Box& box2) const;
};
 
 
enum class IterateDirection {
    TOP_DOWN,  
    BOTTOM_UP, 
};
 
template <IterateDirection Dir, typename TNode, typename TMetric, typename TFunctor>
void iterateTree(KdTree<TNode, TMetric>& tree,
    IScheduler& scheduler,
    const TFunctor& functor,
    const Size nodeIdx = 0,
    const Size depthLimit = Size(-1));
 
template <IterateDirection Dir, typename TNode, typename TMetric, typename TFunctor>
void iterateTree(const KdTree<TNode, TMetric>& tree,
    IScheduler& scheduler,
    const TFunctor& functor,
    const Size nodeIdx = 0,
    const Size depthLimit = Size(-1));
 
NAMESPACE_SPH_END
 
#include "objects/finders/KdTree.inl.h"