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 <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
template <typename DataPoint>
struct KdTree : public Ponca::KdTreeBase<KdTreeDefaultTraits<DataPoint>>{};
#else
template <typename DataPoint>
using KdTree = KdTreeBase<KdTreeDefaultTraits<DataPoint>>; // prefer alias to avoid redefining methods
#endif
 
 
#ifdef PARSED_WITH_DOXYGEN
template <typename DataPoint>
struct KdTreeDense : public Ponca::KdTreeDenseBase<KdTreeDefaultTraits<DataPoint>>{};
#else
template <typename DataPoint>
using KdTreeDense = KdTreeDenseBase<KdTreeDefaultTraits<DataPoint>>; // prefer alias to avoid redefining methods
#endif
 
#ifdef PARSED_WITH_DOXYGEN
template <typename DataPoint>
struct KdTreeSparse : Ponca::KdTreeSparseBase<KdTreeDefaultTraits<DataPoint>>{};
#else
template <typename DataPoint>
using KdTreeSparse = KdTreeSparseBase<KdTreeDefaultTraits<DataPoint>>;
#endif
 
template <typename Traits>
class KdTreeBase
{
public:
    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; 
 
    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;
 
    static_assert(std::is_same<typename PointContainer::value_type, DataPoint>::value,
        "PointContainer must contain DataPoints");
    
    // Queries use a value of -1 for invalid indices
    static_assert(std::is_signed<IndexType>::value, "Index type must be signed");
 
    static_assert(std::is_same<typename IndexContainer::value_type, IndexType>::value, "Index type mismatch");
    static_assert(std::is_same<typename NodeContainer::value_type, NodeType>::value, "Node type mismatch");
 
    static_assert(MAX_DEPTH > 0, "Max depth must be strictly positive");
 
    // Construction ------------------------------------------------------------
public:
    template<typename PointUserContainer, typename Converter>
    inline void build(PointUserContainer&& points, Converter c);
 
    struct DefaultConverter
    {
        template <typename Input>
        inline void operator()(Input&& i, PointContainer& o)
        {
            using InputContainer = typename std::remove_reference<Input>::type;
            if constexpr (std::is_same<InputContainer, PointContainer>::value)
                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>
    inline void build(PointUserContainer&& points)
    {
        build(std::forward<PointUserContainer>(points), DefaultConverter());
    }
 
    inline void clear();
 
    // Accessors ---------------------------------------------------------------
public:
    inline NodeIndexType node_count() const
    {
        return m_nodes.size();
    }
 
    inline IndexType sample_count() const
    {
        return (IndexType)m_indices.size();
    }
 
    inline IndexType pointCount() const
    {
        return (IndexType)m_points.size();
    }
 
    inline NodeIndexType leaf_count() const
    {
        return m_leaf_count;
    }
 
    inline PointContainer& points()
    {
        return m_points;
    };
 
    inline const PointContainer& points() const
    {
        return m_points;
    };
 
    inline const NodeContainer& nodes() const
    {
        return m_nodes;
    }
 
    inline const IndexContainer& samples() const
    {
        return m_indices;
    }
 
    // Parameters --------------------------------------------------------------
public:
    inline LeafSizeType minCellSize() const
    {
        return m_min_cell_size;
    }
 
    inline void setMinCellSize(LeafSizeType min_cell_size)
    {
        PONCA_DEBUG_ASSERT(min_cell_size > 0);
        m_min_cell_size = min_cell_size;
    }
 
    // Index mapping -----------------------------------------------------------
public:
    inline IndexType pointFromSample(IndexType sample_index) const
    {
        return m_indices[sample_index];
    }
 
    inline DataPoint& pointDataFromSample(IndexType sample_index)
    {
        return m_points[pointFromSample(sample_index)];
    }
    
    inline const DataPoint& pointDataFromSample(IndexType sample_index) const
    {
        return m_points[pointFromSample(sample_index)];
    }
 
    // Query -------------------------------------------------------------------
public :
    KdTreeKNearestPointQuery<Traits> kNearestNeighbors(const VectorType& point, IndexType k) const
    {
        return KdTreeKNearestPointQuery<Traits>(this, k, point);
    }
 
    KdTreeKNearestIndexQuery<Traits> kNearestNeighbors(IndexType index, IndexType k) const
    {
        return KdTreeKNearestIndexQuery<Traits>(this, k, index);
    }
 
    KdTreeKNearestPointQuery<Traits> kNearestNeighborsQuery() const
    {
        return KdTreeKNearestPointQuery<Traits>(this, 0, VectorType::Zero());
    }
 
    KdTreeKNearestIndexQuery<Traits> kNearestNeighborsIndexQuery() const
    {
        return KdTreeKNearestIndexQuery<Traits>(this, 0, 0);
    }
 
    KdTreeNearestPointQuery<Traits> nearestNeighbor(const VectorType& point) const
    {
        return KdTreeNearestPointQuery<Traits>(this, point);
    }
 
 
    KdTreeNearestIndexQuery<Traits> nearestNeighbor(IndexType index) const
    {
        return KdTreeNearestIndexQuery<Traits>(this, index);
    }
 
    KdTreeNearestIndexQuery<Traits> nearestNeighborQuery() const
    {
        return KdTreeNearestIndexQuery<Traits>(this, VectorType::Zero());
    }
 
    KdTreeNearestIndexQuery<Traits> nearestNeighborIndexQuery() const
    {
        return KdTreeNearestIndexQuery<Traits>(this, 0);
    }
 
    KdTreeRangePointQuery<Traits> rangeNeighbors(const VectorType& point, Scalar r) const
    {
        return KdTreeRangePointQuery<Traits>(this, r, point);
    }
 
    KdTreeRangeIndexQuery<Traits> rangeNeighbors(IndexType index, Scalar r) const
    {
        return KdTreeRangeIndexQuery<Traits>(this, r, index);
    }
 
    KdTreeRangePointQuery<Traits> rangeNeighborsQuery() const
    {
        return KdTreeRangePointQuery<Traits>(this, 0, VectorType::Zero());
    }
 
    KdTreeRangeIndexQuery<Traits> rangeNeighborsIndexQuery() const
    {
        return KdTreeRangeIndexQuery<Traits>(this, 0, 0);
    }
    
    // Utilities ---------------------------------------------------------------
public:
    [[nodiscard]] inline bool valid() const;
    inline void print(std::ostream& os, bool verbose = false) const;
 
    // Data --------------------------------------------------------------------
protected:
    PointContainer m_points;
    NodeContainer m_nodes;
    IndexContainer m_indices;
 
    LeafSizeType m_min_cell_size {64}; 
    NodeIndexType m_leaf_count {0}; 
 
    // Internal ----------------------------------------------------------------
protected:
    inline KdTreeBase() = default;
 
    template<typename PointUserContainer, typename IndexUserContainer, typename Converter>
    inline void buildWithSampling(PointUserContainer&& points,
                                  IndexUserContainer sampling,
                                  Converter c);
 
    template<typename PointUserContainer, typename IndexUserContainer>
    inline void buildWithSampling(PointUserContainer&& points,
                                  IndexUserContainer sampling)
    {
        buildWithSampling(std::forward<PointUserContainer>(points), std::move(sampling), DefaultConverter());
    }
 
private:
    inline void buildRec(NodeIndexType node_id, IndexType start, IndexType end, int level);
    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:
    KdTreeDenseBase() = default;
 
    template<typename PointUserContainer>
    inline explicit KdTreeDenseBase(PointUserContainer&& points)
        : Base()
    {
        this->build(std::forward<PointUserContainer>(points));
    }
};
 
template <typename Traits>
class KdTreeSparseBase : public KdTreeBase<Traits>
{
private:
    using Base = KdTreeBase<Traits>;
 
public:
    static constexpr bool SUPPORTS_SUBSAMPLING = false;
 
    KdTreeSparseBase() = default;
 
    template<typename PointUserContainer>
    inline explicit KdTreeSparseBase(PointUserContainer&& points)
        : Base()
    {
        this->build(std::forward<PointUserContainer>(points));
    }
 
    template<typename PointUserContainer, typename IndexUserContainer>
    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>
std::ostream& operator<<(std::ostream& os, const Ponca::KdTreeBase<Traits>& kdtree)
{
    kdtree.print(os);
    return os;
}