alternative method for horizontal/vertical segements, don't calculate…

… segments which are unlikely to be closer

src/legendhpges/base.py

@@ -18,6 +18,9 @@
 u = get_application_registry()
 
 
+log = logging.getLogger(__name__)
+
+
 class HPGe(ABC, geant4.LogicalVolume):
     """An High-Purity Germanium detector.
 
@@ -151,38 +154,15 @@ def is_inside(self, coords: ArrayLike, tol: float = 1e-11) -> NDArray[bool]:
             on the order of numerical precision of the floating point representation.
         """
 
-        if not isinstance(self.solid, geant4.solid.GenericPolycone):
-            msg = f"distance_to_surface is not implemented for {type(self.solid)} yet"
-            raise NotImplementedError(msg)
-
-        if not isinstance(coords, np.ndarray):
-            coords = np.array(coords)
-
-        if np.shape(coords)[1] != 3:
-            msg = "coords must be provided as a 2D array with x,y,z coordinates for each point."
-            raise ValueError(msg)
-
-        coords_rz = utils.convert_coords(coords)
-
-        # get the profile
-        r, z = self.get_profile()
-        s1, s2 = utils.get_line_segments(r, z, surface_indices=None)
-
-        # convert coords
-        coords_rz = utils.convert_coords(coords)
-
-        # compute shortest distances
-        dists = utils.shortest_distance(s1, s2, coords_rz, tol=tol)
-
-        # fnd the sign of the id with the lowest distance
-        ids = np.argmin(abs(dists), axis=1)
-        return np.where(dists[np.arange(dists.shape[0]), ids] > 0, True, False)
+        dists = self.distance_to_surface(coords, tol=tol, signed=True)
+        return np.where(dists >= 0, True, False)
 
     def distance_to_surface(
         self,
         coords: ArrayLike,
         surface_indices: ArrayLike | None = None,
         tol: float = 1e-11,
+        signed=False,
     ) -> NDArray:
         """Compute the distance of a set of points to the nearest detector surface.
 
@@ -222,10 +202,58 @@ def distance_to_surface(
         # convert coords
         coords_rz = utils.convert_coords(coords)
 
-        # compute shortest distances to every surface
-        dists = utils.shortest_distance(s1, s2, coords_rz, tol=tol, signed=False)
-
-        return np.min(abs(dists), axis=1)
+        diffs = s1 - s2
+        perps = np.abs(diffs[:, 0] * diffs[:, 1]) < tol
+
+        dists = np.full(len(coords_rz), np.inf)
+
+        # get shortest distance to vertical/horizontal surfaces
+        if np.any(perps):
+            dists = utils.shortest_distance(
+                s1[perps], s2[perps], coords_rz, signed=signed
+            )
+            ids = np.argmin(abs(dists), axis=1)
+            dists = dists[np.arange(dists.shape[0]), ids]
+        # compute shortest distances to remaining surfaces
+        # this condition could probably be tightened
+        for start, end in zip(s1[~perps], s2[~perps]):
+            # Calculate distance to endpoints
+            dist_to_start_sq = np.sum((coords_rz - start) ** 2, axis=1)
+            dist_to_end_sq = np.sum((coords_rz - end) ** 2, axis=1)
+
+            # Calculate segment length
+            segment_length_sq = np.sum((end - start) ** 2)
+
+            # Triangle inequality test:
+            # If a point is closer to the segment than its current min distance,
+            # then at least one of these must be true:
+            # 1. It's within (current_dist + segment_length) of the start point
+            # 2. It's within (current_dist + segment_length) of the end point
+
+            threshold_dist_sq = (
+                dists**2 + segment_length_sq + (2 * dists * segment_length_sq)
+            )
+
+            candidates = np.where(
+                (dist_to_start_sq < threshold_dist_sq)
+                | (dist_to_end_sq < threshold_dist_sq)
+            )[0]
+
+            if len(candidates) > 0:
+                dist_candidates = utils.shortest_distance(
+                    np.array([start]),
+                    np.array([end]),
+                    coords_rz[candidates],
+                    tol=tol,
+                    signed=True,
+                )
+                dists[candidates] = np.where(
+                    np.abs(dist_candidates) < np.abs(dists[candidates]),
+                    dist_candidates,
+                    dists[candidates],
+                )
+
+        return dists
 
     @property
     def volume(self) -> Quantity:

src/legendhpges/utils.py

@@ -122,7 +122,7 @@ def get_line_segments(
     return s1, s2
 
 
-@numba.njit(cache=True)
+# @numba.njit(cache=True)
 def shortest_distance(
     s1_list: NDArray,
     s2_list: NDArray,
@@ -186,36 +186,129 @@ def _norm(a):
     for segment in range(n_segments):
         s1 = s1_list[segment]
         s2 = s2_list[segment]
+        # check if vertical or horizontal
+        # Ensure s2's coordinate is always >= s1's coordinate for simpler logic
+        if (s1[0] > s2[0]) or (s1[1] > s2[1]):
+            s1, s2 = s2, s1
+            sign_factor = -1
+        else:
+            sign_factor = 1
+
+        if abs(s1[0] - s2[0]) < tol:
+            # Compute distances for a vertical segment
+            x_level = s1[0]  # x-coordinate of the vertical line
+            y_min, y_max = s1[1], s2[1]
+
+            # Points that project onto the segment (y between y_min and y_max)
+            mask_on_segment = (points[:, 1] >= y_min) & (points[:, 1] <= y_max)
+
+            # Initialize distance vector
+            dist_vec = np.zeros_like(points)
+
+            # For points that project onto the segment (y within bounds)
+            x_diff = points[:, 0] - x_level
+            dist_vec[mask_on_segment, 0] = x_diff[mask_on_segment]
+            dist_vec[mask_on_segment, 1] = 0
+
+            # For points below the segment
+            mask_below = points[:, 1] < y_min
+            dist_vec[mask_below] = points[mask_below] - s1
+
+            # For points above the segment
+            mask_above = points[:, 1] > y_max
+            dist_vec[mask_above] = points[mask_above] - s2
+
+            # Compute sign for vertical segment (positive to the right if s2>s1)
+            if signed:
+                sign_vec_norm = np.ones(len(points))
+                # If points are to the left of the line, sign is negative
+                sign_vec_norm[x_diff > 0 if sign_factor == 1 else x_diff < 0] = -1
+            else:
+                sign_vec_norm = np.ones(len(points))
+
+            normed_dist = np.abs(
+                np.where(
+                    dist_vec[:, 1] == 0,
+                    dist_vec[:, 0],
+                    _norm(dist_vec),
+                )
+            )
+
+        # check if horizontal segment
+        elif abs(s1[1] - s2[1]) < tol:
+            # Compute distances for a horizontal segment
+            y_level = s1[1]  # y-coordinate of the horizontal line
+            x_min, x_max = s1[0], s2[0]
+
+            # Points that project onto the segment (x between x_min and x_max)
+            mask_on_segment = (points[:, 0] >= x_min) & (points[:, 0] <= x_max)
+
+            # Initialize distance vector
+            dist_vec = np.zeros_like(points)
+
+            # For points that project onto the segment (x within bounds)
+            y_diff = points[:, 1] - y_level
+            dist_vec[mask_on_segment, 0] = 0
+            dist_vec[mask_on_segment, 1] = y_diff[mask_on_segment]
+
+            # For points to the left of the segment
+            mask_left = points[:, 0] < x_min
+            dist_vec[mask_left] = points[mask_left] - s1
+
+            # For points to the right of the segment
+            mask_right = points[:, 0] > x_max
+            dist_vec[mask_right] = points[mask_right] - s2
+            # Compute sign for horizontal segment (positive above if s2>s1)
+            if signed:
+                sign_vec_norm = np.ones(len(points))
+                # If points are below the line, sign is negative
+                sign_vec_norm[y_diff < 0 if sign_factor == 1 else y_diff > 0] = -1
+            else:
+                sign_vec_norm = np.ones(len(points))
+
+            normed_dist = np.abs(
+                np.where(
+                    dist_vec[:, 0] == 0,
+                    dist_vec[:, 1],
+                    _norm(dist_vec),
+                )
+            )
 
-        n = (s2 - s1) / _norm(s2 - s1)
+        else:
+            n = (s2 - s1) / _norm(s2 - s1)
 
-        proj_dist = -_dot(n, (n * _dot(s1 - points, n)[:, np.newaxis]))
+            proj_dist = -_dot(n, (n * _dot(s1 - points, n)[:, np.newaxis]))
 
-        dist_vec = np.empty_like(s1 - points)
+            dist_vec = np.empty_like(s1 - points)
 
-        condition1 = proj_dist < 0
-        condition2 = proj_dist > _norm(s2 - s1)
-        condition3 = (~condition1) & (~condition2)
+            condition1 = proj_dist < 0
+            condition2 = proj_dist > _norm(s2 - s1)
+            condition3 = (~condition1) & (~condition2)
 
-        diff_s1 = s1 - points
-        dist_vec[condition1] = diff_s1[condition1]
-        dist_vec[condition2] = s2 - points[condition2]
-        dist_vec[condition3] = (
-            diff_s1[condition3] - n * _dot(diff_s1, n)[condition3, np.newaxis]
-        )
+            diff_s1 = s1 - points
+            dist_vec[condition1] = diff_s1[condition1]
+            dist_vec[condition2] = s2 - points[condition2]
+            dist_vec[condition3] = (
+                diff_s1[condition3] - n * _dot(diff_s1, n)[condition3, np.newaxis]
+            )
 
-        # make this signed so inside is positive and outside negative
-        if signed:
-            sign_vec = n[0] * dist_vec[:, 1] - n[1] * dist_vec[:, 0]
+            # make this signed so inside is positive and outside negative
+            if signed:
+                sign_vec = n[0] * dist_vec[:, 1] - n[1] * dist_vec[:, 0]
 
-            # push points on surface inside
-            sign_vec = np.where(np.abs(sign_vec) < tol, -tol, sign_vec)
-            sign_vec_norm = -sign_vec / np.abs(sign_vec)
+                # push points on surface inside
+                sign_vec = (
+                    np.where(np.abs(sign_vec) < tol, -tol, sign_vec) * sign_factor
+                )
+                sign_vec_norm = -sign_vec / np.abs(sign_vec)
 
-        else:
-            sign_vec_norm = np.ones(len(dist_vec))
+            else:
+                sign_vec_norm = np.ones(len(dist_vec))
+            normed_dist = np.abs(_norm(dist_vec))
 
         dists[:, segment] = np.where(
-            np.abs(_norm(dist_vec)) < tol, tol, np.abs(_norm(dist_vec)) * sign_vec_norm
+            normed_dist < tol,
+            tol,
+            normed_dist * sign_vec_norm,
         )
     return dists

tests/test_utils.py

@@ -30,6 +30,47 @@ def test_distances():
     res = np.array([0.0, 5.0, 7.0, 5.0, 5.0])
     assert np.allclose(utils.shortest_distance(s1, s2, points)[:, 0], res)
 
+    # test 90 degree rotation
+    rot = np.array(
+        [
+            [np.cos(np.deg2rad(90)), -np.sin(np.deg2rad(90))],
+            [np.sin(np.deg2rad(90)), np.cos(np.deg2rad(90))],
+        ]
+    )
+
+    points_new = np.array([rot @ (p_tmp) for p_tmp in points])
+    s1_new = np.array([rot @ (s1_t) for s1_t in s1])
+    s2_new = np.array([rot @ (s2_t) for s2_t in s2])
+    assert np.allclose(
+        utils.shortest_distance(s1_new, s2_new, points_new)[:, 0],
+        res,
+    )
+
+    # test 180 degree rotation of surface but not points
+    s1_new = np.array([[1, 0]])
+    s2_new = np.array([[0, 0]])
+
+    assert np.allclose(
+        utils.shortest_distance(s1_new, s2_new, points)[:, 0],
+        -res,
+    )
+
+    # test "tapered" segment
+    rot = np.array(
+        [
+            [np.cos(np.deg2rad(45)), -np.sin(np.deg2rad(45))],
+            [np.sin(np.deg2rad(45)), np.cos(np.deg2rad(45))],
+        ]
+    )
+
+    points_new = np.array([rot @ (p_tmp) for p_tmp in points])
+    s1_new = np.array([rot @ (s1_t) for s1_t in s1])
+    s2_new = np.array([rot @ (s2_t) for s2_t in s2])
+    assert np.allclose(
+        utils.shortest_distance(s1_new, s2_new, points_new)[:, 0],
+        res,
+    )
+
     # all distances shouldn't be affected by a global offset and rotation
     offset = np.array([107, -203])
     rot = np.array(