diff --git a/Ponca/src/SpatialPartitioning/KdTree/kdTreeTraits.h b/Ponca/src/SpatialPartitioning/KdTree/kdTreeTraits.h index 88363f4c2..e96675cbb 100644 --- a/Ponca/src/SpatialPartitioning/KdTree/kdTreeTraits.h +++ b/Ponca/src/SpatialPartitioning/KdTree/kdTreeTraits.h @@ -83,6 +83,24 @@ struct KdTreeDefaultLeafNode /*! * \brief The node type used by default by the kd-tree. + * + * It is possible to modify the Inner and Leaf node types by inheritance. For instance, to add a Bounding box to inner + * nodes, define a custom inner node type: + * + * \snippet ponca_customize_kdtree.cpp CustomInnerNodeDefinition + * + * Define a custom node type to use it, and expose custom data (inner/leaf node are not exposed directly): + * + * \snippet ponca_customize_kdtree.cpp CustomNodeDefinition + * + * To use in the KdTree, define a type using the custom node: + * + * \snippet ponca_customize_kdtree.cpp KdTreeTypeWithCustomNode + * + * The added attribute can be accessed + * + * \snippet ponca_customize_kdtree.cpp ReadCustomProperties + * */ template +#include +#include +#include + +struct DataPoint +{ + enum {Dim = 3}; + using Scalar = float; + using VectorType = Eigen::Vector; + inline const auto& pos() const {return m_pos;} + VectorType m_pos; +}; + +//! [CustomInnerNodeDefinition] +template > +struct MyKdTreeInnerNode : public Ponca::KdTreeDefaultInnerNode { + using AabbType = _AabbType; + AabbType m_aabb{}; +}; +//! [CustomInnerNodeDefinition] + +//! [CustomNodeDefinition] +template +struct MyKdTreeNode : Ponca::KdTreeCustomizableNode> { + + using Base = Ponca::KdTreeCustomizableNode>; + using AabbType = typename Base::AabbType; + + void configure_range(Index start, Index size, const AabbType &aabb) + { + Base::configure_range(start, size, aabb); + if (! Base::is_leaf() ) + { + Base::getAsInner().m_aabb = aabb; + } + } + [[nodiscard]] inline std::optional getAabb() const { + if (! Base::is_leaf()) + return Base::getAsInner().m_aabb; + else + return std::optional(); + } +}; +//! [CustomNodeDefinition] + +int main() +{ + // generate N random points + constexpr int N = 1e5; + std::vector points(N); + std::generate(points.begin(), points.end(), [](){ + return DataPoint{100 * DataPoint::VectorType::Random()};}); + +//! [KdTreeTypeWithCustomNode] + using CustomKdTree = Ponca::KdTreeBase>; +//! [KdTreeTypeWithCustomNode] + + // build the k-d tree + const CustomKdTree kdtree(points); + + // neighbor searches are done below from these arbitrary index and point + const int query_idx = 10; + const DataPoint::VectorType query_pt{-10.0, 0.5, 75.0}; + + // + // nearest neighbor search + // + std::cout << "the nearest neighbor of the point at index " << query_idx << " is at index " + << *kdtree.nearest_neighbor(query_idx).begin() << std::endl; + std::cout << "the nearest neighbor of the point (" << query_pt.transpose() << ") is at index " + << *kdtree.nearest_neighbor(query_pt).begin() << std::endl; + + //! [ReadCustomProperties] + auto bbox = kdtree.node_data()[0].getAabb(); + if (bbox) { + std::cout << "Root bounding box is as follows: \n" + << " Center: " << bbox->center() + << " Diagonal: " << bbox->diagonal() + << std::endl; + } + //! [ReadCustomProperties] + + return 1; +}