kdtree2.h

Go to the documentation of this file.

#ifndef included_AMP_kdtree2
#define included_AMP_kdtree2
 
#include <array>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <tuple>
#include <vector>
 
 
namespace AMP {
 
 
template<uint8_t NDIM, class TYPE>
class kdtree2
{
public: // Convience typedef
    using Point = std::array<double, NDIM>;
 
public:
    kdtree2( const std::vector<Point> &x, const std::vector<TYPE> &data );
 
    std::array<double, 2 * NDIM> box() const;
 
    inline size_t size() const { return d_N; }
 
    inline bool empty() const { return d_N == 0; }
 
    std::vector<Point> getPoints() const;
 
    std::vector<std::pair<Point, TYPE>> getPointsAndData() const;
 
    void add( const Point &x, const TYPE &data );
 
    std::tuple<Point, TYPE> findNearest( const Point &x ) const;
 
    std::vector<std::tuple<Point, TYPE>> findNearest( const Point &x, int N ) const;
 
    std::vector<std::tuple<Point, TYPE>> findNearest( const Point &x, double dist ) const;
 
    std::vector<std::tuple<Point, TYPE>>
    findNearestRay( const Point &x, const Point &dir, double dist ) const;
 
 
public: // Copy/assignment operators
    kdtree2()                  = default;
    ~kdtree2()                 = default;
    kdtree2( const kdtree2 & ) = delete;
    kdtree2( kdtree2 && )      = default;
    kdtree2 &operator=( const kdtree2 & ) = delete;
    kdtree2 &operator=( kdtree2 && ) = default;
 
 
private: // Internal data
    // Structure used to store point data in the lowest leaf
    struct data_struct {
        data_struct( const data_struct & ) = delete;
        data_struct &operator=( const data_struct & ) = delete;
        size_t N                                      = 0;
        Point *x                                      = nullptr;
        TYPE *data                                    = nullptr;
        data_struct( size_t N );
        ~data_struct();
        void add( const Point &x2, const TYPE &d2 );
    };
 
    // Internal data
    size_t d_N          = 0;
    uint8_t d_split_dim = 0;
    double d_split      = 0;
    Point d_lb          = { 0 };
    Point d_ub          = { 0 };
    std::unique_ptr<kdtree2> d_left;
    std::unique_ptr<kdtree2> d_right;
    std::unique_ptr<data_struct> d_data;
 
 
private: // Internal functions
    kdtree2( size_t N, Point *x, TYPE *data );
    void initialize( size_t N, Point *x, TYPE *data );
    void splitData( size_t N, Point *x, TYPE *data );
    bool intersect( const Point &x, double dist2 ) const;
    void getPoints( std::vector<Point> &x ) const;
    void getPoints( std::vector<std::pair<Point, TYPE>> &x ) const;
    void
    findNearest( const Point &x, size_t N, std::tuple<Point, TYPE> *nearest, double *dist ) const;
    void findNearest( const Point &x,
                      double dist2,
                      std::vector<std::tuple<Point, TYPE>> &nearest ) const;
    void findNearestRay( const Point &x,
                         const Point &dir,
                         const Point &n_inv,
                         double dist2,
                         std::vector<std::tuple<Point, TYPE>> &nearest ) const;
    void
    checkNearest( const Point &x, size_t N, std::tuple<Point, TYPE> *nearest, double *dist ) const;
};
 
 
} // namespace AMP
 
 
#endif