KdTreePoint.h

#pragma once
// ===========================================================================
// Copyright (c) 2018 Toolchefs Ltd. All rights reserved.
//
// Use of this software is subject to the terms of the Toolchefs license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy form.
// ===========================================================================
 
#include <AtomsUtils/Globals.h>
#include <AtomsUtils/AtomsMath.h>
#include <vector>
 
namespace AtomsUtils
{
    class KdTreePoint;
 
    class ATOMSUTILS_EXPORT KdTreePoint
    {
    public:
 
        typedef int DataType;
 
        struct Sample
        {
            DataType data;
            AtomsMath::Vector3 pos;
        };
 
        class KdTreePointNode
        {
        public:
 
            inline unsigned int numSamples() const { return high - low + 1; }
 
            inline bool isLeaf() const { return left == -1 && right == -1; }
 
        public:
 
            AtomsMath::Box3 bbox;
            unsigned int low, high;
            int left, right;
        };
 
        KdTreePoint(unsigned int maxSamplesPerNode = 10, unsigned int maxDepth = 10);
 
        ~KdTreePoint();
 
        KdTreePoint(const KdTreePoint& rhs);
 
        void addSample(const AtomsMath::Vector3 &sample, const int value);
 
        void buildTree();
 
        void clear();
 
        void reserve(size_t maxNumSamples);
 
        void findInRange(const AtomsMath::Vector3& position, const double radius, std::vector<DataType>& neighbours) const;
 
        void findSampleInRange(const AtomsMath::Vector3& position, const double radius, std::vector<int>& neighbours) const;
 
        bool findClosestSample(const AtomsMath::Vector3 &searchPoint, int &closestId, double* maxdist = nullptr) const;
 
        void findClosestSamples(const AtomsMath::Vector3 &searchPoint, std::vector<std::pair<int, double>> &closestIds, unsigned int numSamples) const;
 
        void findClosestSamples(const AtomsMath::Vector3 &searchPoint, std::vector<std::pair<int, double>> &closestIds, unsigned int numSamples, double maxDistance) const;
 
        bool intersect(const AtomsMath::Vector3& start, const AtomsMath::Vector3& end, const bool disableY) const;
 
        void findInRayRange(const AtomsMath::Line3& line, double radius, std::vector<DataType>& neighbours) const;
 
        void findSampleInRayRange(const AtomsMath::Line3& line, double radius, std::vector<int>& neighbours) const;
 
        inline std::vector<Sample>& samples() { return m_samples; }
 
        inline const std::vector<Sample>& samples() const { return m_samples; }
 
        inline std::vector<int>& comb() { return m_comb; }
 
        inline std::vector<KdTreePointNode>& nodes() { return m_nodes; }
 
    private:
 
        void buildNode(int nodeIndex, const AtomsMath::Box3& box, unsigned int depth, int low, int high);
 
        unsigned int allocateNode();
 
        void findInRangeRecurse(unsigned int index, const AtomsMath::Vector3& position, const double radius, std::vector<DataType>& neighbours) const;
 
        bool intersectRecurse(unsigned int index, const AtomsMath::Line3d& line, const bool disableY) const;
 
        void findSampleInRangeRecurse(unsigned int index, const AtomsMath::Vector3& position, const double radius, std::vector<int>& neighbours) const;
 
        void findInRayRangeRecurse(unsigned int index, const AtomsMath::Line3& line, const double radius, std::vector<DataType>& neighbours) const;
 
        void findSampleInRayRangeRecurse(unsigned int index, const AtomsMath::Line3& line, const double radius, std::vector<int>& neighbours) const;
 
        void walkCells(unsigned int cellIndex, const AtomsMath::Vector3 &searchPoint, int &closestId, double& maxDist) const;
 
        void walkCells(unsigned int cellIndex, const AtomsMath::Vector3 &searchPoint, std::vector<std::pair<int, double>> & neighbours, unsigned int numSamples, double &maxDist) const;
 
    private:
 
        std::vector<KdTreePointNode> m_nodes;
 
        std::vector<Sample> m_samples;
 
        std::vector<int> m_comb;
 
        unsigned int m_maxSamplesPerNode;
 
        unsigned int m_maxDepth;
    };
}