diff --git a/src/legendhpges/base.py b/src/legendhpges/base.py index 50ff2d7..21b1806 100644 --- a/src/legendhpges/base.py +++ b/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: diff --git a/src/legendhpges/utils.py b/src/legendhpges/utils.py index 2712cf8..2e87280 100644 --- a/src/legendhpges/utils.py +++ b/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 diff --git a/tests/test_utils.py b/tests/test_utils.py index 442c031..0735b7c 100644 --- a/tests/test_utils.py +++ b/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(