geometry functions: added unit tests

plotpy/mathutils/geometry.py

@@ -42,138 +42,217 @@
 
 # pylint: disable=C0103
 
-from numpy import (
-    arccos,
-    arctan,
-    array,
-    cos,
-    fabs,
-    linalg,
-    matrix,
-    pi,
-    sign,
-    sin,
-    sqrt,
-    vdot,
-)
+from __future__ import annotations
+
+import numpy as np
 
 # ===============================================================================
 # Transform matrix functions
 # ===============================================================================
 
 
-def translate(tx, ty):
-    """Return translation matrix (NumPy matrix object)"""
-    return matrix([[1, 0, tx], [0, 1, ty], [0, 0, 1]], float)
+def translate(tx: float, ty: float) -> np.matrix:
+    """Return translation matrix (NumPy matrix object)
+
+    Args:
+        tx: Translation along X-axis
+        ty: Translation along Y-axis
+
+    Returns:
+        Translation matrix
+    """
+    return np.matrix([[1, 0, tx], [0, 1, ty], [0, 0, 1]], float)
+
+
+def scale(sx: float, sy: float) -> np.matrix:
+    """Return scale matrix (NumPy matrix object)
+
+    Args:
+        sx: Scale along X-axis
+        sy: Scale along Y-axis
+
+    Returns:
+        Scale matrix
+    """
+    return np.matrix([[sx, 0, 0], [0, sy, 0], [0, 0, 1]], float)
 
 
-def scale(sx, sy):
-    """Return scale matrix (NumPy matrix object)"""
-    return matrix([[sx, 0, 0], [0, sy, 0], [0, 0, 1]], float)
+def rotate(alpha: float) -> np.matrix:
+    """Return rotation matrix (NumPy matrix object)
 
+    Args:
+        alpha: Rotation angle (in radians)
 
-def rotate(alpha):
-    """Return rotation matrix (NumPy matrix object)"""
-    return matrix(
-        [[cos(alpha), -sin(alpha), 0], [sin(alpha), cos(alpha), 0], [0, 0, 1]], float
+    Returns:
+        Rotation matrix
+    """
+    return np.matrix(
+        [
+            [np.cos(alpha), -np.sin(alpha), 0],
+            [np.sin(alpha), np.cos(alpha), 0],
+            [0, 0, 1],
+        ],
+        float,
     )
 
 
-def colvector(x, y):
-    """Return vector (NumPy matrix object) from coordinates"""
-    return matrix([x, y, 1]).T
+def colvector(x: float, y: float) -> np.matrix:
+    """Return vector (NumPy matrix object) from coordinates
+
+    Args:
+        x: x-coordinate
+        y: y-coordinate
+
+    Returns:
+        Vector (NumPy matrix object)
+    """
+    return np.matrix([x, y, 1]).T
 
 
 # ===============================================================================
 # Operations on vectors (from coordinates)
 # ===============================================================================
 
 
-def vector_norm(xa, ya, xb, yb):
-    """Return vector norm: (xa, xb)-->(ya, yb)"""
-    return linalg.norm(array((xb - xa, yb - ya)))
+def vector_norm(xa: float, ya: float, xb: float, yb: float) -> float:
+    """Return vector norm
+
+    Args:
+        xa: x-coordinate of first point
+        ya: y-coordinate of first point
+        xb: x-coordinate of second point
+        yb: y-coordinate of second point
+
+    Returns:
+        Norm of vector (xa, xb)-->(ya, yb)
+    """
+    return np.linalg.norm(np.array((xb - xa, yb - ya)))
+
 
+def vector_projection(dv: np.ndarray, xa: float, ya: float, xb: float, yb: float):
+    """Return vector projection
 
-def vector_projection(dv, xa, ya, xb, yb):
-    """Return vector projection on *dv*: (xa, xb)-->(ya, yb)"""
+    Args:
+        dv: vector to project
+        xa: x-coordinate of first point
+        ya: y-coordinate of first point
+        xb: x-coordinate of second point
+        yb: y-coordinate of second point
+
+    Returns:
+        Projection of *dv* on vector (xa, xb)-->(ya, yb)
+    """
     assert dv.shape == (2,)
-    v_ab = array((xb - xa, yb - ya))
-    u_ab = v_ab / linalg.norm(v_ab)
-    return vdot(u_ab, dv) * u_ab + array((xb, yb))
+    v_ab = np.array((xb - xa, yb - ya))
+    u_ab = v_ab / np.linalg.norm(v_ab)
+    return np.vdot(u_ab, dv) * u_ab + np.array((xb, yb))
+
+
+def vector_rotation(theta: float, dx: float, dy: float) -> tuple[float, float]:
+    """Compute theta-rotation on vector
 
+    Args:
+        theta: Rotation angle
+        dx: x-coordinate of vector
+        dy: y-coordinate of vector
+
+    Returns:
+        Tuple of (x, y) coordinates of rotated vector
+    """
+    return np.array(rotate(theta) * colvector(dx, dy)).ravel()[:2]
 
-def vector_rotation(theta, dx, dy):
-    """Compute theta-rotation on vector *v*, returns vector coordinates"""
-    return array(rotate(theta) * colvector(dx, dy)).ravel()[:2]
 
+def vector_angle(dx: float, dy: float) -> float:
+    """Return vector angle with X-axis
 
-def vector_angle(dx, dy):
-    """Return vector angle with X-axis"""
+    Args:
+        dx: x-coordinate of vector
+        dy: y-coordinate of vector
+
+    Returns:
+        Angle between vector and X-axis (in radians)
+    """
     # sign(dy) ==  1 --> return Arccos()
     # sign(dy) ==  0 --> return  0 if sign(dx) ==  1
     # sign(dy) ==  0 --> return pi if sign(dx) == -1
     # sign(dy) == -1 --> return 2pi-Arccos()
     if dx == 0 and dy == 0:
         return 0.0
     else:
-        sx, sy = sign(dx), sign(dy)
-        acos = arccos(dx / sqrt(dx**2 + dy**2))
-        return sy * (pi * (sy - 1) + acos) + pi * (1 - sy**2) * (1 - sx) * 0.5
+        sx, sy = np.sign(dx), np.sign(dy)
+        acos = np.arccos(dx / np.sqrt(dx**2 + dy**2))
+        return sy * (np.pi * (sy - 1) + acos) + np.pi * (1 - sy**2) * (1 - sx) * 0.5
 
 
 # ===============================================================================
 # Misc.
 # ===============================================================================
 
 
-def compute_center(x1, y1, x2, y2):
-    """
+def compute_center(x1: float, y1: float, x2: float, y2: float) -> tuple[float, float]:
+    """Compute center of rectangle
+
+    Args:
+        x1: x-coordinate of top-left corner
+        y1: y-coordinate of top-left corner
+        x2: x-coordinate of bottom-right corner
+        y2: y-coordinate of bottom-right corner
 
-    :param x1:
-    :param y1:
-    :param x2:
-    :param y2:
-    :return:
+    Returns:
+        Tuple of (x, y) coordinates of center
     """
     return 0.5 * (x1 + x2), 0.5 * (y1 + y2)
 
 
-def compute_rect_size(x1, y1, x2, y2):
-    """
+def compute_rect_size(
+    x1: float, y1: float, x2: float, y2: float
+) -> tuple[float, float]:
+    """Compute rectangle size
+
+    Args:
+        x1: x-coordinate of top-left corner
+        y1: y-coordinate of top-left corner
+        x2: x-coordinate of bottom-right corner
+        y2: y-coordinate of bottom-right corner
 
-    :param x1:
-    :param y1:
-    :param x2:
-    :param y2:
-    :return:
+    Returns:
+        Tuple of (width, height)
     """
-    return x2 - x1, fabs(y2 - y1)
+    return x2 - x1, np.fabs(y2 - y1)
 
 
-def compute_distance(x1, y1, x2, y2):
-    """
+def compute_distance(x1: float, y1: float, x2: float, y2: float) -> float:
+    """Compute distance between two points
+
+    Args:
+        x1: x-coordinate of first point
+        y1: y-coordinate of first point
+        x2: x-coordinate of second point
+        y2: y-coordinate of second point
 
-    :param x1:
-    :param y1:
-    :param x2:
-    :param y2:
-    :return:
+    Returns:
+        Distance between points
     """
-    return sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+    return np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
 
 
-def compute_angle(x1, y1, x2, y2, reverse=False):
-    """
+def compute_angle(
+    x1: float, y1: float, x2: float, y2: float, reverse: bool = False
+) -> float:
+    """Compute angle between two points
+
+    Args:
+        x1: x-coordinate of first point
+        y1: y-coordinate of first point
+        x2: x-coordinate of second point
+        y2: y-coordinate of second point
+        reverse: If True, return the angle in the opposite direction
 
-    :param x1:
-    :param y1:
-    :param x2:
-    :param y2:
-    :param reverse:
-    :return:
+    Returns:
+        Angle between points (in degrees)
     """
     sign = -1 if reverse else 1
     if x2 == x1:
         return 0.0
     else:
-        return arctan(-sign * (y2 - y1) / (x2 - x1)) * 180 / pi
+        return np.arctan(-sign * (y2 - y1) / (x2 - x1)) * 180.0 / np.pi

plotpy/tests/unit/test_geometry.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the terms of the BSD 3-Clause
+# (see plotpy/LICENSE for details)
+
+"""
+Unit tests for geometry module
+"""
+
+import numpy as np
+
+from plotpy.mathutils import geometry as geom
+
+
+def test_transform_matrix() -> None:
+    """Testing transform matrix functions"""
+    x0 = 1.0
+    y0 = 2.0
+    scale_x = 2.0
+    scale_y = 3.0
+    delta_x = 15.3
+    delta_y = 10.4
+    angle = 30.0
+    vect0 = geom.colvector(x0, y0)
+    # print("vect0:", repr(vect0))
+    vect1 = geom.scale(scale_x, scale_y) * vect0
+    assert np.allclose(vect1, np.matrix([[2.0], [6.0], [1.0]]))
+    # print("vect1:", repr(vect1))
+    vect2 = geom.translate(delta_x, delta_y) * vect1
+    assert np.allclose(vect2, np.matrix([[17.3], [16.4], [1.0]]))
+    # print("vect2:", repr(vect2))
+    vect3 = geom.rotate(angle) * vect2
+    assert np.allclose(vect3, np.matrix([[18.87226872], [-14.56322332], [1.0]]))
+    # print("vect3:", repr(vect3))
+    trmat = (
+        geom.scale(1.0 / scale_x, 1.0 / scale_y)
+        * geom.translate(-delta_x, -delta_y)
+        * geom.rotate(-angle)
+    )
+    vect4 = trmat * vect3
+    # print("vect4:", repr(vect4))
+    assert np.allclose(vect0, vect4)
+
+
+def test_vector_operations() -> None:
+    """Test vector operation functions"""
+    xa = 1.0
+    ya = 2.0
+    xb = 3.0
+    yb = 4.0
+    vect = np.array([5.0, 6.0])
+    angle = np.deg2rad(30.0)
+    norm = geom.vector_norm(xa, ya, xb, yb)
+    assert np.allclose(norm, np.sqrt(8.0))
+    proj = geom.vector_projection(vect, xa, ya, xb, yb)
+    # print("proj:", repr(proj))
+    assert np.allclose(proj, np.array([8.5, 9.5]))
+    vrot = geom.vector_rotation(angle, vect[0], vect[1])
+    # print("vrot:", repr(vrot))
+    assert np.allclose(vrot, np.array([1.33012702, 7.69615242]))
+    angle = geom.vector_angle(vect[0], vect[1])
+    # print("angle:", repr(angle))
+    assert np.allclose(angle, np.arctan2(vect[1], vect[0]))
+
+
+def test_coordinates_computations() -> None:
+    """Testing coordinates computation functions"""
+    x1 = 1.0
+    y1 = 2.0
+    x2 = 3.0
+    y2 = 4.0
+    xc, yc = geom.compute_center(x1, y1, x2, y2)
+    assert np.allclose(xc, 2.0) and np.allclose(yc, 3.0)
+    width, height = geom.compute_rect_size(x1, y1, x2, y2)
+    assert np.allclose(width, 2.0) and np.allclose(height, 2.0)
+    dist = geom.compute_distance(x1, y1, x2, y2)
+    assert np.allclose(dist, np.sqrt(8.0))
+    angle = geom.compute_angle(x1, y1, x2, y2)
+    assert np.allclose(angle, np.rad2deg(-np.arctan2(2.0, 2.0)))
+    angle_inv = geom.compute_angle(x1, y1, x2, y2, True)
+    assert np.allclose(angle_inv, np.rad2deg(np.arctan2(2.0, 2.0)))
+
+
+if __name__ == "__main__":
+    test_transform_matrix()
+    test_vector_operations()
+    test_coordinates_computations()