[WIP] Refactor geometry

pytransform3d/geometry.py

@@ -0,0 +1,77 @@
+"""Basic functionality for geometrical shapes."""
+import numpy as np
+from .transformations import check_transform
+
+
+def unit_sphere_surface_grid(n_steps):
+    """Create grid on the surface of a unit sphere in 3D.
+
+    Parameters
+    ----------
+    n_steps : int
+        Number of discrete steps in each dimension of the surface.
+
+    Returns
+    -------
+    x : array, shape (n_steps, n_steps)
+        x-coordinates of grid points.
+
+    y : array, shape (n_steps, n_steps)
+        y-coordinates of grid points.
+
+    z : array, shape (n_steps, n_steps)
+        z-coordinates of grid points.
+    """
+    phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j,
+                          0.0:2.0 * np.pi:n_steps * 1j]
+    sin_phi = np.sin(phi)
+
+    x = sin_phi * np.cos(theta)
+    y = sin_phi * np.sin(theta)
+    z = np.cos(phi)
+
+    return x, y, z
+
+
+def transform_surface(surface2origin, x, y, z):
+    """Transform surface grid.
+
+    Parameters
+    ----------
+    surface2origin : array-like, shape (4, 4)
+        Pose: transformation that will be applied to the surface grid.
+
+    x : array, shape (n_steps, n_steps)
+        x-coordinates of grid points.
+
+    y : array, shape (n_steps, n_steps)
+        y-coordinates of grid points.
+
+    z : array, shape (n_steps, n_steps)
+        z-coordinates of grid points.
+
+    Returns
+    -------
+    x : array, shape (n_steps, n_steps)
+        x-coordinates of transformed grid points.
+
+    y : array, shape (n_steps, n_steps)
+        y-coordinates of transformed grid points.
+
+    z : array, shape (n_steps, n_steps)
+        z-coordinates of transformed grid points.
+    """
+    surface2origin = check_transform(surface2origin)
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+
+    shape = x.shape
+
+    P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
+    P = P.dot(surface2origin[:3, :3].T) + surface2origin[np.newaxis, :3, 3]
+
+    x = P[:, 0].reshape(*shape)
+    y = P[:, 1].reshape(*shape)
+    z = P[:, 2].reshape(*shape)
+    return x, y, z

pytransform3d/geometry.pyi

@@ -0,0 +1,12 @@
+from typing import Tuple
+import numpy as np
+import numpy.typing as npt
+
+
+def unit_sphere_surface_grid(
+        n_steps: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: ...
+
+
+def transform_surface(
+        pose: npt.ArrayLike, x: npt.ArrayLike, y: npt.ArrayLike,
+        z: npt.ArrayLike) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: ...

pytransform3d/plot_utils/_plot_functions.py

@@ -4,9 +4,10 @@
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
 from ._layout import make_3d_axis
 from ._artists import Arrow3D
-from ..transformations import transform, vectors_to_points
+from ..transformations import transform
 from ..rotations import unitx, unitz, perpendicular_to_vectors, norm_vector
 from ..mesh_loader import load_mesh
+from ..geometry import unit_sphere_surface_grid, transform_surface
 
 
 def plot_box(ax=None, size=np.ones(3), A2B=np.eye(4), ax_s=1, wireframe=True,
@@ -131,10 +132,10 @@ def plot_sphere(ax=None, radius=1.0, p=np.zeros(3), ax_s=1, wireframe=True,
     if ax is None:
         ax = make_3d_axis(ax_s)
 
-    phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
-    x = p[0] + radius * np.sin(phi) * np.cos(theta)
-    y = p[1] + radius * np.sin(phi) * np.sin(theta)
-    z = p[2] + radius * np.cos(phi)
+    x, y, z = unit_sphere_surface_grid(n_steps)
+    x = p[0] + radius * x
+    y = p[1] + radius * y
+    z = p[2] + radius * z
 
     if wireframe:
         ax.plot_wireframe(
@@ -424,19 +425,12 @@ def plot_ellipsoid(ax=None, radii=np.ones(3), A2B=np.eye(4), ax_s=1,
 
     radius_x, radius_y, radius_z = radii
 
-    phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
-    x = radius_x * np.sin(phi) * np.cos(theta)
-    y = radius_y * np.sin(phi) * np.sin(theta)
-    z = radius_z * np.cos(phi)
-
-    shape = x.shape
-
-    P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
-    P = transform(A2B, vectors_to_points(P))[:, :3]
+    x, y, z = unit_sphere_surface_grid(n_steps)
+    x *= radius_x
+    y *= radius_y
+    z *= radius_z
 
-    x = P[:, 0].reshape(*shape)
-    y = P[:, 1].reshape(*shape)
-    z = P[:, 2].reshape(*shape)
+    x, y, z = transform_surface(A2B, x, y, z)
 
     if wireframe:
         ax.plot_wireframe(
@@ -490,21 +484,14 @@ def plot_capsule(ax=None, A2B=np.eye(4), height=1.0, radius=1.0,
     if ax is None:
         ax = make_3d_axis(ax_s)
 
-    phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
-    x = radius * np.sin(phi) * np.cos(theta)
-    y = radius * np.sin(phi) * np.sin(theta)
-    z = radius * np.cos(phi)
+    x, y, z = unit_sphere_surface_grid(n_steps)
+    x *= radius
+    y *= radius
+    z *= radius
     z[len(z) // 2:] -= 0.5 * height
     z[:len(z) // 2] += 0.5 * height
 
-    shape = x.shape
-
-    P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
-    P = transform(A2B, vectors_to_points(P))[:, :3]
-
-    x = P[:, 0].reshape(*shape)
-    y = P[:, 1].reshape(*shape)
-    z = P[:, 2].reshape(*shape)
+    x, y, z = transform_surface(A2B, x, y, z)
 
     if wireframe:
         ax.plot_wireframe(

pytransform3d/plot_utils/_plot_functions.pyi

@@ -12,14 +12,14 @@ def plot_box(ax: Union[None, Axes3D] = ..., size: npt.ArrayLike = ...,
 def plot_sphere(
         ax: Union[None, Axes3D] = ..., radius: float = ...,
         p: npt.ArrayLike = ..., ax_s: float = ..., wireframe: bool = ...,
-        n_steps: float = ..., alpha: float = ...,
+        n_steps: int = ..., alpha: float = ...,
         color: str = ...) -> Axes3D: ...
 
 
 def plot_spheres(
         ax: Union[None, Axes3D] = ..., radius: npt.ArrayLike = ...,
         p: npt.ArrayLike = ..., ax_s: float = ..., wireframe: bool = ...,
-        n_steps: float = ..., alpha: npt.ArrayLike = ...,
+        n_steps: int = ..., alpha: npt.ArrayLike = ...,
         color: npt.ArrayLike = ...) -> Axes3D: ...
 
 

pytransform3d/test/test_geometry.py

@@ -0,0 +1,28 @@
+import numpy as np
+
+import pytransform3d.geometry as pg
+import pytransform3d.transformations as pt
+from numpy.testing import assert_array_equal, assert_array_almost_equal
+
+
+def test_unit_sphere():
+    x, y, z = pg.unit_sphere_surface_grid(10)
+    assert_array_equal(x.shape, (10, 10))
+    assert_array_equal(y.shape, (10, 10))
+    assert_array_equal(z.shape, (10, 10))
+
+    P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
+    norms = np.linalg.norm(P, axis=1)
+    assert_array_almost_equal(norms, np.ones_like(norms))
+
+
+def test_transform_surface():
+    x, y, z = pg.unit_sphere_surface_grid(10)
+
+    p = np.array([0.2, -0.5, 0.7])
+    pose = pt.transform_from(R=np.eye(3), p=p)
+    x, y, z = pg.transform_surface(pose, x, y, z)
+
+    P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
+    norms = np.linalg.norm(P - p[np.newaxis], axis=1)
+    assert_array_almost_equal(norms, np.ones_like(norms))

pytransform3d/uncertainty.py

@@ -7,6 +7,7 @@
 from .trajectories import (exponential_coordinates_from_transforms,
                            transforms_from_exponential_coordinates,
                            concat_many_to_one)
+from .geometry import unit_sphere_surface_grid
 
 
 def estimate_gaussian_transform_from_samples(samples):
@@ -452,10 +453,10 @@ def to_projected_ellipsoid(mean, cov, factor=1.96, n_steps=20):
 
     # Grid on ellipsoid in exponential coordinate space
     radius_x, radius_y, radius_z = radii
-    phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
-    x = radius_x * np.sin(phi) * np.cos(theta)
-    y = radius_y * np.sin(phi) * np.sin(theta)
-    z = radius_z * np.cos(phi)
+    x, y, z = unit_sphere_surface_grid(n_steps)
+    x *= radius_x
+    y *= radius_y
+    z *= radius_z
     P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
     P = np.dot(P, vecs[:, :3].T)