Extend kdtree with k-nearest and bounding box queries

Also add test to verify it

src/libslic3r/CMakeLists.txt

@@ -212,6 +212,7 @@ set(SLIC3R_SOURCES
     PrintObject.cpp
     PrintObjectSlice.cpp
     PrintRegion.cpp
+    PointGrid.hpp
     PNGReadWrite.hpp
     PNGReadWrite.cpp
     QuadricEdgeCollapse.cpp

src/libslic3r/KDTreeIndirect.hpp

@@ -11,6 +11,12 @@
 
 namespace Slic3r {
 
+enum class VisitorReturnMask : unsigned int {
+    CONTINUE_LEFT  = 1,
+    CONTINUE_RIGHT = 2,
+    STOP           = 4,
+};
+
 // KD tree for N-dimensional closest point search.
 template<size_t ANumDimensions, typename ACoordType, typename ACoordinateFn>
 class KDTreeIndirect
@@ -25,8 +31,7 @@ public:
     };
 
     KDTreeIndirect(CoordinateFn coordinate) : coordinate(coordinate) {}
-	KDTreeIndirect(CoordinateFn coordinate, std::vector<size_t>   indices) : coordinate(coordinate) { this->build(std::move(indices)); }
+    KDTreeIndirect(CoordinateFn coordinate, std::vector<size_t> indices) : coordinate(coordinate) { this->build(indices); }
-	KDTreeIndirect(CoordinateFn coordinate, std::vector<size_t> &&indices) : coordinate(coordinate) { this->build(std::move(indices)); }
     KDTreeIndirect(CoordinateFn coordinate, size_t num_indices) : coordinate(coordinate) { this->build(num_indices); }
     KDTreeIndirect(KDTreeIndirect &&rhs) : m_nodes(std::move(rhs.m_nodes)), coordinate(std::move(rhs.coordinate)) {}
     KDTreeIndirect& operator=(KDTreeIndirect &&rhs) { m_nodes = std::move(rhs.m_nodes); coordinate = std::move(rhs.coordinate); return *this; }
@@ -38,10 +43,10 @@ public:
         indices.reserve(num_indices);
         for (size_t i = 0; i < num_indices; ++ i)
             indices.emplace_back(i);
-		this->build(std::move(indices));
+        this->build(indices);
     }
 
-	void build(std::vector<size_t> &&indices)
+    void build(std::vector<size_t> &indices)
     {
         if (indices.empty())
             clear();
@@ -53,12 +58,6 @@ public:
         indices.clear();
     }
 
-	enum class VisitorReturnMask : unsigned int
-	{
-		CONTINUE_LEFT   = 1,
-		CONTINUE_RIGHT  = 2,
-		STOP 			= 4,
-	};
     template<typename CoordType>
     unsigned int descent_mask(const CoordType &point_coord, const CoordType &search_radius, size_t idx, size_t dimension) const
     {
@@ -192,37 +191,83 @@ private:
 
 // Find a closest point using Euclidian metrics.
 // Returns npos if not found.
-template<typename KDTreeIndirectType, typename PointType, typename FilterFn>
+template<size_t K,
-size_t find_closest_point(const KDTreeIndirectType &kdtree, const PointType &point, FilterFn filter)
+         typename PointType,
+         typename FilterFn,
+         size_t D,
+         typename CoordT,
+         typename CoordFn>
+std::array<size_t, K> find_closest_points(
+    const KDTreeIndirect<D, CoordT, CoordFn> &kdtree,
+    const PointType                          &point,
+    FilterFn                                  filter)
 {
-	using CoordType = typename KDTreeIndirectType::CoordType;
+    using Tree = KDTreeIndirect<D, CoordT, CoordFn>;
 
-	struct Visitor {
+    struct Visitor
-		const KDTreeIndirectType   &kdtree;
+    {
+        const Tree      &kdtree;
         const PointType &point;
         const FilterFn   filter;
-		size_t 						min_idx  = KDTreeIndirectType::npos;
-		CoordType					min_dist = std::numeric_limits<CoordType>::max();
 
-		Visitor(const KDTreeIndirectType &kdtree, const PointType &point, FilterFn filter) : kdtree(kdtree), point(point), filter(filter) {}
+        std::array<std::pair<size_t, CoordT>, K> results;
-		unsigned int operator()(size_t idx, size_t dimension) {
+
+        Visitor(const Tree &kdtree, const PointType &point, FilterFn filter)
+            : kdtree(kdtree), point(point), filter(filter)
+        {
+            results.fill(std::make_pair(Tree::npos,
+                                        std::numeric_limits<CoordT>::max()));
+        }
+        unsigned int operator()(size_t idx, size_t dimension)
+        {
             if (this->filter(idx)) {
-				auto dist = CoordType(0);
+                auto dist = CoordT(0);
-				for (size_t i = 0; i < KDTreeIndirectType::NumDimensions; ++ i) {
+                for (size_t i = 0; i < D; ++i) {
-					CoordType d = point[i] - kdtree.coordinate(idx, i);
+                    CoordT d = point[i] - kdtree.coordinate(idx, i);
                     dist += d * d;
                 }
-				if (dist < min_dist) {
+
-					min_dist = dist;
+                auto res = std::make_pair(idx, dist);
-					min_idx  = idx;
+                auto it  = std::lower_bound(results.begin(), results.end(),
+                                            res, [](auto &r1, auto &r2) {
+                                               return r1.second < r2.second;
+                                            });
+
+                if (it != results.end()) {
+                    std::rotate(it, std::prev(results.end()), results.end());
+                    *it = res;
                 }
             }
-			return kdtree.descent_mask(point[dimension], min_dist, idx, dimension);
+            return kdtree.descent_mask(point[dimension],
+                                       results.front().second, idx,
+                                       dimension);
         }
     } visitor(kdtree, point, filter);
 
     kdtree.visit(visitor);
-	return visitor.min_idx;
+    std::array<size_t, K> ret;
+    for (size_t i = 0; i < K; i++) ret[i] = visitor.results[i].first;
+
+    return ret;
+}
+
+template<size_t K, typename PointType, size_t D, typename CoordT, typename CoordFn>
+std::array<size_t, K> find_closest_points(
+    const KDTreeIndirect<D, CoordT, CoordFn> &kdtree, const PointType &point)
+{
+    return find_closest_points<K>(kdtree, point, [](size_t) { return true; });
+}
+
+template<typename PointType,
+         typename FilterFn,
+         size_t D,
+         typename CoordT,
+         typename CoordFn>
+size_t find_closest_point(const KDTreeIndirect<D, CoordT, CoordFn> &kdtree,
+                          const PointType                          &point,
+                          FilterFn                                  filter)
+{
+    return find_closest_points<1>(kdtree, point, filter)[0];
 }
 
 template<typename KDTreeIndirectType, typename PointType>
@@ -235,7 +280,7 @@ size_t find_closest_point(const KDTreeIndirectType& kdtree, const PointType& poi
 template<typename KDTreeIndirectType, typename PointType, typename FilterFn>
 std::vector<size_t> find_nearby_points(const KDTreeIndirectType &kdtree, const PointType &center,
                                        const typename KDTreeIndirectType::CoordType& max_distance, FilterFn filter)
-        {
+{
     using CoordType = typename KDTreeIndirectType::CoordType;
 
     struct Visitor {
@@ -271,12 +316,58 @@ std::vector<size_t> find_nearby_points(const KDTreeIndirectType &kdtree, const P
 template<typename KDTreeIndirectType, typename PointType>
 std::vector<size_t> find_nearby_points(const KDTreeIndirectType &kdtree, const PointType &center,
                                        const typename KDTreeIndirectType::CoordType& max_distance)
-        {
+{
     return find_nearby_points(kdtree, center, max_distance, [](size_t) {
         return true;
     });
 }
 
+// Find nearby points (spherical neighbourhood) using Euclidian metrics.
+template<typename KDTreeIndirectType, typename PointType, typename FilterFn>
+std::vector<size_t> find_nearby_points(const KDTreeIndirectType &kdtree,
+                                       const PointType          &bb_min,
+                                       const PointType          &bb_max,
+                                       FilterFn                  filter)
+{
+    struct Visitor {
+        const KDTreeIndirectType &kdtree;
+        const PointType &bb_min, &bb_max;
+        const FilterFn filter;
+        std::vector<size_t> result;
+
+        Visitor(const KDTreeIndirectType &kdtree, const PointType& bbmin, const PointType& bbmax,
+                FilterFn filter) :
+            kdtree(kdtree), bb_min{bbmin}, bb_max{bbmax}, filter(filter) {
+        }
+        unsigned int operator()(size_t idx, size_t dimension) {
+            unsigned int ret =
+                static_cast<unsigned int>(VisitorReturnMask::CONTINUE_LEFT) |
+                static_cast<unsigned int>(VisitorReturnMask::CONTINUE_RIGHT);
+
+            if (this->filter(idx)) {
+                PointType p;
+                bool contains = true;
+                for (size_t i = 0; i < KDTreeIndirectType::NumDimensions; ++i) {
+                    p(i) = kdtree.coordinate(idx, i);
+                    contains = contains && bb_min(i) <= p(i) && p(i) <= bb_max(i);
+                }
+
+                if (p(dimension) < bb_min(dimension))
+                    ret = static_cast<unsigned int>(VisitorReturnMask::CONTINUE_RIGHT);
+                if (p(dimension) > bb_max(dimension))
+                    ret = static_cast<unsigned int>(VisitorReturnMask::CONTINUE_LEFT);
+
+                if (contains)
+                    result.emplace_back(idx);
+            }
+
+            return ret;
+        }
+    } visitor(kdtree, bb_min, bb_max, filter);
+
+    kdtree.visit(visitor);
+    return visitor.result;
+}
 
 } // namespace Slic3r
 

src/libslic3r/PointGrid.hpp

@@ -0,0 +1,74 @@
+#ifndef POINTGRID_HPP
+#define POINTGRID_HPP
+
+#include <libslic3r/Execution/Execution.hpp>
+#include <libslic3r/Point.hpp>
+#include <libslic3r/BoundingBox.hpp>
+
+namespace Slic3r {
+
+template<class T>
+class PointGrid {
+    Vec3i m_size;
+    std::vector<Vec<3, T>> m_data;
+    const int XY;
+
+public:
+    explicit PointGrid(std::vector<Vec<3, T>> data, const Vec3i &size)
+        : m_data(std::move(data)), m_size{size}, XY{m_size.x() * m_size.y()}
+    {}
+
+    const Vec<3, T> & get(size_t idx) const { return m_data[idx]; }
+    const Vec<3, T> & get(const Vec3i &coord) const
+    {
+        return m_data[get_idx(coord)];
+    }
+
+    size_t get_idx(const Vec3i &coord) const
+    {
+        size_t ret = coord.z() * XY + coord.y() * m_size.x() + coord.x();
+
+        return ret;
+    }
+
+    Vec3i get_coord(size_t idx) const {
+        size_t iz = idx / XY;
+        size_t iy = (idx / m_size.x()) % m_size.y();
+        size_t ix = idx % m_size.x();
+
+        return {ix, iy, iz};
+    }
+
+    const std::vector<Vec<3, T>> & data() const { return m_data; }
+    size_t point_count() const { return m_data.size(); }
+    bool empty() const { return m_data.empty(); }
+};
+
+template<class Ex, class CoordT>
+PointGrid<CoordT> point_grid(Ex                                      policy,
+                             const BoundingBox3Base<Vec<3, CoordT>> &bounds,
+                             const Vec<3, CoordT>                   &stride)
+{
+    Vec3i numpts = Vec3i::Zero();
+
+    for (int n = 0; n < 3; ++n)
+        numpts(n) = (bounds.max(n) - bounds.min(n)) / stride(n);
+
+    std::vector<Vec<3, CoordT>> out(numpts.x() * numpts.y() * numpts.z());
+
+    size_t XY = numpts[X] * numpts[Y];
+
+    execution::for_each(policy, size_t(0), out.size(), [&](size_t i) {
+        size_t iz = i / XY;
+        size_t iy = (i / numpts[X]) % numpts[Y];
+        size_t ix = i % numpts[X];
+
+        out[i] = Vec<3, CoordT>(ix * stride.x(), iy * stride.y(), iz * stride.z());
+    });
+
+    return PointGrid{std::move(out), numpts};
+}
+
+} // namespace Slic3r
+
+#endif // POINTGRID_HPP

tests/libslic3r/CMakeLists.txt

@@ -4,6 +4,7 @@ add_executable(${_TEST_NAME}_tests
 	${_TEST_NAME}_tests.cpp
 	test_3mf.cpp
 	test_aabbindirect.cpp
+	test_kdtreeindirect.cpp
 	test_clipper_offset.cpp
 	test_clipper_utils.cpp
 	test_color.cpp

tests/libslic3r/test_kdtreeindirect.cpp

@@ -0,0 +1,142 @@
+#include <catch2/catch.hpp>
+
+#include "libslic3r/KDTreeIndirect.hpp"
+#include "libslic3r/Execution/ExecutionSeq.hpp"
+#include "libslic3r/BoundingBox.hpp"
+#include "libslic3r/PointGrid.hpp"
+
+using namespace Slic3r;
+
+//template<class G>
+//struct Within { // Wrapper for the `within` predicate that counts calls.
+
+//    kdtree::Within<G> pred;
+
+//    Within(G box): pred{box} {}
+
+//       // Number of times the predicate was called
+//    mutable size_t call_count = 0;
+
+//    std::pair<bool, unsigned int> operator() (const Vec3f &p, size_t dim)
+//    {
+//        ++call_count;
+
+//        return pred(p, dim);
+//    }
+//};
+
+static double volume(const BoundingBox3Base<Vec3f> &box)
+{
+    auto sz = box.size();
+    return sz.x() * sz.y() * sz.z();
+}
+
+static double volume(const Eigen::AlignedBox<float, 3> &box)
+{
+    return box.volume();
+}
+
+TEST_CASE("Test kdtree query for a Box", "[KDTreeIndirect]")
+{
+    auto vol = BoundingBox3Base<Vec3f>{{0.f, 0.f, 0.f}, {10.f, 10.f, 10.f}};
+
+    auto pgrid = point_grid(ex_seq, vol, Vec3f{0.1f, 0.1f, 0.1f});
+
+    REQUIRE(!pgrid.empty());
+
+    auto coordfn = [&pgrid] (size_t i, size_t D) { return pgrid.get(i)(int(D)); };
+    KDTreeIndirect<3, float, decltype(coordfn)> tree{coordfn, pgrid.point_count()};
+
+    std::vector<size_t> out;
+
+    auto qbox = BoundingBox3Base{Vec3f{0.f, 0.f, 0.f}, Vec3f{.5f, .5f, .5f}};
+
+    size_t call_count = 0;
+    out = find_nearby_points(tree, qbox.min, qbox.max, [&call_count](size_t) {
+        call_count++;
+        return true;
+    });
+
+    // Output shall be non-empty
+    REQUIRE(!out.empty());
+
+    std::sort(out.begin(), out.end());
+
+    // No duplicates allowed in the output
+    auto it = std::unique(out.begin(), out.end());
+    REQUIRE(it == out.end());
+
+    // Test if inside points are in the output and outside points are not.
+    bool succ = true;
+    for (size_t i = 0; i < pgrid.point_count(); ++i) {
+        auto foundit = std::find(out.begin(), out.end(), i);
+        bool contains = qbox.contains(pgrid.get(i));
+        succ = succ && contains ? foundit != out.end() : foundit == out.end();
+
+        if (!succ) {
+            std::cout << "invalid point: " << i << " " << pgrid.get(i).transpose()
+                      << std::endl;
+            break;
+        }
+    }
+
+    REQUIRE(succ);
+
+    // Test for the expected cost of the query.
+    double gridvolume = volume(vol);
+    double queryvolume = volume(qbox);
+    double volratio = (queryvolume / gridvolume);
+    REQUIRE(call_count < 3 * volratio * pgrid.point_count());
+    REQUIRE(call_count < pgrid.point_count());
+}
+
+//TEST_CASE("Test kdtree query for a Sphere", "[KDTreeIndirect]") {
+//    auto vol = BoundingBox3Base<Vec3f>{{0.f, 0.f, 0.f}, {10.f, 10.f, 10.f}};
+
+//    auto pgrid = point_grid(ex_seq, vol, Vec3f{0.1f, 0.1f, 0.1f});
+
+//    REQUIRE(!pgrid.empty());
+
+//    auto coordfn = [&pgrid] (size_t i, size_t D) { return pgrid.get(i)(int(D)); };
+//    kdtree::KDTreeIndirect<3, float, decltype(coordfn)> tree{coordfn, pgrid.point_count()};
+
+//    std::vector<size_t> out;
+
+//    auto querysphere = kdtree::Sphere{Vec3f{5.f, 5.f, 5.f}, 2.f};
+
+//    auto pred = Within(querysphere);
+
+//    kdtree::query(tree, pred, std::back_inserter(out));
+
+//    // Output shall be non-empty
+//    REQUIRE(!out.empty());
+
+//    std::sort(out.begin(), out.end());
+
+//    // No duplicates allowed in the output
+//    auto it = std::unique(out.begin(), out.end());
+//    REQUIRE(it == out.end());
+
+//    // Test if inside points are in the output and outside points are not.
+//    bool succ = true;
+//    for (size_t i = 0; i < pgrid.point_count(); ++i) {
+//        auto foundit = std::find(out.begin(), out.end(), i);
+//        bool contains = (querysphere.center - pgrid.get(i)).squaredNorm() < pred.pred.r2;
+//        succ = succ && contains ? foundit != out.end() : foundit == out.end();
+
+//        if (!succ) {
+//            std::cout << "invalid point: " << i << " " << pgrid.get(i).transpose()
+//                      << std::endl;
+//            break;
+//        }
+//    }
+
+//    REQUIRE(succ);
+
+//    // Test for the expected cost of the query.
+//    double gridvolume = volume(vol);
+//    double queryvolume = volume(querysphere);
+//    double volratio = (queryvolume / gridvolume);
+//    REQUIRE(pred.call_count < 3 * volratio * pgrid.point_count());
+//    REQUIRE(pred.call_count < pgrid.point_count());
+//}