From ca09f0cdd247c3d1959cd117b4709db4df11823c Mon Sep 17 00:00:00 2001
From: Korijn van Golen <korijn@gmail.com>
Date: Sat, 22 Feb 2025 13:41:24 +0100
Subject: [PATCH] Optimize vec_transform (#102)

* optimize vec_transform

* bump version
---
 pylinalg/vector.py    | 28 +++++++++++++++++++---------
 pyproject.toml        |  2 +-
 tests/test_vectors.py | 30 +++++++++++++++++++++++++++++-
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/pylinalg/vector.py b/pylinalg/vector.py
index 805d236..eb85576 100644
--- a/pylinalg/vector.py
+++ b/pylinalg/vector.py
@@ -72,7 +72,9 @@ def vec_homogeneous(vectors, /, *, w=1, out=None, dtype=None) -> np.ndarray:
     return out
 
 
-def vec_transform(vectors, matrix, /, *, w=1, out=None, dtype=None) -> np.ndarray:
+def vec_transform(
+    vectors, matrix, /, *, w=1, projection=True, out=None, dtype=None
+) -> np.ndarray:
     """
     Apply a transformation matrix to a vector.
 
@@ -86,6 +88,9 @@ def vec_transform(vectors, matrix, /, *, w=1, out=None, dtype=None) -> np.ndarra
         The value of the scale component of the homogeneous coordinate. This
         affects the result of translation transforms. use 0 (vectors) if the
         translation component should not be applied, 1 (positions) otherwise.
+    projection : bool, optional
+        If False, the matrix is assumed to be purely affine
+        and the homogeneous component is not applied. Default is True.
     out : ndarray, optional
         A location into which the result is stored. If provided, it must have a
         shape that the inputs broadcast to. If not provided or None, a
@@ -102,17 +107,22 @@ def vec_transform(vectors, matrix, /, *, w=1, out=None, dtype=None) -> np.ndarra
 
     matrix = np.asarray(matrix)
 
-    vectors = vec_homogeneous(vectors, w=w, dtype=float)
-    result = matrix @ vectors[..., None]
-    result /= result[..., -1, :][..., None, :]
-    result = result[..., :-1, 0]
-
-    if out is not None:
-        out[:] = result
+    if projection:
+        vectors = vec_homogeneous(vectors, w=w, dtype=float)
+        vectors @= matrix.T
+        vectors[..., :-1] /= vectors[..., -1, None]
+        vectors = vectors[..., :-1]
     else:
-        out = result
+        vectors = np.asarray(vectors, dtype=float, copy=True)
+        vectors @= matrix[:-1, :-1].T
+        vectors += matrix[:-1, -1]
+
+    if out is None:
+        out = vectors
         if dtype is not None:
             out = out.astype(dtype, copy=False)
+    else:
+        out[:] = vectors
 
     return out
 
diff --git a/pyproject.toml b/pyproject.toml
index 27c0c8b..2554418 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,7 +2,7 @@
 
 [project]
 name = "pylinalg"
-version = "0.6.4"
+version = "0.6.5"
 description = "Linear algebra utilities for Python"
 readme = "README.md"
 license = { file = "LICENSE" }
diff --git a/tests/test_vectors.py b/tests/test_vectors.py
index b5e923b..dc69931 100644
--- a/tests/test_vectors.py
+++ b/tests/test_vectors.py
@@ -67,7 +67,7 @@ def test_vector_apply_translation():
     vectors = np.array([[1, 0, 0]])
     expected = np.array([[0, 2, 2]])
     matrix = la.mat_from_translation([-1, 2, 2])
-    result = la.vec_transform(vectors, matrix)
+    result = la.vec_transform(vectors, matrix, projection=False)
 
     npt.assert_array_almost_equal(
         result,
@@ -92,6 +92,34 @@ def test_vec_transform_out():
     assert result is out
 
 
+def test_vec_transform_projection_flag():
+    vectors = np.array(
+        [
+            [1, 0, 0],
+            [1, 2, 3],
+            [1, 1, 1],
+            [0, 0, 0],
+            [7, 8, -9],
+        ],
+        dtype="f4",
+    )
+    translation = np.array([-1, 2, 2], dtype=float)
+    expected = vectors + translation[None, :]
+
+    matrix = la.mat_from_translation(translation)
+
+    for projection in [True, False]:
+        for batch in [True, False]:
+            if batch:
+                vectors_in = vectors
+                expected_out = expected
+            else:
+                vectors_in = vectors[0]
+                expected_out = expected[0]
+            result = la.vec_transform(vectors_in, matrix, projection=projection)
+            npt.assert_array_equal(result, expected_out)
+
+
 @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))