kdTree.h

/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/
 
#pragma once
 
#include "./kdTreeTraits.h"
 
#include <algorithm>
#include <iostream>
#include <memory>
#include <numeric>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <Eigen/Geometry> // aabb
 
#include "../../Common/Assert.h"
 
#include "Query/kdTreeNearestQueries.h"
#include "Query/kdTreeKNearestQueries.h"
#include "Query/kdTreeRangeQueries.h"
 
namespace Ponca
{
    template <typename Traits>
    class KdTreeBase;
    template <typename Traits>
    class KdTreeDenseBase;
    template <typename Traits>
    class KdTreeSparseBase;
 
#ifdef PARSED_WITH_DOXYGEN
    // [KdTree type definition]
    template <typename DataPoint>
    struct KdTree : public Ponca::KdTreeBase<KdTreeDefaultTraits<DataPoint>>
    {
    };
// [KdTree type definition]
#else
    template <typename DataPoint>
    using KdTree = KdTreeBase<KdTreeDefaultTraits<DataPoint>>; // prefer alias to avoid redefining methods
#endif
 
#ifdef PARSED_WITH_DOXYGEN
    // [StaticKdTree type definition]
    template <typename DataPoint>
    struct StaticKdTree : public Ponca::StaticKdTreeBase<KdTreeDefaultTraits<DataPoint>>
    {
    };
// [StaticKdTree type definition]
#else
    template <typename DataPoint>
    using StaticKdTree = StaticKdTreeBase<KdTreeDefaultTraits<DataPoint>>; // prefer alias to avoid redefining methods
#endif
 
#ifdef PARSED_WITH_DOXYGEN
    // [KdTreeDense type definition]
    template <typename DataPoint>
    struct KdTreeDense : public Ponca::KdTreeDenseBase<KdTreeDefaultTraits<DataPoint>>
    {
    };
// [KdTreeDense type definition]
#else
    template <typename DataPoint>
    using KdTreeDense = KdTreeDenseBase<KdTreeDefaultTraits<DataPoint>>; // prefer alias to avoid redefining methods
#endif
 
#ifdef PARSED_WITH_DOXYGEN
    // [KdTreeSparse type definition]
    template <typename DataPoint>
    struct KdTreeSparse : Ponca::KdTreeSparseBase<KdTreeDefaultTraits<DataPoint>>
    {
    };
// [KdTreeSparse type definition]
#else
    template <typename DataPoint>
    using KdTreeSparse = KdTreeSparseBase<KdTreeDefaultTraits<DataPoint>>;
#endif
 
    template <typename _Traits>
    class StaticKdTreeBase
    {
    public:
#define WRITE_TRAITS                                                                                                   \
    using Traits         = _Traits;                         \
    using DataPoint      = typename Traits::DataPoint;      \
    using IndexType      = typename Traits::IndexType;      \
    using LeafSizeType   = typename Traits::LeafSizeType;   \
    using PointContainer = typename Traits::PointContainer;      \
    using IndexContainer = typename Traits::IndexContainer;        \
    using NodeIndexType  = typename Traits::NodeIndexType;  \
    using NodeType       = typename Traits::NodeType;       \
    using NodeContainer  = typename Traits::NodeContainer;  \
    using Scalar         = typename DataPoint::Scalar;      \
    using VectorType     = typename DataPoint::VectorType;  \
    using AabbType       = typename NodeType::AabbType;    
        WRITE_TRAITS
 
        struct Buffers
        {
            PointContainer points;  
            NodeContainer nodes;    
            IndexContainer indices; 
 
            size_t points_size{0};
            size_t nodes_size{0};
            size_t indices_size{0};
 
            PONCA_MULTIARCH inline Buffers() = default;
 
            PONCA_MULTIARCH inline Buffers(PointContainer _points, NodeContainer _nodes, IndexContainer _indices,
                                                 const size_t _points_size, const size_t _nodes_size,
                                                 const size_t _indices_size)
                : points(_points), nodes(_nodes), indices(_indices), points_size(_points_size), nodes_size(_nodes_size),
                  indices_size(_indices_size)
            {
            }
        };
 
        static constexpr std::size_t MAX_NODE_COUNT = NodeType::MAX_COUNT;
        static constexpr std::size_t MAX_POINT_COUNT = std::size_t(2) << sizeof(IndexType) * 8;
 
        static constexpr int MAX_DEPTH = Traits::MAX_DEPTH;
 
        static constexpr bool SUPPORTS_SUBSAMPLING = false;
 
        // Queries use a value of -1 for invalid indices
        static_assert(std::is_signed_v<IndexType>, "Index type must be signed");
        static_assert(MAX_DEPTH > 0, "Max depth must be strictly positive");
 
        // Construction ------------------------------------------------------------
    protected:
        PONCA_MULTIARCH inline StaticKdTreeBase() = default;
 
    public:
        PONCA_MULTIARCH inline StaticKdTreeBase(Buffers& buf) : m_bufs(buf) {}
 
        // Accessors ---------------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] inline NodeIndexType nodeCount() const
        {
            return (NodeIndexType)m_bufs.nodes_size;
        }
 
        PONCA_MULTIARCH [[nodiscard]] inline IndexType sampleCount() const { return (IndexType)m_bufs.indices_size; }
 
        PONCA_MULTIARCH [[nodiscard]] inline IndexType pointCount() const { return (IndexType)m_bufs.points_size; }
 
        PONCA_MULTIARCH [[nodiscard]] inline NodeIndexType leafCount() const { return m_leaf_count; }
 
        PONCA_MULTIARCH [[nodiscard]] inline PointContainer& points() { return m_bufs.points; };
 
        PONCA_MULTIARCH [[nodiscard]] inline const PointContainer& points() const { return m_bufs.points; };
 
        PONCA_MULTIARCH [[nodiscard]] inline const NodeContainer& nodes() const { return m_bufs.nodes; }
 
        PONCA_MULTIARCH [[nodiscard]] inline const IndexContainer& samples() const { return m_bufs.indices; }
 
        PONCA_MULTIARCH [[nodiscard]] inline const Buffers& buffers() const { return m_bufs; }
 
        // Parameters --------------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] inline LeafSizeType minCellSize() const { return m_min_cell_size; }
 
        PONCA_MULTIARCH inline void setMinCellSize(LeafSizeType min_cell_size)
        {
            PONCA_DEBUG_ASSERT(min_cell_size > 0);
            m_min_cell_size = min_cell_size;
        }
 
        // Index mapping -----------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] inline IndexType pointFromSample(IndexType sample_index) const
        {
            return m_bufs.indices[sample_index];
        }
 
        PONCA_MULTIARCH [[nodiscard]] inline DataPoint& pointDataFromSample(IndexType sample_index)
        {
            return m_bufs.points[pointFromSample(sample_index)];
        }
 
        PONCA_MULTIARCH [[nodiscard]] inline const DataPoint& pointDataFromSample(IndexType sample_index) const
        {
            return m_bufs.points[pointFromSample(sample_index)];
        }
 
        // Query -------------------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] KdTreeKNearestPointQuery<Traits> kNearestNeighbors(const VectorType& point,
                                                                                               IndexType k) const
        {
            return KdTreeKNearestPointQuery<Traits>(this, k, point);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeKNearestIndexQuery<Traits> kNearestNeighbors(IndexType index,
                                                                                               IndexType k) const
        {
            return KdTreeKNearestIndexQuery<Traits>(this, k, index);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeKNearestPointQuery<Traits> kNearestNeighborsQuery() const
        {
            return KdTreeKNearestPointQuery<Traits>(this, 0, VectorType::Zero());
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeKNearestIndexQuery<Traits> kNearestNeighborsIndexQuery() const
        {
            return KdTreeKNearestIndexQuery<Traits>(this, 0, 0);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeNearestPointQuery<Traits> nearestNeighbor(const VectorType& point) const
        {
            return KdTreeNearestPointQuery<Traits>(this, point);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeNearestIndexQuery<Traits> nearestNeighbor(IndexType index) const
        {
            return KdTreeNearestIndexQuery<Traits>(this, index);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeNearestIndexQuery<Traits> nearestNeighborQuery() const
        {
            return KdTreeNearestIndexQuery<Traits>(this, VectorType::Zero());
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeNearestIndexQuery<Traits> nearestNeighborIndexQuery() const
        {
            return KdTreeNearestIndexQuery<Traits>(this, 0);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeRangePointQuery<Traits> rangeNeighbors(const VectorType& point,
                                                                                         Scalar r) const
        {
            return KdTreeRangePointQuery<Traits>(this, r, point);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeRangeIndexQuery<Traits> rangeNeighbors(IndexType index, Scalar r) const
        {
            return KdTreeRangeIndexQuery<Traits>(this, r, index);
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeRangePointQuery<Traits> rangeNeighborsQuery() const
        {
            return KdTreeRangePointQuery<Traits>(this, 0, VectorType::Zero());
        }
 
        PONCA_MULTIARCH [[nodiscard]] KdTreeRangeIndexQuery<Traits> rangeNeighborsIndexQuery() const
        {
            return KdTreeRangeIndexQuery<Traits>(this, 0, 0);
        }
 
        // Utilities ---------------------------------------------------------------
    public:
        PONCA_MULTIARCH_HOST [[nodiscard]] inline bool valid() const;
        PONCA_MULTIARCH_HOST inline void print(std::ostream& os, bool verbose = false) const;
 
        // Data --------------------------------------------------------------------
    protected:
        Buffers m_bufs;                   
        LeafSizeType m_min_cell_size{64}; 
        NodeIndexType m_leaf_count{0};    
    };
 
    template <typename _Traits>
    class KdTreeBase : public StaticKdTreeBase<_Traits>
    {
    public:
        using Base = StaticKdTreeBase<_Traits>;
        WRITE_TRAITS
        template <typename PointUserContainer, typename PointConverter>
        PONCA_MULTIARCH_HOST inline void build(PointUserContainer&& points, PointConverter c);
 
        struct DefaultConverter
        {
            template <typename Input>
            PONCA_MULTIARCH_HOST inline void operator()(Input&& i, PointContainer& o)
            {
                using InputContainer = std::remove_reference_t<Input>;
                if constexpr (std::is_same_v<InputContainer, PointContainer> && std::is_copy_assignable_v<DataPoint>)
                    o = std::forward<Input>(i); // Either move or copy
                else
                    std::transform(
                        i.cbegin(), i.cend(), std::back_inserter(o),
                        [](const typename InputContainer::value_type& p) -> DataPoint { return DataPoint(p); });
            }
        };
 
        template <typename PointUserContainer>
        PONCA_MULTIARCH_HOST inline void build(PointUserContainer&& points)
        {
            build(std::forward<PointUserContainer>(points), DefaultConverter());
        }
 
        // Internal ----------------------------------------------------------------
    protected:
        template <typename PointUserContainer, typename IndexUserContainer, typename PointConverter>
        PONCA_MULTIARCH_HOST inline void buildWithSampling(PointUserContainer&& points, IndexUserContainer&& sampling,
                                                           PointConverter c);
 
        template <typename PointUserContainer, typename IndexUserContainer>
        PONCA_MULTIARCH_HOST inline void buildWithSampling(PointUserContainer&& points, IndexUserContainer&& sampling)
        {
            buildWithSampling(std::forward<PointUserContainer>(points), std::forward<IndexUserContainer>(sampling),
                              DefaultConverter());
        }
 
    private:
        PONCA_MULTIARCH_HOST inline void buildRec(NodeIndexType node_id, IndexType start, IndexType end, int level);
        PONCA_MULTIARCH_HOST [[nodiscard]] inline IndexType partition(IndexType start, IndexType end, int dim,
                                                                      Scalar value);
    };
 
    template <typename Traits>
    class KdTreeDenseBase : public KdTreeBase<Traits>
    {
    private:
        using Base = KdTreeBase<Traits>;
 
    public:
        PONCA_MULTIARCH KdTreeDenseBase() = default;
 
        template <typename PointUserContainer>
        PONCA_MULTIARCH_HOST inline explicit KdTreeDenseBase(PointUserContainer&& points) : Base()
        {
            Base::build(std::forward<PointUserContainer>(points));
        }
    };
 
    template <typename Traits>
    class KdTreeSparseBase : public KdTreeBase<Traits>
    {
    private:
        using Base = KdTreeBase<Traits>;
 
    public:
        static constexpr bool SUPPORTS_SUBSAMPLING = true;
 
        PONCA_MULTIARCH KdTreeSparseBase() = default;
 
        template <typename PointUserContainer>
        PONCA_MULTIARCH_HOST inline explicit KdTreeSparseBase(PointUserContainer&& points) : Base()
        {
            this->build(std::forward<PointUserContainer>(points));
        }
 
        template <typename PointUserContainer, typename IndexUserContainer>
        PONCA_MULTIARCH_HOST inline KdTreeSparseBase(PointUserContainer&& points, IndexUserContainer sampling) : Base()
        {
            this->buildWithSampling(std::forward<PointUserContainer>(points), std::move(sampling));
        }
 
        using Base::buildWithSampling;
    };
 
#include "./kdTree.hpp"
} // namespace Ponca
 
template <typename Traits>
PONCA_MULTIARCH_HOST std::ostream& operator<<(std::ostream& os, const Ponca::KdTreeBase<Traits>& kdtree)
{
    kdtree.print(os);
    return os;
}
 
#undef WRITE_TRAITS