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.

264 lines
10 KiB

12 months ago
#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
12 months ago
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.
12 months ago
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
12 months ago
template <typename T, int N>
class KDTree
{
std::vector<uvector<T, N>> points;
std::vector<int> index;
static constexpr int leaf_size = 16;
12 months ago
// 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;
};
12 months ago
std::vector<Node> nodes;
std::vector<uvector<uvector<T, N>, N>> transforms;
12 months ago
// 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];
12 months ago
// 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);
12 months ago
}
}
mean /= static_cast<T>(ub - lb);
12 months ago
// 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;
}
12 months ago
// Splitting node: default to splitting along greatest extent
node.type = -2;
int axis = argmax(node.xmax - node.xmin);
12 months ago
// 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)
12 months ago
n(0) = 1.0;
else
n /= sqrt(msqr);
12 months ago
// 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;
}
12 months ago
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);
12 months ago
// 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));
12 months ago
// 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;
12 months ago
}
}
node.type = static_cast<int>(transforms.size());
transforms.push_back(axes);
axis = argmax(bmax - bmin);
hasTransformed = true;
12 months ago
}
}
12 months ago
// Use median as the split
int m = (lb + ub) / 2;
12 months ago
// 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);
});
12 months ago
// 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);
}
12 months ago
struct ClosestPoint {
uvector<T, N> x;
T distsqr;
int ind;
};
12 months ago
// 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;
12 months ago
}
}
} 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);
}
12 months ago
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);
12 months ago
}
} else {
if (dright < cp.distsqr) {
search(nodes[node.i1], cp);
if (dleft < cp.distsqr) search(nodes[node.i0], cp);
12 months ago
}
}
12 months ago
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);
12 months ago
}
}
}
12 months ago
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;
12 months ago
// Copy points
points = p;
12 months ago
// Build initial index array, which shall soon be reordered
index.resize(len);
for (int i = 0; i < len; ++i) index[i] = i;
12 months ago
// Recursively build tree starting from root node. This only manipulates index array
nodes.push_back(Node());
build_tree(0, 0, len, false, 0);
12 months ago
// 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]];
}
12 months ago
// 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;
}
};
12 months ago
} // namespace algoim
#endif