Move overlaps

neurom/core/soma.py

@@ -89,12 +89,7 @@ def volume(self):
 
     def overlaps(self, points, exclude_boundary=False):
         """Check that the given points are located inside the soma."""
-        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
-        if exclude_boundary:
-            mask = np.linalg.norm(points - self.center, axis=1) < self.radius
-        else:
-            mask = np.linalg.norm(points - self.center, axis=1) <= self.radius
-        return mask
+        return soma_overlaps(self, points, exclude_boundary=exclude_boundary)
 
 
 class SomaSinglePoint(Soma):
@@ -144,29 +139,6 @@ def __str__(self):
             self.radius,
         )
 
-    def overlaps(self, points, exclude_boundary=False):
-        """Check that the given points are located inside the soma."""
-        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
-        mask = np.ones(len(points)).astype(bool)
-        for p1, p2 in zip(self.points[:-1], self.points[1:]):
-            vec = p2[COLS.XYZ] - p1[COLS.XYZ]
-            vec_norm = np.linalg.norm(vec)
-            dot = (points[mask] - p1[COLS.XYZ]).dot(vec) / vec_norm
-
-            cross = np.linalg.norm(np.cross(vec, points[mask]), axis=1) / vec_norm
-            dot_clipped = np.clip(dot / vec_norm, a_min=0, a_max=1)
-            radii = p1[COLS.R] * (1 - dot_clipped) + p2[COLS.R] * dot_clipped
-
-            if exclude_boundary:
-                in_cylinder = (dot > 0) & (dot < vec_norm) & (cross < radii)
-            else:
-                in_cylinder = (dot >= 0) & (dot <= vec_norm) & (cross <= radii)
-            mask[np.where(mask)] = ~in_cylinder
-            if not mask.any():
-                break
-
-        return ~mask
-
 
 class SomaNeuromorphoThreePointCylinders(SomaCylinders):
     """NeuroMorpho compatible soma.
@@ -230,58 +202,6 @@ def __str__(self):
             self.radius,
         )
 
-    def overlaps(self, points, exclude_boundary=False):
-        """Check that the given points are located inside the soma.
-
-        The contour is supposed to be in the plane XY, the Z component is ignored.
-        """
-        # pylint: disable=too-many-locals
-        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
-
-        # Convert points to angles from the center
-        relative_pts = points - self.center
-        pt_angles = np.arctan2(relative_pts[:, COLS.Y], relative_pts[:, COLS.X])
-
-        # Convert soma points to angles from the center
-        relative_soma_pts = self.points[:, COLS.XYZ] - self.center
-        soma_angles = np.arctan2(relative_soma_pts[:, COLS.Y], relative_soma_pts[:, COLS.X])
-
-        # Order the soma points by ascending angles
-        soma_angle_order = np.argsort(soma_angles)
-        ordered_soma_angles = soma_angles[soma_angle_order]
-        ordered_relative_soma_pts = relative_soma_pts[soma_angle_order]
-
-        # Find the two soma points which form the segment crossed by the one from the center
-        # to the point
-        angles = np.atleast_2d(pt_angles).T - ordered_soma_angles
-        closest_indices = np.argmin(np.abs(angles), axis=1)
-        neighbors = np.ones_like(closest_indices)
-        neighbors[angles[np.arange(len(closest_indices)), closest_indices] < 0] = -1
-        signs = (neighbors == 1) * 2.0 - 1.0
-        neighbors[(closest_indices >= len(relative_soma_pts) - 1) & (neighbors == 1)] = (
-            -len(relative_soma_pts) + 1
-        )
-
-        # Compute the cross product and multiply by neighbors to get the same result as if all
-        # vectors were clockwise
-        cross_z = (
-            np.cross(
-                (
-                    ordered_relative_soma_pts[closest_indices + neighbors]
-                    - ordered_relative_soma_pts[closest_indices]
-                ),
-                relative_pts - ordered_relative_soma_pts[closest_indices],
-            )[:, COLS.Z]
-            * signs
-        )
-
-        if exclude_boundary:
-            interior_side = cross_z > 0
-        else:
-            interior_side = cross_z >= 0
-
-        return interior_side
-
 
 def soma_center(soma):
     """Calculate soma center."""
@@ -443,6 +363,127 @@ def _soma_three_point_cylinders_volume(morphio_soma):
     return soma_algo(morphio_soma)
 
 
+def soma_overlaps(soma, points, exclude_boundary=False):
+    """Check if soma overlaps with points."""
+
+    if hasattr(soma, "to_morphio"):
+        morphio_soma = soma.to_morphio()
+    else:
+        morphio_soma = soma
+
+    def _soma_undefined_overlaps(morphio_soma, points, exclude_boundary):
+        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
+
+        center = soma_center(morphio_soma)
+        radius = soma_radius(morphio_soma)
+
+        if exclude_boundary:
+            return np.linalg.norm(points - center, axis=1) < radius
+
+        return np.linalg.norm(points - center, axis=1) <= radius
+
+    def _soma_cylinders_overlaps(morphio_soma, points, exclude_boundary):
+
+        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
+
+        soma_points = np.concatenate(
+            (
+                morphio_soma.points,
+                0.5 * morphio_soma.diameters[:, np.newaxis],
+            ),
+            axis=1,
+        )
+
+        mask = np.ones(len(points)).astype(bool)
+        for p1, p2 in zip(soma_points[:-1], soma_points[1:]):
+            vec = p2[COLS.XYZ] - p1[COLS.XYZ]
+            vec_norm = np.linalg.norm(vec)
+            dot = (points[mask] - p1[COLS.XYZ]).dot(vec) / vec_norm
+
+            cross = np.linalg.norm(np.cross(vec, points[mask]), axis=1) / vec_norm
+            dot_clipped = np.clip(dot / vec_norm, a_min=0, a_max=1)
+            radii = p1[COLS.R] * (1 - dot_clipped) + p2[COLS.R] * dot_clipped
+
+            if exclude_boundary:
+                in_cylinder = (dot > 0) & (dot < vec_norm) & (cross < radii)
+            else:
+                in_cylinder = (dot >= 0) & (dot <= vec_norm) & (cross <= radii)
+            mask[np.where(mask)] = ~in_cylinder
+            if not mask.any():
+                break
+
+        return ~mask
+
+    def _soma_simple_contour_overlaps(morphio_soma, points, exclude_boundary):
+
+        soma_points = np.concatenate(
+            (
+                morphio_soma.points,
+                0.5 * morphio_soma.diameters[:, np.newaxis],
+            ),
+            axis=1,
+        )
+        center = soma_center(morphio_soma)
+
+        # pylint: disable=too-many-locals
+        points = np.atleast_2d(np.asarray(points, dtype=np.float64))
+
+        # Convert points to angles from the center
+        relative_pts = points - center
+        pt_angles = np.arctan2(relative_pts[:, COLS.Y], relative_pts[:, COLS.X])
+
+        # Convert soma points to angles from the center
+        relative_soma_pts = soma_points[:, COLS.XYZ] - center
+        soma_angles = np.arctan2(relative_soma_pts[:, COLS.Y], relative_soma_pts[:, COLS.X])
+
+        # Order the soma points by ascending angles
+        soma_angle_order = np.argsort(soma_angles)
+        ordered_soma_angles = soma_angles[soma_angle_order]
+        ordered_relative_soma_pts = relative_soma_pts[soma_angle_order]
+
+        # Find the two soma points which form the segment crossed by the one from the center
+        # to the point
+        angles = np.atleast_2d(pt_angles).T - ordered_soma_angles
+        closest_indices = np.argmin(np.abs(angles), axis=1)
+        neighbors = np.ones_like(closest_indices)
+        neighbors[angles[np.arange(len(closest_indices)), closest_indices] < 0] = -1
+        signs = (neighbors == 1) * 2.0 - 1.0
+        neighbors[(closest_indices >= len(relative_soma_pts) - 1) & (neighbors == 1)] = (
+            -len(relative_soma_pts) + 1
+        )
+
+        # Compute the cross product and multiply by neighbors to get the same result as if all
+        # vectors were clockwise
+        cross_z = (
+            np.cross(
+                (
+                    ordered_relative_soma_pts[closest_indices + neighbors]
+                    - ordered_relative_soma_pts[closest_indices]
+                ),
+                relative_pts - ordered_relative_soma_pts[closest_indices],
+            )[:, COLS.Z]
+            * signs
+        )
+
+        if exclude_boundary:
+            interior_side = cross_z > 0
+        else:
+            interior_side = cross_z >= 0
+
+        return interior_side
+
+    soma_algo = {
+        SomaType.SOMA_UNDEFINED: _soma_undefined_overlaps,
+        SomaType.SOMA_SINGLE_POINT: _soma_undefined_overlaps,
+        SomaType.SOMA_CYLINDERS: _soma_cylinders_overlaps,
+        SomaType.SOMA_NEUROMORPHO_THREE_POINT_CYLINDERS: _soma_cylinders_overlaps,
+        SomaType.SOMA_SIMPLE_CONTOUR: _soma_simple_contour_overlaps,
+    }[morphio_soma.type]
+
+    return soma_algo(morphio_soma, points, exclude_boundary=exclude_boundary)
+
+
+
 
 def make_soma(morphio_soma):
     """Make a soma object from a MorphIO soma.