Skip to content

Commit

Permalink
[SpatialPartitioning] Change the interface of kd-tree nodes
Browse files Browse the repository at this point in the history
  • Loading branch information
Amael Marquez committed Nov 30, 2023
1 parent 822a725 commit e8fe4df
Show file tree
Hide file tree
Showing 4 changed files with 148 additions and 153 deletions.
14 changes: 7 additions & 7 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeQuery.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,8 @@ class KdTreeQuery
if(node.is_leaf())
{
m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
IndexType start = node.leaf_start();
IndexType end = node.leaf_start() + node.leaf_size();
prepareLeafTraversal(start, end);
for(IndexType i=start; i<end; ++i)
{
Expand All @@ -80,17 +80,17 @@ class KdTreeQuery
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
Scalar newOff = point[node.inner_split_dim()] - node.inner_split_value();
m_stack.push();
if(newOff < 0)
{
m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
m_stack.top().index = node.inner_first_child_id();
qnode.index = node.inner_first_child_id()+1;
}
else
{
m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
m_stack.top().index = node.inner_first_child_id()+1;
qnode.index = node.inner_first_child_id();
}
m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
Expand Down
33 changes: 18 additions & 15 deletions Ponca/src/SpatialPartitioning/KdTree/kdTree.h
Original file line number Diff line number Diff line change
Expand Up @@ -52,35 +52,38 @@ template <typename Traits>
class KdTreeBase
{
public:
using DataPoint = typename Traits::DataPoint; ///< DataPoint given by user via Traits
using Scalar = typename DataPoint::Scalar; ///< Scalar given by user via DataPoint
using VectorType = typename DataPoint::VectorType; ///< VectorType given by user via DataPoint
using AabbType = typename Traits::AabbType; ///< Bounding box type given by user via DataPoint
using DataPoint = typename Traits::DataPoint; ///< DataPoint given by user via Traits
using IndexType = typename Traits::IndexType; ///< Type used to index points into the PointContainer
using LeafSizeType = typename Traits::LeafSizeType; ///< Type used to store the size of leaf nodes

using IndexType = typename Traits::IndexType;
using PointContainer = typename Traits::PointContainer; ///< Container for DataPoint used inside the KdTree
using IndexContainer = typename Traits::IndexContainer; ///< Container for indices used inside the KdTree

using NodeIndexType = typename Traits::NodeIndexType; ///< Type used to index nodes into the NodeContainer
using NodeType = typename Traits::NodeType; ///< Type of nodes used inside the KdTree
using NodeContainer = typename Traits::NodeContainer; ///< Container for nodes used inside the KdTree

using NodeType = typename NodeContainer::value_type;
using NodeCountType = typename NodeContainer::size_type;
using LeafSizeType = typename NodeType::LeafSizeType;
using Scalar = typename DataPoint::Scalar; ///< Scalar given by user via DataPoint
using VectorType = typename DataPoint::VectorType; ///< VectorType given by user via DataPoint
using AabbType = typename NodeType::AabbType; ///< Bounding box type given by user via NodeType

enum
{
/*!
* The maximum number of points that can be stored in the kd-tree, considering how many
* bits the inner nodes use to store their children indices.
* The maximum number of points that can be stored in the kd-tree,
* considering the bit width of the index type.
*/
MAX_POINT_COUNT = 2 << NodeType::InnerType::INDEX_BITS,
MAX_POINT_COUNT = std::size_t(2) << sizeof(IndexType)*8,
};

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(Traits::MAX_DEPTH > 0, "Max depth must be strictly positive");

Expand Down Expand Up @@ -188,7 +191,7 @@ class KdTreeBase

// Accessors ---------------------------------------------------------------
public:
inline NodeCountType node_count() const
inline NodeIndexType node_count() const
{
return m_nodes.size();
}
Expand All @@ -203,7 +206,7 @@ class KdTreeBase
return (IndexType)m_points.size();
}

inline NodeCountType leaf_count() const
inline NodeIndexType leaf_count() const
{
return m_leaf_count;
}
Expand Down Expand Up @@ -255,7 +258,7 @@ class KdTreeBase

// Internal ----------------------------------------------------------------
protected:
inline void build_rec(NodeCountType node_id, IndexType start, IndexType end, int level);
inline void build_rec(NodeIndexType node_id, IndexType start, IndexType end, int level);
inline IndexType partition(IndexType start, IndexType end, int dim, Scalar value);

// Query -------------------------------------------------------------------
Expand Down Expand Up @@ -298,7 +301,7 @@ public :
IndexContainer m_indices;

LeafSizeType m_min_cell_size; ///< Minimal number of points per leaf
NodeCountType m_leaf_count; ///< Number of leaves in the Kdtree (computed during construction)
NodeIndexType m_leaf_count; ///< Number of leaves in the Kdtree (computed during construction)
};

#include "./kdTree.hpp"
Expand Down
36 changes: 16 additions & 20 deletions Ponca/src/SpatialPartitioning/KdTree/kdTree.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -95,25 +95,25 @@ bool KdTreeBase<Traits>::valid() const
b[idx] = true;
}

for(NodeCountType n=0;n<node_count();++n)
for(NodeIndexType n=0;n<node_count();++n)
{
const NodeType& node = m_nodes.operator[](n);
if(node.is_leaf())
{
if(index_count() <= node.leaf.start || index_count() < node.leaf.start+node.leaf.size)
if(index_count() <= node.leaf_start() || index_count() < node.leaf_start()+node.leaf_size())
{
PONCA_DEBUG_ERROR;
return false;
}
}
else
{
if(node.inner.dim < 0 || 2 < node.inner.dim)
if(node.inner_dim() < 0 || 2 < node.inner_dim())
{
PONCA_DEBUG_ERROR;
return false;
}
if(node_count() <= node.inner.first_child_id || node_count() <= node.inner.first_child_id+1)
if(node_count() <= node.inner_first_child_id() || node_count() <= node.inner_first_child_id()+1)
{
PONCA_DEBUG_ERROR;
return false;
Expand All @@ -136,51 +136,47 @@ std::string KdTreeBase<Traits>::to_string() const
str << " " << i << ": " << m_indices.operator[](i) << "\n";
}
str << "nodes (" << node_count() << ") :\n";
for(NodeCountType n=0; n< node_count(); ++n)
for(NodeIndexType n=0; n< node_count(); ++n)
{
const NodeType& node = m_nodes.operator[](n);
if(node.is_leaf())
{
auto end = node.leaf.start + node.leaf.size;
str << " leaf: start=" << node.leaf.start << " end=" << end << " (size=" << node.leaf.size << ")\n";
auto end = node.leaf_start() + node.leaf_size();
str << " leaf: start=" << node.leaf_start() << " end=" << end << " (size=" << node.leaf_size() << ")\n";
}
else
{
str << " node: dim=" << node.inner.dim << " split=" << node.inner.split_value << " child=" << node.inner.first_child_id << "\n";
str << " node: dim=" << node.inner_dim() << " split=" << node.inner_split_value() << " child=" << node.inner_first_child_id() << "\n";
}
}
return str.str();
}

template<typename Traits>
void KdTreeBase<Traits>::build_rec(NodeCountType node_id, IndexType start, IndexType end, int level)
void KdTreeBase<Traits>::build_rec(NodeIndexType node_id, IndexType start, IndexType end, int level)
{
NodeType& node = m_nodes[node_id];
AabbType aabb;
for(IndexType i=start; i<end; ++i)
aabb.extend(m_points[m_indices[i]].pos());

node.set_is_leaf(end-start <= m_min_cell_size || level >= Traits::MAX_DEPTH);
node.configure_range(start, end-start, aabb);
if (node.is_leaf())
{
node.leaf.start = start;
node.leaf.size = static_cast<LeafSizeType>(end-start);
++m_leaf_count;
}
else
{
int dim = 0;
(Scalar(0.5) * (aabb.max() - aabb.min())).maxCoeff(&dim);
node.inner.dim = dim;
node.inner.split_value = aabb.center()[dim];

IndexType mid_id = this->partition(start, end, dim, node.inner.split_value);
node.inner.first_child_id = m_nodes.size();
int split_dim = 0;
(Scalar(0.5) * aabb.diagonal()).maxCoeff(&split_dim);
node.configure_inner(aabb.center()[split_dim], m_nodes.size(), split_dim);
m_nodes.emplace_back();
m_nodes.emplace_back();

build_rec(node.inner.first_child_id, start, mid_id, level+1);
build_rec(node.inner.first_child_id+1, mid_id, end, level+1);
IndexType mid_id = this->partition(start, end, split_dim, node.inner_split_value());
build_rec(node.inner_first_child_id(), start, mid_id, level+1);
build_rec(node.inner_first_child_id()+1, mid_id, end, level+1);
}
}

Expand Down
Loading

0 comments on commit e8fe4df

Please sign in to comment.