You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
263 lines
10 KiB
263 lines
10 KiB
#ifndef ALGOIM_KDTREE_HPP
|
|
#define ALGOIM_KDTREE_HPP
|
|
|
|
/* algoim::KDTree<T,N> constructs a k-d tree data structure for a given collection of
|
|
points of type uvector<T,N>, where T is typically float or double, and N is
|
|
the dimension. This particular implementation of the k-d tree has been optimised for the
|
|
case that the points are situated on a smooth codimension-one surface by using coordinate
|
|
transformations that result in "tight" bounding boxes for some of the nodes in the tree.
|
|
|
|
Refer to the paper
|
|
R. I. Saye, High-order methods for computing distances to implicitly defined surfaces,
|
|
Communications in Applied Mathematics and Computational Science, 9(1), 107-141 (2014),
|
|
http://dx.doi.org/10.2140/camcos.2014.9.107
|
|
for more information. */
|
|
|
|
#include <vector>
|
|
#include <cassert>
|
|
#include "real.hpp"
|
|
#include "uvector.hpp"
|
|
#include "utility.hpp"
|
|
|
|
namespace algoim
|
|
{
|
|
namespace detail
|
|
{
|
|
// Returns the squared-distance from a point x to the hyperrectangle [min,max] (if x is inside the rectangle, the distance is
|
|
// zero)
|
|
template <typename T, int N>
|
|
T sqrDistBox(const uvector<T, N>& x, const uvector<T, N>& min, const uvector<T, N>& max)
|
|
{
|
|
T dsqr = T(0);
|
|
for (int i = 0; i < N; ++i)
|
|
if (x(i) < min(i))
|
|
dsqr += util::sqr(x(i) - min(i));
|
|
else if (x(i) > max(i))
|
|
dsqr += util::sqr(x(i) - max(i));
|
|
return dsqr;
|
|
}
|
|
} // namespace detail
|
|
|
|
template <typename T, int N>
|
|
class KDTree
|
|
{
|
|
std::vector<uvector<T, N>> points;
|
|
std::vector<int> index;
|
|
static constexpr int leaf_size = 16;
|
|
|
|
// A Node of the tree is a leaf node iff type=-1. If type >= 0, then a coordinate transform
|
|
// is applied. If type < -1, the node is a standard splitting node with two children and no transform
|
|
struct Node {
|
|
int type;
|
|
int i0, i1;
|
|
uvector<T, N> xmin, xmax;
|
|
};
|
|
|
|
std::vector<Node> nodes;
|
|
std::vector<uvector<uvector<T, N>, N>> transforms;
|
|
|
|
// Given a node and a range of points [lb,ub), build the tree
|
|
void build_tree(int nodeIndex, int lb, int ub, bool hasTransformed, int level)
|
|
{
|
|
assert(lb < ub);
|
|
Node& node = nodes[nodeIndex];
|
|
|
|
// Compute bounding box and mean
|
|
uvector<T, N> mean = node.xmin = node.xmax = points[index[lb]];
|
|
for (int i = lb + 1; i < ub; ++i) {
|
|
const uvector<T, N>& x = points[index[i]];
|
|
mean += x;
|
|
for (int j = 0; j < N; ++j) {
|
|
if (x(j) < node.xmin(j)) node.xmin(j) = x(j);
|
|
if (x(j) > node.xmax(j)) node.xmax(j) = x(j);
|
|
}
|
|
}
|
|
mean /= static_cast<T>(ub - lb);
|
|
|
|
// The node is a leaf iff point count is sufficiently small
|
|
if (ub - lb <= leaf_size) {
|
|
node.type = -1;
|
|
node.i0 = lb;
|
|
node.i1 = ub;
|
|
return;
|
|
}
|
|
|
|
// Splitting node: default to splitting along greatest extent
|
|
node.type = -2;
|
|
int axis = argmax(node.xmax - node.xmin);
|
|
|
|
// Evaluate possibility for coordinate transformation
|
|
if (!hasTransformed && level > 5 && ub - lb >= leaf_size * (1 << 2)) {
|
|
// Estimate normal
|
|
T holeRadiusSqr = util::sqr(0.05 * max(node.xmax - node.xmin));
|
|
uvector<T, N> n = static_cast<T>(0);
|
|
n(0) = 1.0;
|
|
for (int i = lb; i < ub; ++i) {
|
|
uvector<T, N> x = points[index[i]] - mean;
|
|
T msqr = sqrnorm(x);
|
|
if (msqr > holeRadiusSqr) n -= x * (dot(x, n) / msqr);
|
|
}
|
|
T msqr = sqrnorm(n);
|
|
if (msqr == 0.0)
|
|
n(0) = 1.0;
|
|
else
|
|
n /= sqrt(msqr);
|
|
|
|
// Compute alpha range
|
|
T minAlpha = std::numeric_limits<T>::max();
|
|
T maxAlpha = -std::numeric_limits<T>::max();
|
|
for (int i = lb; i < ub; ++i) {
|
|
T alpha = dot(points[index[i]], n);
|
|
if (alpha > maxAlpha) maxAlpha = alpha;
|
|
if (alpha < minAlpha) minAlpha = alpha;
|
|
}
|
|
|
|
if (maxAlpha - minAlpha < 0.1 * max(node.xmax - node.xmin)) {
|
|
// Perform transformation: calculate an orthonormal basis using the normal as first axis.
|
|
// A stable method for doing so is to compute the Householder matrix which maps n to ej,
|
|
// i.e., P = I - 2 uu^T where u = normalised(n - ej), where ej is the j-th basis vector,
|
|
// j chosen that that n != ej.
|
|
uvector<uvector<T, N>, N> axes;
|
|
int j = argmin(abs(n));
|
|
uvector<T, N> u = n;
|
|
u(j) -= 1.0;
|
|
u /= norm(u);
|
|
for (int dim = 0; dim < N; ++dim)
|
|
for (int i = 0; i < N; ++i) axes(dim)(i) = (dim == i ? 1.0 : 0.0) - 2.0 * u(dim) * u(i);
|
|
|
|
// Swap the first row of axes with j, so that the normal is the first axis. This is likely
|
|
// unnecessary (but done in the name of consistency with old approach).
|
|
std::swap(axes(0), axes(j));
|
|
|
|
// Apply coordinate transformation and calculate new bounding box in order to determine new split direction
|
|
uvector<T, N> bmin = std::numeric_limits<T>::max();
|
|
uvector<T, N> bmax = -std::numeric_limits<T>::max();
|
|
for (int i = lb; i < ub; ++i) {
|
|
uvector<T, N> x = points[index[i]];
|
|
for (int dim = 0; dim < N; ++dim) {
|
|
T alpha = dot(axes(dim), x);
|
|
points[index[i]](dim) = alpha;
|
|
if (alpha < bmin(dim)) bmin(dim) = alpha;
|
|
if (alpha > bmax(dim)) bmax(dim) = alpha;
|
|
}
|
|
}
|
|
|
|
node.type = static_cast<int>(transforms.size());
|
|
transforms.push_back(axes);
|
|
axis = argmax(bmax - bmin);
|
|
hasTransformed = true;
|
|
}
|
|
}
|
|
|
|
// Use median as the split
|
|
int m = (lb + ub) / 2;
|
|
|
|
// Rearrange points
|
|
std::nth_element(index.begin() + lb, index.begin() + m, index.begin() + ub, [&](int i0, int i1) {
|
|
return points[i0](axis) < points[i1](axis);
|
|
});
|
|
|
|
// Build child trees
|
|
int i0 = node.i0 = static_cast<int>(nodes.size());
|
|
int i1 = node.i1 = static_cast<int>(nodes.size() + 1);
|
|
nodes.push_back(Node());
|
|
nodes.push_back(Node());
|
|
build_tree(i0, lb, m, hasTransformed, level + 1);
|
|
build_tree(i1, m, ub, hasTransformed, level + 1);
|
|
}
|
|
|
|
struct ClosestPoint {
|
|
uvector<T, N> x;
|
|
T distsqr;
|
|
int ind;
|
|
};
|
|
|
|
// Recursive function for searching the tree for the closest point to a given point cp.x
|
|
void search(const Node& node, ClosestPoint& cp) const
|
|
{
|
|
if (node.type == -1) {
|
|
// Leaf node
|
|
for (int j = node.i0; j < node.i1; ++j) {
|
|
T dsqr = sqrnorm(points[j] - cp.x);
|
|
if (dsqr < cp.distsqr) {
|
|
cp.distsqr = dsqr;
|
|
cp.ind = j;
|
|
}
|
|
}
|
|
} else {
|
|
// Non-leaf node
|
|
if (node.type >= 0) {
|
|
// Transform query point to new coordinate system
|
|
const uvector<uvector<T, N>, N>& axes = transforms[node.type];
|
|
uvector<T, N> x = cp.x;
|
|
for (int dim = 0; dim < N; ++dim) cp.x(dim) = dot(axes(dim), x);
|
|
}
|
|
|
|
T dleft = detail::sqrDistBox(cp.x, nodes[node.i0].xmin, nodes[node.i0].xmax);
|
|
T dright = detail::sqrDistBox(cp.x, nodes[node.i1].xmin, nodes[node.i1].xmax);
|
|
if (dleft < dright) {
|
|
if (dleft < cp.distsqr) {
|
|
search(nodes[node.i0], cp);
|
|
if (dright < cp.distsqr) search(nodes[node.i1], cp);
|
|
}
|
|
} else {
|
|
if (dright < cp.distsqr) {
|
|
search(nodes[node.i1], cp);
|
|
if (dleft < cp.distsqr) search(nodes[node.i0], cp);
|
|
}
|
|
}
|
|
|
|
if (node.type >= 0) {
|
|
// Transform query point back to old coordinate system. This is about 5% faster than storing
|
|
// the old value of cp.x on the stack
|
|
const uvector<uvector<T, N>, N>& axes = transforms[node.type];
|
|
uvector<T, N> x = cp.x;
|
|
cp.x = axes(0) * x(0);
|
|
for (int dim = 1; dim < N; ++dim) cp.x += axes(dim) * x(dim);
|
|
}
|
|
}
|
|
}
|
|
|
|
public:
|
|
// Construct a KDTree from a given collection of points p (the given points are not overwritten)
|
|
KDTree(const std::vector<uvector<T, N>>& p)
|
|
{
|
|
assert(p.size()
|
|
< std::numeric_limits<int>::max()); // code currently uses int to index but could easily be adapted to size_t
|
|
int len = static_cast<int>(p.size());
|
|
if (len == 0) return;
|
|
|
|
// Copy points
|
|
points = p;
|
|
|
|
// Build initial index array, which shall soon be reordered
|
|
index.resize(len);
|
|
for (int i = 0; i < len; ++i) index[i] = i;
|
|
|
|
// Recursively build tree starting from root node. This only manipulates index array
|
|
nodes.push_back(Node());
|
|
build_tree(0, 0, len, false, 0);
|
|
|
|
// Rearrange so that points in leaf nodes are contiguous in memory. Based on tree
|
|
// construction, both nodes & point ranges will be organised depth-first.
|
|
std::vector<uvector<T, N>> pointscopy(points);
|
|
for (size_t i = 0; i < len; ++i) points[i] = pointscopy[index[i]];
|
|
}
|
|
|
|
// Search the tree for the closest point to x, where the search is restricted to all points that are
|
|
// within a squared distance of rsqr to x. Returns -1 if no such point exists, otherwise the index
|
|
// of the closest point in the original array p passed to the KDTree constructor is returned.
|
|
int nearest(const uvector<T, N>& x, T rsqr = std::numeric_limits<T>::max()) const
|
|
{
|
|
if (nodes.empty()) return -1;
|
|
ClosestPoint cp;
|
|
cp.x = x;
|
|
cp.distsqr = rsqr;
|
|
cp.ind = -1;
|
|
search(nodes[0], cp);
|
|
return cp.ind >= 0 ? index[cp.ind] : -1;
|
|
}
|
|
};
|
|
} // namespace algoim
|
|
|
|
#endif
|
|
|