Skip to content

Commit

Permalink
[spatialpartitionning] Factorize kdtree descent code in kdTreeQuery
Browse files Browse the repository at this point in the history
  • Loading branch information
@nmellado
nmellado committed Jul 26, 2023
1 parent 81130b0 commit 360e4ea
Show file tree
Hide file tree
Showing 7 changed files with 138 additions and 334 deletions.
63 changes: 6 additions & 57 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestIndexQuery.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,61 +24,10 @@ auto KdTreeKNearestIndexQuery<Traits>::end()
template<typename Traits>
void KdTreeKNearestIndexQuery<Traits>::search()
{
const auto& nodes = QueryAccelType::m_kdtree->node_data();
const auto& points = QueryAccelType::m_kdtree->point_data();
const auto& indices = QueryAccelType::m_kdtree->index_data();
const auto& point = points[QueryType::input()].pos();

if (nodes.empty() || points.empty() || indices.empty())
throw std::invalid_argument("Empty KdTree");

while(!QueryAccelType::m_stack.empty())
{
auto& qnode = QueryAccelType::m_stack.top();
const auto& node = nodes[qnode.index];

if(qnode.squared_distance < QueryType::descentDistanceThreshold())
{
if(node.is_leaf())
{
QueryAccelType::m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
for(IndexType i=start; i<end; ++i)
{
IndexType idx = indices[i];
if(QueryType::input() == idx) continue;

Scalar d = (point - points[idx].pos()).squaredNorm();

if(d < QueryType::descentDistanceThreshold())
{
QueryType::m_queue.push({idx, d});
}
}
}
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
QueryAccelType::m_stack.push();
if(newOff < 0)
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
}
else
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
}
QueryAccelType::m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
}
}
else
{
QueryAccelType::m_stack.pop();
}
}
KdTreeQuery<Traits>::search_internal(QueryAccelType::m_kdtree->point_data()[QueryType::input()].pos(),
[](IndexType, IndexType){},
[this](){return QueryType::descentDistanceThreshold();},
[this](IndexType idx){return QueryType::input() == idx;},
[this](IndexType idx, IndexType, Scalar d){QueryType::m_queue.push({idx, d}); return false;}
);
}
61 changes: 6 additions & 55 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestPointQuery.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,59 +24,10 @@ auto KdTreeKNearestPointQuery<Traits>::end()
template <typename Traits>
void KdTreeKNearestPointQuery<Traits>::search()
{
const auto& nodes = QueryAccelType::m_kdtree->node_data();
const auto& points = QueryAccelType::m_kdtree->point_data();
const auto& indices = QueryAccelType::m_kdtree->index_data();
const auto& point = QueryType::input();

if (nodes.empty() || points.empty() || indices.empty())
throw std::invalid_argument("Empty KdTree");

while(!QueryAccelType::m_stack.empty())
{
auto& qnode = QueryAccelType::m_stack.top();
const auto& node = nodes[qnode.index];

if(qnode.squared_distance < QueryType::descentDistanceThreshold())
{
if(node.is_leaf())
{
QueryAccelType::m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
for(IndexType i=start; i<end; ++i)
{
IndexType idx = indices[i];
Scalar d = (point - points[idx].pos()).squaredNorm();

if(d < QueryType::descentDistanceThreshold())
{
QueryType::m_queue.push({idx, d});
}
}
}
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
QueryAccelType::m_stack.push();
if(newOff < 0)
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
}
else
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
}
QueryAccelType::m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
}
}
else
{
QueryAccelType::m_stack.pop();
}
}
KdTreeQuery<Traits>::search_internal(QueryType::input(),
[](IndexType, IndexType){},
[this](){return QueryType::descentDistanceThreshold();},
[](IndexType){return false;},
[this](IndexType idx, IndexType, Scalar d){QueryType::m_queue.push({idx, d}); return false;}
);
}
69 changes: 11 additions & 58 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestIndexQuery.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,62 +24,15 @@ auto KdTreeNearestIndexQuery<Traits>::end()
template <typename Traits>
void KdTreeNearestIndexQuery<Traits>::search()
{
const auto& nodes = QueryAccelType::m_kdtree->node_data();
const auto& points = QueryAccelType::m_kdtree->point_data();
const auto& indices = QueryAccelType::m_kdtree->index_data();
const auto& point = points[QueryType::input()].pos();

if (nodes.empty() || points.empty() || indices.empty())
throw std::invalid_argument("Empty KdTree");

while(!QueryAccelType::m_stack.empty())
{
auto& qnode = QueryAccelType::m_stack.top();
const auto& node = nodes[qnode.index];

if(qnode.squared_distance < QueryType::descentDistanceThreshold())
{
if(node.is_leaf())
{
QueryAccelType::m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
for(IndexType i=start; i<end; ++i)
{
IndexType idx = indices[i];
if(QueryType::input() == idx) continue;

Scalar d = (point - points[idx].pos()).squaredNorm();

if(d < QueryType::descentDistanceThreshold())
{
QueryType::m_nearest = idx;
QueryType::m_squared_distance = d;
}
}
}
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
QueryAccelType::m_stack.push();
if(newOff < 0)
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
}
else
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
}
QueryAccelType::m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
}
}
else
{
QueryAccelType::m_stack.pop();
}
}
KdTreeQuery<Traits>::search_internal(QueryAccelType::m_kdtree->point_data()[QueryType::input()].pos(),
[](IndexType, IndexType){},
[this](){return QueryType::descentDistanceThreshold();},
[this](IndexType idx){return QueryType::input() == idx;},
[this](IndexType idx, IndexType, Scalar d)
{
QueryType::m_nearest = idx;
QueryType::m_squared_distance = d;
return false;
}
);
}
67 changes: 11 additions & 56 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestPointQuery.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,60 +24,15 @@ auto KdTreeNearestPointQuery<Traits>::end()
template <typename Traits>
void KdTreeNearestPointQuery<Traits>::search()
{
const auto& nodes = QueryAccelType::m_kdtree->node_data();
const auto& points = QueryAccelType::m_kdtree->point_data();
const auto& indices = QueryAccelType::m_kdtree->index_data();
const auto& point = QueryType::input();

if (nodes.empty() || points.empty() || indices.empty())
throw std::invalid_argument("Empty KdTree");

while(!QueryAccelType::m_stack.empty())
{
auto& qnode = QueryAccelType::m_stack.top();
const auto& node = nodes[qnode.index];

if(qnode.squared_distance < QueryType::descentDistanceThreshold())
{
if(node.is_leaf())
{
QueryAccelType::m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
for(IndexType i=start; i<end; ++i)
{
IndexType idx = indices[i];
Scalar d = (point - points[idx].pos()).squaredNorm();

if(d < QueryType::descentDistanceThreshold())
{
QueryType::m_nearest = idx;
QueryType::m_squared_distance = d;
}
}
}
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
QueryAccelType::m_stack.push();
if(newOff < 0)
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
}
else
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
}
QueryAccelType::m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
}
}
else
{
QueryAccelType::m_stack.pop();
}
}
KdTreeQuery<Traits>::search_internal(QueryType::input(),
[](IndexType, IndexType){},
[this](){return QueryType::descentDistanceThreshold();},
[](IndexType){return false;},
[this](IndexType idx, IndexType, Scalar d)
{
QueryType::m_nearest = idx;
QueryType::m_squared_distance = d;
return false;
}
);
}
73 changes: 17 additions & 56 deletions Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeRangeIndexQuery.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,11 @@ auto KdTreeRangeIndexQuery<Traits>::end() -> Iterator
template<typename Traits>
void KdTreeRangeIndexQuery<Traits>::advance(Iterator& it)
{
const auto& nodes = QueryAccelType::m_kdtree->node_data();
const auto& points = QueryAccelType::m_kdtree->point_data();
const auto& indices = QueryAccelType::m_kdtree->index_data();
const auto& point = points[QueryType::input()].pos();

if (nodes.empty() || points.empty() || indices.empty())
if (points.empty() || indices.empty())
throw std::invalid_argument("Empty KdTree");

for(IndexType i=it.m_start; i<it.m_end; ++i)
Expand All @@ -45,58 +44,20 @@ void KdTreeRangeIndexQuery<Traits>::advance(Iterator& it)
}
}

while(!QueryAccelType::m_stack.empty())
{
auto& qnode = QueryAccelType::m_stack.top();
const auto& node = nodes[qnode.index];

if(qnode.squared_distance < QueryType::descentDistanceThreshold())
{
if(node.is_leaf())
{
QueryAccelType::m_stack.pop();
IndexType start = node.leaf.start;
IndexType end = node.leaf.start + node.leaf.size;
it.m_start = start;
it.m_end = end;
for(IndexType i=start; i<end; ++i)
{
IndexType idx = indices[i];
if(idx == QueryType::input()) continue;

Scalar d = (point - points[idx].pos()).squaredNorm();

if(d < QueryType::descentDistanceThreshold())
{
it.m_index = idx;
it.m_start = i+1;
return;
}
}
}
else
{
// replace the stack top by the farthest and push the closest
Scalar newOff = point[node.inner.dim] - node.inner.split_value;
QueryAccelType::m_stack.push();
if(newOff < 0)
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id;
qnode.index = node.inner.first_child_id+1;
}
else
{
QueryAccelType::m_stack.top().index = node.inner.first_child_id+1;
qnode.index = node.inner.first_child_id;
}
QueryAccelType::m_stack.top().squared_distance = qnode.squared_distance;
qnode.squared_distance = newOff*newOff;
}
}
else
{
QueryAccelType::m_stack.pop();
}
}
it.m_index = static_cast<IndexType>(points.size());
if (KdTreeQuery<Traits>::search_internal(point,
[&it](IndexType start, IndexType end)
{
it.m_start = start;
it.m_end = end;
},
[this](){return QueryType::descentDistanceThreshold();},
[this](IndexType idx){return QueryType::input() == idx;},
[&it](IndexType idx, IndexType i, Scalar)
{
it.m_index = idx;
it.m_start = i+1;
return true;
}
))
it.m_index = static_cast<IndexType>(points.size());
}
Loading

0 comments on commit 360e4ea

Please sign in to comment.