filter: Add FilterParaboloid

phtree/common/filter.h

@@ -215,6 +215,140 @@ class FilterSphere {
     const DISTANCE distance_function_;
 };
 
+/*
+ * The paraboloid filter can be used to query a point tree for an axis-aligned paraboloid.
+ * This is useful for, e.g., doing a range check for a projectile with a fixed launch speed.
+ * Also see https://en.wikipedia.org/wiki/paraboloid_of_safety for this application.
+ */
+template <
+    typename CONVERTER = ConverterIEEE<3>,
+    typename DISTANCE = DistanceEuclidean<3>>
+class FilterParaboloid {
+    using KeyExternal = typename CONVERTER::KeyExternal;
+    using KeyInternal = typename CONVERTER::KeyInternal;
+    using ScalarInternal = typename CONVERTER::ScalarInternal;
+    using ScalarExternal = typename CONVERTER::ScalarExternal;
+
+    static constexpr auto DIM = CONVERTER::DimInternal;
+
+  public:
+    /*
+     * @param axis The axis along which the paraboloid is aligned.
+     * @param vertex The vertex of the paraboloid.
+     * @param scale_factor Scales the paraboloid along the axis. A positive scale factor
+     * means that the paraboloid opens in the positive direction of the axis, a negative
+     * scale factor means the reverse.
+     */
+    FilterParaboloid(
+        const size_t axis,
+        const KeyExternal& vertex,
+        const ScalarExternal& scale_factor,
+        CONVERTER converter = CONVERTER(),
+        DISTANCE distance_function = DISTANCE())
+    : axis_{axis}
+    , vertex_external_{vertex}
+    , vertex_internal_{converter.pre(vertex)}
+    , scale_factor_{scale_factor}
+    , converter_{converter}
+    , distance_function_{distance_function} {};
+
+    template <typename T>
+    [[nodiscard]] bool IsEntryValid(const KeyInternal& key, const T& value) const {
+        if (
+            (key[axis_] > vertex_internal_[axis_] && scale_factor_ < 0) ||
+            (key[axis_] < vertex_internal_[axis_] && scale_factor_ > 0)) {
+            // point is on the side of the vertex where the paraboloid does not exist
+            return false;
+        }
+
+        KeyExternal point = converter_.post(key);
+
+        // radius of paraboloid at axis offset of point
+        ScalarExternal diff_axis = point[axis_] - vertex_external_[axis_];
+        ScalarExternal radius_sqr = diff_axis / scale_factor_;
+
+        // point projected onto axis of paraboloid
+        KeyExternal projected = vertex_external_;
+        projected[axis_] = point[axis_];
+
+        // distance of point to the axis of the paraboloid
+        ScalarExternal dist_point_to_axis = distance_function_(projected, point);
+
+        return dist_point_to_axis * dist_point_to_axis <= radius_sqr;
+    }
+
+    /*
+     * Calculate whether AABB encompassing all possible points in the node intersects with the
+     * paraboloid.
+     */
+    [[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;
+
+        // coordinate on axis that is inside the paraboloid and furthest away from vertex
+        ScalarInternal coord_max_vertex_dist;
+        if (scale_factor_ < 0) {
+            // lower bound is furthest from the vertex
+            coord_max_vertex_dist = prefix[axis_] & node_min_bits;
+            if (coord_max_vertex_dist > vertex_internal_[axis_]) {
+                // point is on the side of the vertex where the paraboloid does not exist
+                return false;
+            }
+        } else {
+            // upper bound is furthest from the vertex
+            coord_max_vertex_dist = prefix[axis_] | node_max_bits;
+            if (coord_max_vertex_dist < vertex_internal_[axis_]) {
+                // point is on the side of the vertex where the paraboloid does not exist
+                return false;
+            }
+        }
+
+        KeyInternal closest_to_axis;
+        for (size_t i = 0; i < DIM; ++i) {
+            if (i == axis_) {
+                closest_to_axis[i] = coord_max_vertex_dist;
+            } else {
+                // 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 vertex for dimension
+                closest_to_axis[i] = std::clamp(vertex_internal_[i], lo, hi);
+            }
+        }
+
+        // closest_to_axis projected onto axis of paraboloid
+        KeyExternal projected = vertex_internal_;
+        projected[axis_] = coord_max_vertex_dist;
+
+        KeyExternal closest_point = converter_.post(closest_to_axis);
+        KeyExternal projected_point = converter_.post(projected);
+
+        // radius of paraboloid at axis offset of closest_to_axis
+        ScalarExternal diff_axis = closest_point[axis_] - vertex_external_[axis_];
+        ScalarExternal radius_sqr = sqrt(diff_axis / scale_factor_);
+
+        // distance of point to the axis of the paraboloid
+        ScalarExternal dist_point_to_axis = distance_function_(projected_point, closest_point);
+
+        return dist_point_to_axis * dist_point_to_axis <= radius_sqr;
+    }
+
+  private:
+    const size_t axis_;
+    const KeyExternal vertex_external_;
+    const KeyExternal vertex_internal_;
+    const ScalarExternal scale_factor_;
+    const CONVERTER converter_;
+    const DISTANCE distance_function_;
+};
+
 }  // namespace improbable::phtree
 
 #endif  // PHTREE_COMMON_FILTERS_H

phtree/common/filter_test.cc

@@ -44,6 +44,56 @@ TEST(PhTreeFilterTest, FilterSphereTest) {
     ASSERT_FALSE(filter.IsEntryValid({3, 8}, nullptr));
 }
 
+TEST(PhTreeFilterTest, FilterParaboloidTest) {
+    FilterParaboloid<ConverterNoOp<2, scalar_64_t>, DistanceEuclidean<2>> filter{1, {5, 3}, 2};
+    // root is always valid
+    ASSERT_TRUE(filter.IsNodeValid({0, 0}, 63));
+    // valid because node encompasses vertex
+    ASSERT_TRUE(filter.IsNodeValid({1, 1}, 10));
+    // valid because node cuts the parabola
+    ASSERT_TRUE(filter.IsNodeValid({4, 4}, 2));
+    // valid because parabola encompasses the node
+    ASSERT_TRUE(filter.IsNodeValid({6, 20}, 1));
+    // valid because parabola touches the corner of the node
+    ASSERT_TRUE(filter.IsNodeValid({6, 4}, 1));
+    // invalid because node is beside the parabola
+    ASSERT_FALSE(filter.IsNodeValid({2, 8}, 1));
+    // invalid because node is below the parabola
+    ASSERT_FALSE(filter.IsNodeValid({4, 0}, 1));
+
+    // valid because point is the vertex
+    ASSERT_TRUE(filter.IsEntryValid({5, 3}, nullptr));
+    // valid because point is within parabola on axis
+    ASSERT_TRUE(filter.IsEntryValid({5, 4}, nullptr));
+    // valid because point is within parabola
+    ASSERT_TRUE(filter.IsEntryValid({6, 7}, nullptr));
+    // valid because point on parabola
+    ASSERT_TRUE(filter.IsEntryValid({6, 5}, nullptr));
+    // invalid because point is below parabola
+    ASSERT_FALSE(filter.IsEntryValid({4, 2}, nullptr));
+    // invalid because point is outside parabola
+    ASSERT_FALSE(filter.IsEntryValid({2, 6}, nullptr));
+
+    FilterParaboloid<ConverterNoOp<2, scalar_64_t>, DistanceEuclidean<2>> filter2{0, {8, 10}, -1};
+    // valid because parabola encompasses the node
+    ASSERT_TRUE(filter2.IsNodeValid({0, 8}, 2));
+    // valid because node cuts the parabola and contains the vertex
+    ASSERT_TRUE(filter2.IsNodeValid({6, 10}, 2));
+    // invalid because node is outside parabola
+    ASSERT_FALSE(filter2.IsNodeValid({6, 4}, 2));
+    // invalid because node is beyond parabola
+    ASSERT_FALSE(filter2.IsNodeValid({16, 8}, 3));
+
+    // valid because point is the vertex
+    ASSERT_TRUE(filter2.IsEntryValid({8, 10}, nullptr));
+    // valid because point is within parabola
+    ASSERT_TRUE(filter2.IsEntryValid({2, 8}, nullptr));
+    // invalid because point is outside parabola
+    ASSERT_FALSE(filter2.IsEntryValid({4, 6}, nullptr));
+    // invalid because point is beyond parabola
+    ASSERT_FALSE(filter2.IsEntryValid({10, 8}, nullptr));
+}
+
 TEST(PhTreeFilterTest, BoxFilterTest) {
     FilterAABB<ConverterNoOp<2, scalar_64_t>> filter{{3, 3}, {7, 7}};
     // root is always valid