Adding KdTree and PointKdTreeSearcher (#115)

This revision includes:
* Generic KdTree implementation
* PointKdTreeSearcher based on KdTree
* Performance improvements on Anisotropic kernel
* Updating flatbuffers

Author: doyubkim

CMakeLists.txt

@@ -62,7 +62,7 @@ add_subdirectory(external/googletest EXCLUDE_FROM_ALL)
 if ((CMAKE_VERSION VERSION_EQUAL 3.3) OR (CMAKE_VERSION VERSION_GREATER 3.3))
     cmake_policy(SET CMP0063 NEW)
 endif()
-set(BENCHMARK_ENABLE_TESTING 0)
+set(BENCHMARK_ENABLE_TESTING OFF)
 add_subdirectory(external/googlebenchmark)
 set(CMAKE_CXX_VISIBILITY_PRESET hidden)
 set(CMAKE_VISIBILITY_INLINES_HIDDEN 1)

README.md

@@ -15,6 +15,7 @@ The latest code is always available from the [`master`](https://github.com/doyub
 * PIC, FLIP, and APIC fluid simulators
 * Upwind, ENO, and FMM level set solvers
 * Jacobi, Gauss-Seidel, SOR, MG, CG, ICCG, and MGPCG linear system solvers
+* Spherical, SPH, Zhu & Bridson, and Anisotropic kernel for points-to-surface converter 
 * Converters between signed distance function and triangular mesh
 * C++ and Python API
 

include/jet/anisotropic_points_to_implicit2.h

@@ -36,7 +36,7 @@ class AnisotropicPointsToImplicit2 final : public PointsToImplicit2 {
     //!
     AnisotropicPointsToImplicit2(double kernelRadius = 1.0,
                                  double cutOffDensity = 0.5,
-                                 double positionSmoothingFactor = 0.0,
+                                 double positionSmoothingFactor = 0.5,
                                  size_t minNumNeighbors = 8,
                                  bool isOutputSdf = true);
 

include/jet/anisotropic_points_to_implicit3.h

@@ -36,7 +36,7 @@ class AnisotropicPointsToImplicit3 final : public PointsToImplicit3 {
     //!
     AnisotropicPointsToImplicit3(double kernelRadius = 1.0,
                                  double cutOffDensity = 0.5,
-                                 double positionSmoothingFactor = 0.0,
+                                 double positionSmoothingFactor = 0.5,
                                  size_t minNumNeighbors = 25,
                                  bool isOutputSdf = true);
 

include/jet/detail/bvh2-inl.h

@@ -22,7 +22,7 @@ Bvh2<T>::Node::Node() : flags(0) {
 
 template <typename T>
 void Bvh2<T>::Node::initLeaf(size_t it, const BoundingBox2D& b) {
-    flags |= 2;
+    flags = 2;
     item = it;
     bound = b;
 }

include/jet/detail/bvh3-inl.h

@@ -22,7 +22,7 @@ Bvh3<T>::Node::Node() : flags(0) {
 
 template <typename T>
 void Bvh3<T>::Node::initLeaf(size_t it, const BoundingBox3D& b) {
-    flags |= 3;
+    flags = 3;
     item = it;
     bound = b;
 }

include/jet/detail/kdtree-inl.h

@@ -0,0 +1,309 @@
+// Copyright (c) 2017 Doyub Kim
+//
+// I am making my contributions/submissions to this project solely in my
+// personal capacity and am not conveying any rights to any intellectual
+// property of any third parties.
+
+#ifndef INCLUDE_JET_DETAIL_KDTREE_INL_H_
+#define INCLUDE_JET_DETAIL_KDTREE_INL_H_
+
+#include <jet/kdtree.h>
+
+#include <numeric>
+
+namespace jet {
+
+template <typename T, size_t K>
+KdTree<T, K>::Node::Node() {
+    child = kMaxSize;
+}
+
+template <typename T, size_t K>
+void KdTree<T, K>::Node::initLeaf(size_t it, const Point& pt) {
+    flags = K;
+    item = it;
+    child = kMaxSize;
+    point = pt;
+}
+
+template <typename T, size_t K>
+void KdTree<T, K>::Node::initInternal(size_t axis, size_t it, size_t c,
+                                      const Point& pt) {
+    flags = axis;
+    item = it;
+    child = c;
+    point = pt;
+}
+
+template <typename T, size_t K>
+bool KdTree<T, K>::Node::isLeaf() const {
+    return flags == K;
+}
+
+//
+
+template <typename T, size_t K>
+KdTree<T, K>::KdTree() {}
+
+template <typename T, size_t K>
+void KdTree<T, K>::build(const ConstArrayAccessor1<Point>& points) {
+    _points.resize(points.size());
+    std::copy(points.begin(), points.end(), _points.begin());
+
+    if (_points.empty()) {
+        return;
+    }
+
+    _nodes.clear();
+
+    std::vector<size_t> itemIndices(_points.size());
+    std::iota(std::begin(itemIndices), std::end(itemIndices), 0);
+
+    build(0, itemIndices.data(), _points.size(), 0);
+}
+
+template <typename T, size_t K>
+void KdTree<T, K>::forEachNearbyPoint(
+    const Point& origin, T radius,
+    const std::function<void(size_t, const Point&)>& callback) const {
+    const T r2 = radius * radius;
+
+    // prepare to traverse the tree for sphere
+    static const int kMaxTreeDepth = 8 * sizeof(size_t);
+    const Node* todo[kMaxTreeDepth];
+    size_t todoPos = 0;
+
+    // traverse the tree nodes for sphere
+    const Node* node = _nodes.data();
+
+    while (node != nullptr) {
+        if (node->item != kMaxSize &&
+            (node->point - origin).lengthSquared() <= r2) {
+            callback(node->item, node->point);
+        }
+
+        if (node->isLeaf()) {
+            // grab next node to process from todo stack
+            if (todoPos > 0) {
+                // Dequeue
+                --todoPos;
+                node = todo[todoPos];
+            } else {
+                break;
+            }
+        } else {
+            // get node children pointers for sphere
+            const Node* firstChild = node + 1;
+            const Node* secondChild = (Node*)&_nodes[node->child];
+
+            // advance to next child node, possibly enqueue other child
+            const size_t axis = node->flags;
+            const T plane = node->point[axis];
+            if (plane - origin[axis] > radius) {
+                node = firstChild;
+            } else if (origin[axis] - plane > radius) {
+                node = secondChild;
+            } else {
+                // enqueue secondChild in todo stack
+                todo[todoPos] = secondChild;
+                ++todoPos;
+                node = firstChild;
+            }
+        }
+    }
+}
+
+template <typename T, size_t K>
+bool KdTree<T, K>::hasNearbyPoint(const Point& origin, T radius) const {
+    const T r2 = radius * radius;
+
+    // prepare to traverse the tree for sphere
+    static const int kMaxTreeDepth = 8 * sizeof(size_t);
+    const Node* todo[kMaxTreeDepth];
+    size_t todoPos = 0;
+
+    // traverse the tree nodes for sphere
+    const Node* node = _nodes.data();
+
+    while (node != nullptr) {
+        if (node->item != kMaxSize &&
+            (node->point - origin).lengthSquared() <= r2) {
+            return true;
+        }
+
+        if (node->isLeaf()) {
+            // grab next node to process from todo stack
+            if (todoPos > 0) {
+                // Dequeue
+                --todoPos;
+                node = todo[todoPos];
+            } else {
+                break;
+            }
+        } else {
+            // get node children pointers for sphere
+            const Node* firstChild = node + 1;
+            const Node* secondChild = (Node*)&_nodes[node->child];
+
+            // advance to next child node, possibly enqueue other child
+            const size_t axis = node->flags;
+            const T plane = node->point[axis];
+            if (origin[axis] < plane && plane - origin[axis] > radius) {
+                node = firstChild;
+            } else if (origin[axis] > plane && origin[axis] - plane > radius) {
+                node = secondChild;
+            } else {
+                // enqueue secondChild in todo stack
+                todo[todoPos] = secondChild;
+                ++todoPos;
+                node = firstChild;
+            }
+        }
+    }
+
+    return false;
+}
+
+template <typename T, size_t K>
+size_t KdTree<T, K>::nearestPoint(const Point& origin) const {
+    // prepare to traverse the tree for sphere
+    static const int kMaxTreeDepth = 8 * sizeof(size_t);
+    const Node* todo[kMaxTreeDepth];
+    size_t todoPos = 0;
+
+    // traverse the tree nodes for sphere
+    const Node* node = _nodes.data();
+    size_t nearest = 0;
+    T minDist2 = (node->point - origin).lengthSquared();
+
+    while (node != nullptr) {
+        const T newDist2 = (node->point - origin).lengthSquared();
+        if (newDist2 <= minDist2) {
+            nearest = node->item;
+            minDist2 = newDist2;
+        }
+
+        if (node->isLeaf()) {
+            // grab next node to process from todo stack
+            if (todoPos > 0) {
+                // Dequeue
+                --todoPos;
+                node = todo[todoPos];
+            } else {
+                break;
+            }
+        } else {
+            // get node children pointers for sphere
+            const Node* firstChild = node + 1;
+            const Node* secondChild = (Node*)&_nodes[node->child];
+
+            // advance to next child node, possibly enqueue other child
+            const size_t axis = node->flags;
+            const T plane = node->point[axis];
+            const T minDist = std::sqrt(minDist2);
+            if (plane - origin[axis] > minDist) {
+                node = firstChild;
+            } else if (origin[axis] - plane > minDist) {
+                node = secondChild;
+            } else {
+                // enqueue secondChild in todo stack
+                todo[todoPos] = secondChild;
+                ++todoPos;
+                node = firstChild;
+            }
+        }
+    }
+
+    return nearest;
+}
+
+template <typename T, size_t K>
+void KdTree<T, K>::reserve(size_t numPoints, size_t numNodes) {
+    _points.resize(numPoints);
+    _nodes.resize(numNodes);
+}
+
+template <typename T, size_t K>
+typename KdTree<T, K>::Iterator KdTree<T, K>::begin() {
+    return _points.begin();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::Iterator KdTree<T, K>::end() {
+    return _points.end();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::ConstIterator KdTree<T, K>::begin() const {
+    return _points.begin();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::ConstIterator KdTree<T, K>::end() const {
+    return _points.end();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::NodeIterator KdTree<T, K>::beginNode() {
+    return _nodes.begin();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::NodeIterator KdTree<T, K>::endNode() {
+    return _nodes.end();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::ConstNodeIterator KdTree<T, K>::beginNode() const {
+    return _nodes.begin();
+};
+
+template <typename T, size_t K>
+typename KdTree<T, K>::ConstNodeIterator KdTree<T, K>::endNode() const {
+    return _nodes.end();
+};
+
+template <typename T, size_t K>
+size_t KdTree<T, K>::build(size_t nodeIndex, size_t* itemIndices, size_t nItems,
+                           size_t currentDepth) {
+    // add a node
+    _nodes.emplace_back();
+
+    // initialize leaf node if termination criteria met
+    if (nItems == 0) {
+        _nodes[nodeIndex].initLeaf(kMaxSize, {});
+        return currentDepth + 1;
+    }
+    if (nItems == 1) {
+        _nodes[nodeIndex].initLeaf(itemIndices[0], _points[itemIndices[0]]);
+        return currentDepth + 1;
+    }
+
+    // choose which axis to split along
+    BBox nodeBound;
+    for (size_t i = 0; i < nItems; ++i) {
+        nodeBound.merge(_points[itemIndices[i]]);
+    }
+    Point d = nodeBound.upperCorner - nodeBound.lowerCorner;
+    size_t axis = static_cast<size_t>(d.dominantAxis());
+
+    // pick mid point
+    std::nth_element(itemIndices, itemIndices + nItems / 2,
+                     itemIndices + nItems, [&](size_t a, size_t b) {
+                         return _points[a][axis] < _points[b][axis];
+                     });
+    size_t midPoint = nItems / 2;
+
+    // recursively initialize children nodes
+    size_t d0 = build(nodeIndex + 1, itemIndices, midPoint, currentDepth + 1);
+    _nodes[nodeIndex].initInternal(axis, itemIndices[midPoint], _nodes.size(),
+                                   _points[itemIndices[midPoint]]);
+    size_t d1 = build(_nodes[nodeIndex].child, itemIndices + midPoint + 1,
+                      nItems - midPoint - 1, currentDepth + 1);
+
+    return std::max(d0, d1);
+}
+
+}  // namespace jet
+
+#endif  // INCLUDE_JET_DETAIL_KDTREE_INL_H_

include/jet/detail/serialization-inl.h

@@ -14,7 +14,7 @@
 namespace jet {
 
 template <typename T>
-void serialize(const Array1<T>& array, std::vector<uint8_t>* buffer) {
+void serialize(const ConstArrayAccessor1<T>& array, std::vector<uint8_t>* buffer) {
     size_t size = sizeof(T) * array.size();
     serialize(reinterpret_cast<const uint8_t*>(array.data()), size, buffer);
 }

include/jet/jet.h

@@ -134,6 +134,7 @@
 #include <jet/iterative_level_set_solver2.h>
 #include <jet/iterative_level_set_solver3.h>
 #include <jet/jet.h>
+#include <jet/kdtree.h>
 #include <jet/level_set_liquid_solver2.h>
 #include <jet/level_set_liquid_solver3.h>
 #include <jet/level_set_solver2.h>
@@ -180,6 +181,8 @@
 #include <jet/point_generator3.h>
 #include <jet/point_hash_grid_searcher2.h>
 #include <jet/point_hash_grid_searcher3.h>
+#include <jet/point_kdtree_searcher2.h>
+#include <jet/point_kdtree_searcher3.h>
 #include <jet/point_neighbor_searcher2.h>
 #include <jet/point_neighbor_searcher3.h>
 #include <jet/point_parallel_hash_grid_searcher2.h>