Add spherical coordinate utilities to geo

Needed for virtual fitting algorithm (OpenwaterHealth#147)

Author: ebrahimebrahim

src/openlifu/geo.py

@@ -11,6 +11,7 @@
 from openlifu.util.units import getunitconversion
 
 
+# === Tools to work with points ===
 @dataclass
 class Point:
     position: np.ndarray = field(default_factory=lambda: np.array([0.0, 0.0, 0.0])) # mm
@@ -136,3 +137,81 @@ def to_json(self, compact:bool) -> str:
             return json.dumps(self.to_dict(), separators=(',', ':'))
         else:
             return json.dumps(self.to_dict(), indent=4)
+
+
+# === Tools to work with spherical coordinate systems ===
+
+def cartesian_to_spherical(x:float,y:float,z:float) -> Tuple[float, float, float]:
+    """Convert cartesian coordinates to spherical coordinates
+
+    Args: x, y, z are cartesian coordinates
+    Returns: r, theta, phi, where
+        r is the radial spherical coordinate, a nonnegative float.
+        theta is the polar spherical coordinate, aka the angle off the z-axis, aka the non-azimuthal spherical angle.
+            theta is in the range [0,pi].
+        phi is the azimuthal spherical coordinate, in the range [-pi,pi]
+
+    Angles are in radians.
+    """
+    return (np.sqrt(x**2+y**2+z**2), np.arctan2(np.sqrt(x**2+y**2),z), np.arctan2(y,x))
+
+def spherical_to_cartesian(r: float, th:float, ph:float) -> Tuple[float, float, float]:
+    """Convert spherical coordinates to cartesian coordinates
+
+    Args:
+        r: the radial spherical coordinate
+        th: the polar spherical coordinate theta, aka the angle off the z-axis, aka the non-azimuthal spherical angle
+        ph: the azimuthal spherical coordinate phi
+    Returns the cartesian coordinates x,y,z
+
+    Angles are in radians.
+    """
+    return (r*np.sin(th)*np.cos(ph), r*np.sin(th)*np.sin(ph), r*np.cos(th))
+
+def cartesian_to_spherical_vectorized(p:np.ndarray) -> np.ndarray:
+    """Convert cartesian coordinates to spherical coordinates
+
+    Args:
+        p: an array of shape  (...,3), where the last axis describes point cartesian coordinates x,y,z.
+    Returns: An array of shape (...,3), where the last axis describes point spherical coordinates r, theta, phi, where
+        r is the radial spherical coordinate, a nonnegative float.
+        theta is the polar spherical coordinate, aka the angle off the z-axis, aka the non-azimuthal spherical angle.
+            theta is in the range [0,pi].
+        phi is the azimuthal spherical coordinate, in the range [-pi,pi]
+
+    Angles are in radians.
+    """
+    return np.stack([
+        np.sqrt((p**2).sum(axis=-1)),
+        np.arctan2(np.sqrt((p[...,0:2]**2).sum(axis=-1)),p[...,2]),
+        np.arctan2(p[...,1],p[...,0]),
+    ], axis=-1)
+
+def spherical_to_cartesian_vectorized(p:np.ndarray) -> np.ndarray:
+    """Convert spherical coordinates to cartesian coordinates
+
+    Args:
+        p: an array of shape  (...,3), where the last axis describes point spherical coordinates r, theta, phi, where:
+            r is the radial spherical coordinate
+            theta is the polar spherical coordinate, aka the angle off the z-axis, aka the non-azimuthal spherical angle
+            phi is the azimuthal spherical coordinate
+    Returns the cartesian coordinates x,y,z
+
+    Angles are in radians.
+    """
+    return np.stack([
+        p[...,0]*np.sin(p[...,1])*np.cos(p[...,2]),
+        p[...,0]*np.sin(p[...,1])*np.sin(p[...,2]),
+        p[...,0]*np.cos(p[...,1]),
+    ], axis=-1)
+
+def spherical_coordinate_basis(th:float, phi:float) -> np.ndarray:
+    """Return normalized spherical coordinate basis at a location with spherical polar and azimuthal coordinates (th, phi).
+    The coordinate basis is returned as an array `basis` of shape (3,3), where the rows are the basis vectors,
+    in the order r, theta, phi. So `basis[0], basis[1], basis[2]` are the vectors $\\hat{r}$, $\\hat{\\theta}$, $\\hat{\\phi}$.
+    Angles are assumed to be provided in radians."""
+    return np.array([
+        [np.sin(th)*np.cos(phi), np.sin(th)*np.sin(phi), np.cos(th)],
+        [np.cos(th)*np.cos(phi), np.cos(th)*np.sin(phi), -np.sin(th)],
+        [-np.sin(phi), np.cos(phi), 0],
+    ])

src/openlifu/seg/skinseg.py

@@ -23,6 +23,8 @@
 from scipy.ndimage import distance_transform_edt
 from vtk.util.numpy_support import numpy_to_vtk
 
+from openlifu.geo import cartesian_to_spherical
+
 
 def take_largest_connected_component(mask: np.ndarray) -> np.ndarray:
     """Given a boolean image array (or any integer numpy array), return a mask of the largest connected component."""
@@ -223,33 +225,6 @@ def create_closed_surface_from_labelmap(
 
     return surface_mesh
 
-def cartesian_to_spherical(x:float,y:float,z:float) -> Tuple[float, float, float]:
-    """Convert cartesian coordinates to spherical coordinates
-
-    Args: x, y, z are cartesian coordinates
-    Returns: r, theta, phi, where
-        r is the radial spherical coordinate, a nonnegative float.
-        theta is the polar spherical coordinate, aka the angle off the z-axis, aka the non-azimuthal spherical angle.
-            theta is in the range [0,pi].
-        phi is the azimuthal spherical coordinate, in the range [-pi,pi]
-
-    Angles are in radians.
-    """
-    return (np.sqrt(x**2+y**2+z**2), np.arctan2(np.sqrt(x**2+y**2),z), np.arctan2(y,x))
-
-def spherical_to_cartesian(r: float, th:float, ph:float) -> Tuple[float, float, float]:
-    """Convert spherical coordinates to cartesian coordinates
-
-    Args:
-        r: the radial spherical coordinate
-        th: the polar spherical coordinate theta, aka the angle off the z-axis, aka the non-azimuthal spherical angle
-        ph: the azimuthal spherical coordinate phi
-    Returns the cartesian coordinates x,y,z
-
-    Angles are in radians.
-    """
-    return (r*np.sin(th)*np.cos(ph), r*np.sin(th)*np.sin(ph), r*np.cos(th))
-
 def spherical_interpolator_from_mesh(
     surface_mesh: vtk.vtkPolyData,
     origin: Tuple[float, float, float] = (0.,0.,0.),

tests/test_geo.py

@@ -2,9 +2,88 @@
 
 import numpy as np
 
-from openlifu.geo import Point
+from openlifu.geo import (
+    Point,
+    cartesian_to_spherical,
+    cartesian_to_spherical_vectorized,
+    spherical_coordinate_basis,
+    spherical_to_cartesian,
+    spherical_to_cartesian_vectorized,
+)
 
 
 def test_point_from_dict():
     point = Point.from_dict({'position' : [10,20,30],})
     assert (point.position == np.array([10,20,30], dtype=float)).all()
+
+def test_spherical_coordinate_range():
+    """Verify that spherical coordinate output is in the prescribed value ranges"""
+    rng = np.random.default_rng(848)
+    # try all 8 octants of 3D space
+    for sign_x in [-1,1]:
+        for sign_y in [-1,1]:
+            for sign_z in [-1,1]:
+                cartesian_coords = np.array([sign_x, sign_y, sign_z]) * rng.random(size=3)
+                r, th, ph = cartesian_to_spherical(*cartesian_coords)
+                assert r>=0
+                assert 0 <= th <= np.pi
+                assert -np.pi <= ph <= np.pi
+
+def test_spherical_coordinate_conversion_inverse():
+    """Verify that the spherical coordinate conversion forward and backward functions are inverses of one another"""
+    rng = np.random.default_rng(241)
+    # try all 8 octants of 3D space
+    for sign_x in [-1,1]:
+        for sign_y in [-1,1]:
+            for sign_z in [-1,1]:
+                cartesian_coords = np.array([sign_x, sign_y, sign_z]) * rng.random(size=3)
+                np.testing.assert_almost_equal(
+                    spherical_to_cartesian(*cartesian_to_spherical(*cartesian_coords)),
+                    cartesian_coords
+                )
+                np.testing.assert_almost_equal(
+                    cartesian_to_spherical(*spherical_to_cartesian(*cartesian_to_spherical(*cartesian_coords))),
+                    cartesian_to_spherical(*cartesian_coords)
+                )
+
+def test_cartesian_to_spherical_vectorized():
+    rng = np.random.default_rng(35932)
+    points_cartesian = rng.normal(size=(10,3), scale=2) # make 10 random cartesian points
+    points_spherical = cartesian_to_spherical_vectorized(points_cartesian)
+    # Check individual points against the non-vectorized conversion function:
+    for point_cartesian, point_spherical in zip(points_cartesian, points_spherical):
+        assert np.allclose(
+            point_spherical, # result of vectorized converter
+            np.array(cartesian_to_spherical(*point_cartesian)), # non-vectorized converter
+        )
+
+def test_spherical_to_cartesian_vectorized():
+    rng = np.random.default_rng(85932)
+
+    # make 10 random points in spherical coordinates
+    num_pts = 10
+    points_spherical = np.zeros(shape=(num_pts,3))
+    points_spherical[...,0] = rng.random(num_pts)*5 # random r coordinates
+    points_spherical[...,1] = rng.random(num_pts)*np.pi # random theta coordinates
+    points_spherical[...,2] = rng.random(num_pts)*2*np.pi-np.pi # random phi coordinates
+
+    points_cartesian = spherical_to_cartesian_vectorized(points_spherical)
+    # Check individual points against the non-vectorized conversion function:
+    for point_cartesian, point_spherical in zip(points_cartesian, points_spherical):
+        assert np.allclose(
+            point_cartesian, # result of vectorized converter
+            np.array(spherical_to_cartesian(*point_spherical)), # non-vectorized converter
+        )
+
+def test_spherical_coordinate_basis():
+    rng = np.random.default_rng(35235)
+    th = rng.random()*np.pi
+    phi = rng.random()*2*np.pi-np.pi
+    r  = rng.random()*10
+    basis = spherical_coordinate_basis(th,phi)
+    assert np.allclose(basis @ basis.T, np.eye(3)) # verify it is an orthonormal basis
+    r_hat, theta_hat, phi_hat = basis
+    point = np.array(spherical_to_cartesian(r, th, phi))
+    assert np.allclose(np.diff(r_hat / point), 0) # verify that r_hat is a scalar multiple of the cartesian coords
+    assert cartesian_to_spherical_vectorized(point + 0.01*phi_hat)[2] > phi # verify phi_hat points along increasing phi
+    assert cartesian_to_spherical_vectorized(point + 0.01*theta_hat)[1] > th # verify theta_hat points along increasing theta

tests/test_skinseg.py

@@ -13,7 +13,6 @@
     compute_foreground_mask,
     create_closed_surface_from_labelmap,
     spherical_interpolator_from_mesh,
-    spherical_to_cartesian,
     take_largest_connected_component,
     vtk_img_from_array_and_affine,
 )
@@ -89,36 +88,6 @@ def test_create_closed_surface_from_labelmap():
         point_distance_from_sphere_center = np.linalg.norm(point_position - sphere_center, ord=2)
         assert np.abs(point_distance_from_sphere_center - sphere_radius) < 1.
 
-def test_spherical_coordinate_range():
-    """Verify that spherical coordinate output is in the prescribed value ranges"""
-    rng = np.random.default_rng(848)
-    # try all 8 octants of 3D space
-    for sign_x in [-1,1]:
-        for sign_y in [-1,1]:
-            for sign_z in [-1,1]:
-                cartesian_coords = np.array([sign_x, sign_y, sign_z]) * rng.random(size=3)
-                r, th, ph = cartesian_to_spherical(*cartesian_coords)
-                assert r>=0
-                assert 0 <= th <= np.pi
-                assert -np.pi <= ph <= np.pi
-
-def test_spherical_coordinate_conversion_inverse():
-    """Verify that the spherical coordinate conversion forward and backward functions are inverses of one another"""
-    rng = np.random.default_rng(241)
-    # try all 8 octants of 3D space
-    for sign_x in [-1,1]:
-        for sign_y in [-1,1]:
-            for sign_z in [-1,1]:
-                cartesian_coords = np.array([sign_x, sign_y, sign_z]) * rng.random(size=3)
-                np.testing.assert_almost_equal(
-                    spherical_to_cartesian(*cartesian_to_spherical(*cartesian_coords)),
-                    cartesian_coords
-                )
-                np.testing.assert_almost_equal(
-                    cartesian_to_spherical(*spherical_to_cartesian(*cartesian_to_spherical(*cartesian_coords))),
-                    cartesian_to_spherical(*cartesian_coords)
-                )
-
 def test_spherical_interpolator_from_mesh():
     """Check using a torus that the spherical interpolator behaves reasonably"""
     parametric_torus = vtk.vtkParametricTorus()