filter: Add FilterSphere (#16)

* Add FilterSphere

* Use point converter instead of scalar converter

* Use std::clamp

* Remove unused ConverterNoOpScalar

Co-authored-by: Tilmann Zäschke <[email protected]>

Author: ctbur

README.md

@@ -444,7 +444,7 @@ There are numerous ways to improve performance. The following list gives an over
 4) **Use non-box query shapes**. Depending on the use case it may be more suitable to use a custom filter for querier. 
 For example: 
 
-    `tree.for_each(callback, MySphereFIlter(center, radius, tree.converter()));` 
+    `tree.for_each(callback, FilterSphere(center, radius, tree.converter()));`
 
 5) **Use a different data converter**. The default converter of the PH-Tree results in a reasonably fast index. 
 Its biggest advantage is that it provides lossless conversion from floating point coordinates to PH-Tree coordinates 

phtree/common/filter.h

@@ -23,6 +23,7 @@
 #include "flat_array_map.h"
 #include "flat_sparse_map.h"
 #include "tree_stats.h"
+#include <algorithm>
 #include <cassert>
 #include <cmath>
 #include <functional>
@@ -150,6 +151,73 @@ class FilterAABB {
     const CONVERTER converter_;
 };
 
+
+/*
+ * The sphere filter can be used to query a point tree for a sphere.
+ */
+template <
+    typename CONVERTER = ConverterIEEE<3>,
+    typename DISTANCE = DistanceEuclidean<3>>
+class FilterSphere {
+    using KeyExternal = typename CONVERTER::KeyExternal;
+    using KeyInternal = typename CONVERTER::KeyInternal;
+    using ScalarInternal = typename CONVERTER::ScalarInternal;
+    using ScalarExternal = typename CONVERTER::ScalarExternal;
+
+  public:
+    FilterSphere(
+        const KeyExternal& center,
+        const ScalarExternal& radius,
+        CONVERTER converter = CONVERTER(),
+        DISTANCE distance_function = DISTANCE())
+    : center_external_{center}
+    , center_internal_{converter.pre(center)}
+    , radius_{radius}
+    , converter_{converter}
+    , distance_function_{distance_function} {};
+
+    template <typename T>
+    [[nodiscard]] bool IsEntryValid(const KeyInternal& key, const T& value) const {
+        KeyExternal point = converter_.post(key);
+        return distance_function_(center_external_, point) <= radius_;
+    }
+
+    /*
+     * Calculate whether AABB encompassing all possible points in the node intersects with the
+     * sphere.
+     */
+    [[nodiscard]] bool IsNodeValid(const KeyInternal& prefix, int bits_to_ignore) const {
+        // we always want to traverse the root node (bits_to_ignore == 64)
+
+        if (bits_to_ignore >= (MAX_BIT_WIDTH<ScalarInternal> - 1)) {
+            return true;
+        }
+
+        ScalarInternal node_min_bits = MAX_MASK<ScalarInternal> << bits_to_ignore;
+        ScalarInternal node_max_bits = ~node_min_bits;
+
+        KeyInternal closest_in_bounds;
+        for (size_t i = 0; i < prefix.size(); ++i) {
+            // calculate lower and upper bound for dimension for given node
+            ScalarInternal lo = prefix[i] & node_min_bits;
+            ScalarInternal hi = prefix[i] | node_max_bits;
+
+            // choose value closest to center for dimension
+            closest_in_bounds[i] = std::clamp(center_internal_[i], lo, hi);
+        }
+
+        KeyExternal closest_point = converter_.post(closest_in_bounds);
+        return distance_function_(center_external_, closest_point) <= radius_;
+    }
+
+  private:
+    const KeyExternal center_external_;
+    const KeyExternal center_internal_;
+    const ScalarExternal radius_;
+    const CONVERTER converter_;
+    const DISTANCE distance_function_;
+};
+
 }  // namespace improbable::phtree
 
 #endif  // PHTREE_COMMON_FILTERS_H

phtree/common/filter_test.cc

@@ -20,6 +20,30 @@
 
 using namespace improbable::phtree;
 
+TEST(PhTreeFilterTest, FilterSphereTest) {
+    FilterSphere<ConverterNoOp<2, scalar_64_t>, DistanceEuclidean<2>> filter{{5, 3}, 5};
+    // root is always valid
+    ASSERT_TRUE(filter.IsNodeValid({0, 0}, 63));
+    // valid because node encompasses the circle
+    ASSERT_TRUE(filter.IsNodeValid({1, 1}, 10));
+    // valid because circle encompasses the node
+    ASSERT_TRUE(filter.IsNodeValid({5, 5}, 2));
+    // valid because circle encompasses the node AABB
+    ASSERT_TRUE(filter.IsNodeValid({7, 7}, 1));
+    // valid because circle touches the edge of the node AABB
+    ASSERT_TRUE(filter.IsNodeValid({5, 9}, 1));
+    // valid because circle cuts edge of node AABB
+    ASSERT_TRUE(filter.IsNodeValid({12, 7}, 3));
+    ASSERT_TRUE(filter.IsNodeValid({10, 7}, 2));
+    // invalid because node is just outside the circle
+    ASSERT_FALSE(filter.IsNodeValid({5, 10}, 1));
+    ASSERT_FALSE(filter.IsNodeValid({12, 12}, 3));
+
+    ASSERT_TRUE(filter.IsEntryValid({3, 7}, nullptr));
+    ASSERT_TRUE(filter.IsEntryValid({5, 8}, nullptr));
+    ASSERT_FALSE(filter.IsEntryValid({3, 8}, nullptr));
+}
+
 TEST(PhTreeFilterTest, BoxFilterTest) {
     FilterAABB<ConverterNoOp<2, scalar_64_t>> filter{{3, 3}, {7, 7}};
     // root is always valid