Broadcasting support for aabb_transform (#107)

* ruff

* fix broadcasting for aabb_transform

Author: Korijn

pylinalg/misc.py

@@ -67,30 +67,40 @@ def aabb_transform(aabb, matrix, /, *, out=None, dtype=None) -> np.ndarray:
     """
 
     aabb = np.asarray(aabb, dtype=float)
-    matrix = np.asarray(matrix, dtype=float).transpose((-1, -2))
+    matrix = np.asarray(matrix, dtype=float)
+
+    # transpose last two dimensions
+    axes = list(range(matrix.ndim))
+    axes[-2:] = axes[-1], axes[-2]
+    matrix = matrix.transpose(axes)
 
     if out is None:
-        out = np.empty_like(aabb, dtype=dtype)
+        # Compute output shape by broadcasting aabb and matrix shapes (excluding last 2 dims)
+        aabb_shape = aabb.shape[:-2]
+        matrix_shape = matrix.shape[:-2]
+        broadcast_shape = np.broadcast_shapes(aabb_shape, matrix_shape)
+        out = np.empty((*broadcast_shape, *aabb.shape[-2:]), dtype=dtype)
 
     corners = np.full(
-        aabb.shape[:-2] + (8, 4),
+        (*aabb.shape[:-2], 8, 4),
         # Fill value of 1 is used for homogeneous coordinates.
         fill_value=1.0,
         dtype=float,
     )
+
     # x
-    corners[..., 0::2, 0] = aabb[..., 0, 0]
-    corners[..., 1::2, 0] = aabb[..., 1, 0]
+    corners[..., 0::2, 0] = aabb[..., 0, 0, np.newaxis]
+    corners[..., 1::2, 0] = aabb[..., 1, 0, np.newaxis]
 
     # y
-    corners[..., 0::4, 1] = aabb[..., 0, 1]
-    corners[..., 1::4, 1] = aabb[..., 0, 1]
-    corners[..., 2::4, 1] = aabb[..., 1, 1]
-    corners[..., 3::4, 1] = aabb[..., 1, 1]
+    corners[..., 0::4, 1] = aabb[..., 0, 1, np.newaxis]
+    corners[..., 1::4, 1] = aabb[..., 0, 1, np.newaxis]
+    corners[..., 2::4, 1] = aabb[..., 1, 1, np.newaxis]
+    corners[..., 3::4, 1] = aabb[..., 1, 1, np.newaxis]
 
     # z
-    corners[..., 0:4, 2] = aabb[..., 0, 2]
-    corners[..., 4:8, 2] = aabb[..., 1, 2]
+    corners[..., 0:4, 2] = aabb[..., 0, 2, np.newaxis]
+    corners[..., 4:8, 2] = aabb[..., 1, 2, np.newaxis]
 
     corners = corners @ matrix
     out[..., 0, :] = np.min(corners[..., :-1], axis=-2)

pylinalg/quaternion.py

@@ -251,7 +251,7 @@ def quat_from_axis_angle(axis, angle, /, *, out=None, dtype=None) -> np.ndarray:
         out = np.empty((*out_shape, 4), dtype=dtype)
 
     # result should be independent of the length of the given axis
-    lengths_shape = axis.shape[:-1] + (1,)
+    lengths_shape = (*axis.shape[:-1], 1)
     axis = axis / np.linalg.norm(axis, axis=-1).reshape(lengths_shape)
 
     out[..., :3] = axis * np.sin(angle / 2).reshape(lengths_shape)

pylinalg/vector.py

@@ -30,7 +30,7 @@ def vec_normalize(vectors, /, *, out=None, dtype=None) -> np.ndarray:
     if out is None:
         out = np.empty_like(vectors, dtype=dtype)
 
-    lengths_shape = vectors.shape[:-1] + (1,)
+    lengths_shape = (*vectors.shape[:-1], 1)
     lengths = np.linalg.norm(vectors, axis=-1).reshape(lengths_shape)
     return np.divide(vectors, lengths, out=out)
 

pyproject.toml

@@ -2,7 +2,7 @@
 
 [project]
 name = "pylinalg"
-version = "0.6.7"
+version = "0.6.8"
 description = "Linear algebra utilities for Python"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -59,5 +59,5 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"__init__.py" = ["F401", "F403"]
+"__init__.py" = ["F401", "F403", "RUF048"]
 "tests/conftest.py" = ["B008"]

tests/conftest.py

@@ -13,13 +13,16 @@
 def pytest_report_header(config):
     # report the CPU model to allow detecting platform-specific problems
     if platform.system() == "Windows":
-        name = (
-            subprocess.check_output(["wmic", "cpu", "get", "name"])
-            .decode()
-            .strip()
-            .split("\n")[1]
-        )
-        cpu_info = " ".join([name])
+        try:
+            name = (
+                subprocess.check_output(["wmic", "cpu", "get", "name"])
+                .decode()
+                .strip()
+                .split("\n")[1]
+            )
+            cpu_info = " ".join([name])
+        except Exception:
+            cpu_info = "Unknown CPU (wmic not available)"
     elif platform.system() == "Linux":
         info_string = subprocess.check_output(["lscpu"]).decode()
         for line in info_string.split("\n"):

tests/test_misc.py

@@ -49,3 +49,51 @@ def test_aabb_transform(point, translation, scale):
     scale_matrix = la.mat_from_scale(scale)
     result = la.aabb_transform(aabb, scale_matrix)
     assert np.allclose(result, np.sort(aabb * scale, axis=0), atol=1e-10)
+
+
+def test_aabb_transform_single():
+    """Test single transform."""
+    aabb = np.array([[-1, -1, -1], [1, 1, 1]])
+    translation = np.array([1, 0, 0])
+    translation_matrix = la.mat_from_translation(translation)
+
+    expected = aabb + translation
+    result = la.aabb_transform(aabb, translation_matrix)
+    assert np.allclose(result, expected, atol=1e-10)
+
+
+def test_aabb_transform_broadcasting():
+    """Test pairwise broadcasting of AABBs and matrices."""
+    aabbs = np.array(
+        [
+            [[-1, -1, -1], [1, 1, 1]],
+            [[-2, -2, -2], [2, 2, 2]],
+        ]
+    )
+    translations = np.array(
+        [
+            [1, 0, 0],
+            [0, 1, 0],
+        ]
+    )
+    translation_matrices = la.mat_from_translation(translations)
+
+    expected = aabbs + translations[:, np.newaxis, :]
+    result = la.aabb_transform(aabbs, translation_matrices)
+    assert np.allclose(result, expected, atol=1e-10)
+
+
+def test_aabb_transform_broadcasting_2():
+    """Test broadcasting many matrices and one AABB."""
+    aabb = np.array([[-1, -1, -1], [1, 1, 1]])
+    translations = np.array(
+        [
+            [1, 0, 0],
+            [0, 1, 0],
+        ]
+    )
+    translation_matrices = la.mat_from_translation(translations)
+
+    expected = aabb[np.newaxis, ...] + translations[:, np.newaxis, :]
+    result = la.aabb_transform(aabb, translation_matrices)
+    assert np.allclose(result, expected, atol=1e-10)