optimize vec_transform again (#104)

Author: Korijn

pylinalg/vector.py

@@ -105,28 +105,51 @@ def vec_transform(
         transformed vectors
     """
 
-    matrix = np.asarray(matrix, dtype="f8")
+    matrix = np.asarray(matrix)
+    vectors = np.asarray(vectors)
 
-    if projection:
-        vectors = vec_homogeneous(vectors, w=w, dtype="f8")
-        if vectors.ndim == 1:
-            vectors = matrix @ vectors
-            vectors[:-1] /= vectors[-1]
-            vectors = vectors[:-1]
-        elif vectors.ndim == 2:
-            vectors = (matrix @ vectors.T).T
-            vectors = vectors[..., :-1] / vectors[..., -1, None]
-        else:
-            raise ValueError("vectors must be a 1D or 2D array")
+    # this code has been micro-optimized for the 1D and 2D cases
+    vectors_ndim = vectors.ndim
+    if vectors_ndim > 2:
+        raise ValueError("vectors must be a 1D or 2D array")
+
+    # determine if we are working with a batch of vectors
+    batch = vectors_ndim != 1
+
+    # we don't need to work in homogeneous vector space
+    # if matrix is purely affine and vectors is a single vector
+    homogeneous = projection or batch
+
+    if homogeneous:
+        vectors = vec_homogeneous(vectors, w=w)
+        matmul_matrix = matrix
     else:
-        if vectors.ndim == 1:
-            vectors = matrix[:-1, :-1] @ vectors + matrix[:-1, -1]
-        elif vectors.ndim == 2:
-            vectors = vec_homogeneous(vectors, w=w, dtype="f8")
-            vectors = (matrix @ vectors.T).T
-            vectors = vectors[..., :-1]
-        else:
-            raise ValueError("vectors must be a 1D or 2D array")
+        # if we are not working in homogeneous space, it's
+        # more efficient to matmul the 3x3 (scale + rotation)
+        # part of the matrix with the vectors and then add
+        # the translation part after
+        matmul_matrix = matrix[:-1, :-1]
+
+    if batch:
+        # transposing the vectors array performs better
+        # than transposing the matrix
+        vectors = (matmul_matrix @ vectors.T).T
+    else:
+        vectors = matmul_matrix @ vectors
+        if not homogeneous:
+            # as alluded to before, we add the translation
+            # part of the matrix after the matmul
+            # if we are not working in homogeneous space
+            vectors = vectors + matrix[:-1, -1]
+
+    if projection:
+        # if we are projecting, we divide by the last
+        # element of the vectors array
+        vectors = vectors[..., :-1] / vectors[..., -1, None]
+    elif homogeneous:
+        # if we are NOT projecting but we are working in
+        # homogeneous space, just drop the last element
+        vectors = vectors[..., :-1]
 
     if out is None:
         out = vectors

tests/test_vectors.py

@@ -101,9 +101,9 @@ def test_vec_transform_projection_flag():
             [0, 0, 0],
             [7, 8, -9],
         ],
-        dtype="f4",
+        dtype="f8",
     )
-    translation = np.array([-1, 2, 2], dtype=float)
+    translation = np.array([-1, 2, 2], dtype="f8")
     expected = vectors + translation[None, :]
 
     matrix = la.mat_from_translation(translation)