Skip to content

Commit

Permalink
Merge pull request #128 from poncateam/easeKdTreeNodesChange
Browse files Browse the repository at this point in the history
Simplify KdTree Node type specialization
  • Loading branch information
nmellado authored Dec 15, 2023
2 parents 280968b + 10631b9 commit 632e0e2
Show file tree
Hide file tree
Showing 5 changed files with 180 additions and 28 deletions.
1 change: 1 addition & 0 deletions CHANGELOG
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ Current head (v.1.2 RC)
- API
- [spatialPartitioning] Optimize memory use in knngraph queries (#104)
- [spatialPartitioning] Clean KdTree API (#122)
- [spatialPartitioning] Simplify node customization (#128)
- [fitting] Mark `Base` type as protected instead of private in CRTP classes (#119)
- [fitting] Improve KdTreeNodes API by hiding internal memory layout, improve methods naming (#120)

Expand Down
101 changes: 75 additions & 26 deletions Ponca/src/SpatialPartitioning/KdTree/kdTreeTraits.h
Original file line number Diff line number Diff line change
Expand Up @@ -69,29 +69,49 @@ struct KdTreeDefaultInnerNode
INDEX_BITS = sizeof(UIndex)*8 - DIM_BITS,
};

Scalar split_value;
Scalar split_value{0};
UIndex first_child_id : INDEX_BITS;
UIndex split_dim : DIM_BITS;
};

template <typename Index, typename Size>
struct KdTreeDefaultLeafNode
{
Index start;
Size size;
Index start{0};
Size size{0};
};

/*!
* \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 <typename Index, typename NodeIndex, typename DataPoint,
typename LeafSize = Index>
class KdTreeDefaultNode
typename LeafSize = Index,
typename _InnerNodeType = KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
typename _LeafNodeType = KdTreeDefaultLeafNode<Index, LeafSize>>
class KdTreeCustomizableNode
{
private:
using Scalar = typename DataPoint::Scalar;
using LeafType = KdTreeDefaultLeafNode<Index, LeafSize>;
using InnerType = KdTreeDefaultInnerNode<NodeIndex, Scalar, DataPoint::Dim>;
using InnerType = _InnerNodeType;
using LeafType = _LeafNodeType;

public:
enum
Expand All @@ -110,10 +130,8 @@ class KdTreeDefaultNode
* `DataPoint::VectorType`.
*/
using AabbType = Eigen::AlignedBox<Scalar, DataPoint::Dim>;

KdTreeDefaultNode() = default;

bool is_leaf() const { return m_is_leaf; }
[[nodiscard]] bool is_leaf() const { return m_is_leaf; }
void set_is_leaf(bool is_leaf) { m_is_leaf = is_leaf; }

/*!
Expand All @@ -132,8 +150,8 @@ class KdTreeDefaultNode
{
if (m_is_leaf)
{
m_leaf.start = start;
m_leaf.size = (LeafSize)size;
data.m_leaf.start = start;
data.m_leaf.size = (LeafSize)size;
}
}

Expand All @@ -150,32 +168,32 @@ class KdTreeDefaultNode
{
if (!m_is_leaf)
{
m_inner.split_value = split_value;
m_inner.first_child_id = first_child_id;
m_inner.split_dim = split_dim;
data.m_inner.split_value = split_value;
data.m_inner.first_child_id = first_child_id;
data.m_inner.split_dim = split_dim;
}
}

/*!
* \brief The start index of the range of the leaf node in the sample
* index array.
*/
Index leaf_start() const { return m_leaf.start; }
[[nodiscard]] Index leaf_start() const { return data.m_leaf.start; }

/*!
* \brief The size of the range of the leaf node in the sample index array.
*/
LeafSize leaf_size() const { return m_leaf.size; }
[[nodiscard]] LeafSize leaf_size() const { return data.m_leaf.size; }

/*!
* \brief The position of the AABB split of the inner node.
*/
Scalar inner_split_value() const { return m_inner.split_value; }
[[nodiscard]] Scalar inner_split_value() const { return data.m_inner.split_value; }

/*!
* \brief Which axis the split of the AABB of the inner node was done on.
*/
int inner_split_dim() const { return (int)m_inner.split_dim; }
[[nodiscard]] int inner_split_dim() const { return (int)data.m_inner.split_dim; }

/*!
* \brief The index of the first child of the node in the node array of the
Expand All @@ -184,21 +202,52 @@ class KdTreeDefaultNode
* \note The second child is stored directly after the first in the array
* (i.e. `first_child_id + 1`).
*/
Index inner_first_child_id() const { return (Index)m_inner.first_child_id; }
[[nodiscard]] Index inner_first_child_id() const { return (Index)data.m_inner.first_child_id; }

protected:
[[nodiscard]] inline LeafType& getAsLeaf() { return data.m_leaf; }
[[nodiscard]] inline InnerType& getAsInner() { return data.m_inner; }
[[nodiscard]] inline const LeafType& getAsLeaf() const { return data.m_leaf; }
[[nodiscard]] inline const InnerType& getAsInner() const { return data.m_inner; }

private:
bool m_is_leaf;
union
bool m_is_leaf{true};
union Data
{
KdTreeDefaultLeafNode<Index, LeafSize> m_leaf;
KdTreeDefaultInnerNode<NodeIndex, Scalar, DataPoint::Dim> m_inner;
// We need an explicit constructor here, see https://stackoverflow.com/a/70428826
constexpr Data() : m_leaf() {}
// Needed to satisfy MoveInsertable requirement https://en.cppreference.com/w/cpp/named_req/MoveInsertable
constexpr Data(const Data&d) : m_leaf(d.m_leaf) {}

~Data() {}
LeafType m_leaf;
InnerType m_inner;
};
Data data;
};

template <typename Index, typename NodeIndex, typename DataPoint,
typename LeafSize = Index>
struct KdTreeDefaultNode : public KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
KdTreeDefaultLeafNode<Index, LeafSize>> {
using Base = KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
KdTreeDefaultLeafNode<Index, LeafSize>>;
};

/*!
* \brief The default traits type used by the kd-tree.
*
* \see KdTreeCustomizableNode Helper class to modify Inner/Leaf nodes without redefining a Trait class
*
* \tparam _NodeType Type used to store nodes, set by default to #KdTreeDefaultNode
*/
template <typename _DataPoint>
template <typename _DataPoint,
template <typename /*Index*/,
typename /*NodeIndex*/,
typename /*DataPoint*/,
typename /*LeafSize*/> typename _NodeType = KdTreeDefaultNode>
struct KdTreeDefaultTraits
{
enum
Expand Down Expand Up @@ -228,7 +277,7 @@ struct KdTreeDefaultTraits

// Nodes
using NodeIndexType = std::size_t;
using NodeType = KdTreeDefaultNode<IndexType, NodeIndexType, DataPoint, LeafSizeType>;
using NodeType = _NodeType<IndexType, NodeIndexType, DataPoint, LeafSizeType>;
using NodeContainer = std::vector<NodeType>;
};
} // namespace Ponca
3 changes: 2 additions & 1 deletion doc/src/ponca_module_spatialpartitioning.mdoc
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,8 @@ fit.computeWithIds( myDataStructure.range_neighbors(fitInitPos, scale), vectorPo

\subsection spatialpartitioning_kdtree_extending Extending KdTree
Ponca::KdTreeBase is a customizable version of Ponca::KdTree, which can be controlled using `Traits`.
See KdTreeDefaultTraits for customization API.
See KdTreeDefaultTraits and KdTreeCustomizableNode for customization APIs. See also:
- `examples/cpp/ponca_customize_kdtree.cpp`

\subsection spatialpartitioning_kdtree_implementation Implementation details

Expand Down
10 changes: 9 additions & 1 deletion examples/cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,20 @@ add_dependencies(ponca-examples ponca_fit_line)
ponca_handle_eigen_dependency(ponca_fit_line)

set(ponca_neighbor_search_SRCS
ponca_neighbor_search.cpp
ponca_neighbor_search.cpp
)
add_executable(ponca_neighbor_search ${ponca_neighbor_search_SRCS})
target_include_directories(ponca_neighbor_search PRIVATE ${PONCA_src_ROOT})
add_dependencies(ponca-examples ponca_neighbor_search)
ponca_handle_eigen_dependency(ponca_neighbor_search)

set(ponca_customize_kdtree_SRCS
ponca_customize_kdtree.cpp
)
add_executable(ponca_customize_kdtree ${ponca_customize_kdtree_SRCS})
target_include_directories(ponca_customize_kdtree PRIVATE ${PONCA_src_ROOT})
add_dependencies(ponca-examples ponca_customize_kdtree)
ponca_handle_eigen_dependency(ponca_customize_kdtree)

add_subdirectory(pcl)
add_subdirectory(nanoflann)
93 changes: 93 additions & 0 deletions examples/cpp/ponca_customize_kdtree.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
/*
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/.
*/

#include <iostream>
#include <optional>
#include <Ponca/SpatialPartitioning>
#include <Eigen/Core>

struct DataPoint
{
enum {Dim = 3};
using Scalar = float;
using VectorType = Eigen::Vector<Scalar,Dim>;
inline const auto& pos() const {return m_pos;}
VectorType m_pos;
};

//! [CustomInnerNodeDefinition]
template <typename NodeIndex, typename Scalar, int DIM, typename _AabbType = Eigen::AlignedBox<Scalar, DIM>>
struct MyKdTreeInnerNode : public Ponca::KdTreeDefaultInnerNode<NodeIndex, Scalar, DIM> {
using AabbType = _AabbType;
AabbType m_aabb{};
};
//! [CustomInnerNodeDefinition]

//! [CustomNodeDefinition]
template <typename Index, typename NodeIndex, typename DataPoint, typename LeafSize = Index>
struct MyKdTreeNode : Ponca::KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
MyKdTreeInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>> {

using Base = Ponca::KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
MyKdTreeInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>>;
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<AabbType> getAabb() const {
if (! Base::is_leaf())
return Base::getAsInner().m_aabb;
else
return std::optional<AabbType>();
}
};
//! [CustomNodeDefinition]

int main()
{
// generate N random points
constexpr int N = 1e5;
std::vector<DataPoint> points(N);
std::generate(points.begin(), points.end(), [](){
return DataPoint{100 * DataPoint::VectorType::Random()};});

//! [KdTreeTypeWithCustomNode]
using CustomKdTree = Ponca::KdTreeBase<Ponca::KdTreeDefaultTraits<DataPoint,MyKdTreeNode>>;
//! [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;
}

0 comments on commit 632e0e2

Please sign in to comment.