Implement faster annulus check (#1100)

Implement annulus check from global coordinates

Author: niermann999

core/include/detray/geometry/shapes/annulus2D.hpp

@@ -14,6 +14,7 @@
 #include "detray/definitions/indexing.hpp"
 #include "detray/definitions/math.hpp"
 #include "detray/definitions/units.hpp"
+#include "detray/geometry/coordinates/cartesian3D.hpp"
 #include "detray/geometry/coordinates/polar2D.hpp"
 #include "detray/geometry/detail/shape_utils.hpp"
 #include "detray/geometry/detail/vertexer.hpp"
@@ -160,11 +161,60 @@ class annulus2D {
             std::numeric_limits<dscalar<algebra_t>>::epsilon(),
         const dscalar<algebra_t> edge_tol = 0.f) const {
 
-        // Get the full local position
-        const dpoint2D<algebra_t> loc_p =
-            local_frame_type<algebra_t>::global_to_local(trf, glob_p, {});
+        using scalar_t = dscalar<algebra_t>;
+
+        // Move point to local plane: Focal frame in cartesian coordinates
+        const auto loc_3D{
+            cartesian3D<algebra_t>::global_to_local(trf, glob_p, {})};
+
+        // Shift local 3D position into beam frame to check the radius
+        const scalar_t new_x{loc_3D[0] + bounds[e_shift_x]};
+        const scalar_t new_y{loc_3D[1] + bounds[e_shift_y]};
+
+        const scalar_t r_beam{
+            math::sqrt(math::fma(new_x, new_x, new_y * new_y))};
+
+        auto inside_mask = ((r_beam + tol) >= bounds[e_min_r]) &&
+                           (r_beam <= (bounds[e_max_r] + tol));
+
+        // Try to avoid the costly phi calculation
+        auto phi_focal{detail::invalid_value<scalar_t>()};
+        if (detail::any_of(inside_mask)) {
+            // Get phi for phi-bounds check and rotate by average phi
+            phi_focal = vector::phi(loc_3D) - bounds[e_average_phi];
+            // Estimate angular tolerance along r
+            const scalar_t phi_tol{detail::phi_tolerance(tol, r_beam)};
+
+            inside_mask = (phi_focal >= (bounds[e_min_phi_rel] - phi_tol)) &&
+                          (phi_focal <= (bounds[e_max_phi_rel] + phi_tol)) &&
+                          inside_mask;
+        }
+
+        decltype(inside_mask) inside_edge{false};
+        if (detail::any_of(edge_tol > 0.f)) {
+            // Edge tolerance
+            const scalar_t full_tol{tol + edge_tol};
+
+            inside_edge = ((r_beam + full_tol) >= bounds[e_min_r]) &&
+                          (r_beam <= (bounds[e_max_r] + full_tol));
+
+            if (detail::any_of(inside_edge)) {
+                // If phi had not been calculated before, do it now
+                if (detail::is_invalid_value(phi_focal)) {
+                    phi_focal = vector::phi(loc_3D) - bounds[e_average_phi];
+                }
+
+                const scalar_t phi_tol_full{
+                    detail::phi_tolerance(full_tol, r_beam)};
+
+                inside_edge =
+                    (phi_focal >= (bounds[e_min_phi_rel] - phi_tol_full)) &&
+                    (phi_focal <= (bounds[e_max_phi_rel] + phi_tol_full)) &&
+                    inside_edge;
+            }
+        }
 
-        return check_boundaries(bounds, loc_p, tol, edge_tol);
+        return result_type<decltype(inside_mask)>{inside_mask, inside_edge};
     }
 
     /// @note the point is expected to be given in local coordinates by the
@@ -188,36 +238,54 @@ class annulus2D {
 
         // Check phi boundaries, which are well def. in focal frame
         const scalar_t phi_tol = detail::phi_tolerance(tol, loc_p[0]);
-        const auto phi_check =
-            !((phi_focal < (bounds[e_min_phi_rel] - phi_tol)) ||
-              (phi_focal > (bounds[e_max_phi_rel] + phi_tol)));
+        auto inside_mask = !((phi_focal < (bounds[e_min_phi_rel] - phi_tol)) ||
+                             (phi_focal > (bounds[e_max_phi_rel] + phi_tol)));
+
+        // Try to avoid the costly r_beam calculation
+        auto r_beam2{detail::invalid_value<scalar_t>()};
+        if (detail::any_of(inside_mask)) {
 
-        const auto r_beam2 = get_r2_beam_frame(bounds, loc_p);
+            r_beam2 = get_r2_beam_frame(bounds, loc_p);
 
-        // Apply tolerances as squares: 0 <= a, 0 <= b: a^2 <= b^2 <=> a <= b
-        const scalar_t minR_tol =
-            math::max(bounds[e_min_r] - tol, scalar_t(0.f));
-        const scalar_t maxR_tol = bounds[e_max_r] + tol;
+            // Apply tolerances as squares: 0 <= a, 0 <= b: a^2 <= b^2 <=> a <=
+            // b
+            const scalar_t minR_tol =
+                math::max(bounds[e_min_r] - tol, scalar_t(0.f));
+            const scalar_t maxR_tol = bounds[e_max_r] + tol;
 
-        assert(detail::all_of(minR_tol >= scalar_t(0.f)));
+            assert(detail::all_of(minR_tol >= scalar_t(0.f)));
 
-        auto inside_mask{((r_beam2 >= (minR_tol * minR_tol)) &&
-                          (r_beam2 <= (maxR_tol * maxR_tol))) &&
-                         phi_check};
+            inside_mask = (r_beam2 >= (minR_tol * minR_tol)) &&
+                          (r_beam2 <= (maxR_tol * maxR_tol)) && inside_mask;
+        }
 
         decltype(inside_mask) inside_edge{false};
         if (detail::any_of(edge_tol > 0.f)) {
             // Edge tolerance
             const scalar_t full_tol{tol + edge_tol};
-            const scalar_t minR_tol_edge =
-                math::max(bounds[e_min_r] - full_tol, scalar_t(0.f));
-            const scalar_t maxR_tol_edge = bounds[e_max_r] + full_tol;
+            const scalar_t phi_tol_full =
+                detail::phi_tolerance(full_tol, loc_p[0]);
+
+            const auto phi_check_edge =
+                (phi_focal >= (bounds[e_min_phi_rel] - phi_tol_full)) &&
+                (phi_focal <= (bounds[e_max_phi_rel] + phi_tol_full));
+
+            if (detail::any_of(inside_edge)) {
+                // If phi had not been calculated before, do it now
+                if (detail::is_invalid_value(r_beam2)) {
+                    r_beam2 = get_r2_beam_frame(bounds, loc_p);
+                }
+
+                const scalar_t minR_tol_edge =
+                    math::max(bounds[e_min_r] - full_tol, scalar_t(0.f));
+                const scalar_t maxR_tol_edge = bounds[e_max_r] + full_tol;
 
-            assert(detail::all_of(minR_tol_edge >= scalar_t(0.f)));
+                assert(detail::all_of(minR_tol_edge >= scalar_t(0.f)));
 
-            inside_edge = ((r_beam2 >= (minR_tol_edge * minR_tol_edge)) &&
-                           (r_beam2 <= (maxR_tol_edge * maxR_tol_edge))) &&
-                          phi_check;
+                inside_edge = (r_beam2 >= (minR_tol_edge * minR_tol_edge)) &&
+                              (r_beam2 <= (maxR_tol_edge * maxR_tol_edge)) &&
+                              phi_check_edge;
+            }
         }
 
         return result_type<decltype(inside_mask)>{inside_mask, inside_edge};

tests/unit_tests/cpu/geometry/masks/annulus2D.cpp

@@ -133,9 +133,7 @@ GTEST_TEST(detray_masks, annulus2D_ratio_test) {
         bool operator()(const point3 &p,
                         const mask<annulus2D, test_algebra> &ann,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{ann.to_local_frame3D(trf, p)};
-            return ann.is_inside(loc_p, t);
+            return ann.is_inside(trf, p, t);
         }
     };
 

tests/unit_tests/cpu/geometry/masks/cylinder.cpp

@@ -76,9 +76,7 @@ GTEST_TEST(detray_masks, cylinder2D_ratio_test) {
         bool operator()(const point3 &p,
                         const mask<cylinder2D, test_algebra> &cyl,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{cyl.to_local_frame3D(trf, p)};
-            return cyl.is_inside(loc_p, t);
+            return cyl.is_inside(trf, p, t);
         }
     };
 
@@ -156,9 +154,7 @@ GTEST_TEST(detray_masks, cylinder3D_ratio_test) {
         bool operator()(const point3 &p,
                         const mask<cylinder3D, test_algebra> &cyl,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{cyl.to_local_frame3D(trf, p)};
-            return cyl.is_inside(loc_p, t);
+            return cyl.is_inside(trf, p, t);
         }
     };
 

tests/unit_tests/cpu/geometry/masks/line.cpp

@@ -79,11 +79,9 @@ GTEST_TEST(detray_masks, line_circular_ratio_test) {
     struct mask_check {
         bool operator()(const point3 &p,
                         const mask<line_circular, test_algebra> &st,
-                        const test::transform3 &trf, const test::vector3 &dir,
-                        const scalar t) {
-
-            const point3 loc_p{st.to_local_frame3D(trf, p, dir)};
-            return st.is_inside(loc_p, t);
+                        const test::transform3 &trf,
+                        const test::vector3 & /*dir*/, const scalar t) {
+            return st.is_inside(trf, p, t);
         }
     };
 
@@ -155,11 +153,9 @@ GTEST_TEST(detray_masks, line_square_ratio_test) {
     struct mask_check {
         bool operator()(const point3 &p,
                         const mask<line_square, test_algebra> &dcl,
-                        const test::transform3 &trf, const test::vector3 &dir,
-                        const scalar t) {
-
-            const point3 loc_p{dcl.to_local_frame3D(trf, p, dir)};
-            return dcl.is_inside(loc_p, t);
+                        const test::transform3 &trf,
+                        const test::vector3 & /*dir*/, const scalar t) {
+            return dcl.is_inside(trf, p, t);
         }
     };
 

tests/unit_tests/cpu/geometry/masks/rectangle2D.cpp

@@ -76,9 +76,7 @@ GTEST_TEST(detray_masks, rectangle2D_ratio_test) {
         bool operator()(const point3 &p,
                         const mask<rectangle2D, test_algebra> &r,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{r.to_local_frame3D(trf, p)};
-            return r.is_inside(loc_p, t);
+            return r.is_inside(trf, p, t);
         }
     };
 

tests/unit_tests/cpu/geometry/masks/ring2D.cpp

@@ -74,9 +74,7 @@ GTEST_TEST(detray_masks, ring2D_ratio_test) {
     struct mask_check {
         bool operator()(const point3 &p, const mask<ring2D, test_algebra> &r,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{r.to_local_frame3D(trf, p)};
-            return r.is_inside(loc_p, t);
+            return r.is_inside(trf, p, t);
         }
     };
 

tests/unit_tests/cpu/geometry/masks/trapezoid2D.cpp

@@ -78,9 +78,7 @@ GTEST_TEST(detray_masks, trapezoid2D_ratio_test) {
         bool operator()(const point3 &p,
                         const mask<trapezoid2D, test_algebra> &tp,
                         const test::transform3 &trf, const scalar t) {
-
-            const point3 loc_p{tp.to_local_frame3D(trf, p)};
-            return tp.is_inside(loc_p, t);
+            return tp.is_inside(trf, p, t);
         }
     };