Restore support for ndim>2 vector arrays in vec_transform (#105)

* support ndim 3+ in vec_transform

* refactor again

* bump version

pylinalg/vector.py

@@ -106,50 +106,30 @@ def vec_transform(
     """
 
     matrix = np.asarray(matrix)
-    vectors = np.asarray(vectors)
-
-    # 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 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:
+    vectors = vec_homogeneous(vectors, w=w)
+
+    # yes, the ndim > 2 version can also handle ndim=1
+    # and ndim=2, but it's slower
+    if vectors.ndim == 1:
+        vectors = matrix @ vectors
+        if projection:
+            vectors = vectors[:-1] / vectors[-1]
+        else:
+            vectors = vectors[:-1]
+    elif vectors.ndim == 2:
         # transposing the vectors array performs better
         # than transposing the matrix
-        vectors = (matmul_matrix @ vectors.T).T
+        vectors = (matrix @ vectors.T).T
+        if projection:
+            vectors = vectors[:, :-1] / vectors[:, -1, None]
+        else:
+            vectors = vectors[:, :-1]
     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]
+        vectors = matrix @ vectors[..., None]
+        if projection:
+            vectors = vectors[..., :-1, 0] / vectors[..., -1, :]
+        else:
+            vectors = vectors[..., :-1, 0]
 
     if out is None:
         out = vectors

pyproject.toml

@@ -2,7 +2,7 @@
 
 [project]
 name = "pylinalg"
-version = "0.6.6"
+version = "0.6.7"
 description = "Linear algebra utilities for Python"
 readme = "README.md"
 license = { file = "LICENSE" }

tests/test_vectors.py

@@ -120,6 +120,36 @@ def test_vec_transform_projection_flag():
             npt.assert_array_equal(result, expected_out)
 
 
+def test_vec_transform_ndim():
+    vectors_2d = np.array(
+        [
+            [1, 0, 0],
+            [1, 2, 3],
+            [1, 1, 1],
+            [1, 1, -1],
+            [0, 0, 0],
+            [7, 8, -9],
+        ],
+        dtype="f8",
+    )
+    translation = np.array([-1, 2, 2], dtype="f8")
+
+    vectors_3d = vectors_2d.reshape((3, 2, 3))
+    vectors_4d = vectors_2d.reshape((6, 1, 1, 3))
+
+    expected_3d = vectors_3d + translation[None, None, :]
+    expected_4d = vectors_4d + translation[None, None, None, :]
+
+    matrix = la.mat_from_translation(translation)
+
+    for projection in [True, False]:
+        result = la.vec_transform(vectors_3d, matrix, projection=projection)
+        npt.assert_array_equal(result, expected_3d)
+
+        result = la.vec_transform(vectors_4d, matrix, projection=projection)
+        npt.assert_array_equal(result, expected_4d)
+
+
 @given(ct.test_spherical, none())
 @example((1, 0, np.pi / 2), (0, 0, 1))
 @example((1, np.pi / 2, np.pi / 2), (1, 0, 0))