add support & test for scalar scaling argument to mat_compose

pylinalg/matrix.py

@@ -323,27 +323,28 @@ def mat_compose(translation, rotation, scaling, /, *, out=None, dtype=None):
     wy = w * y2
     wz = w * z2
 
+    scaling = np.asarray(scaling)
+    if scaling.ndim == 0:
+        scaling = np.array([scaling, scaling, scaling])
+    elif scaling.size == 1:
+        scaling = np.array([scaling[0], scaling[0], scaling[0]])
     sx, sy, sz = scaling
 
-    out.flat[0] = (1 - (yy + zz)) * sx
-    out.flat[4] = (xy + wz) * sx
-    out.flat[8] = (xz - wy) * sx
-    out.flat[12] = 0
-
-    out.flat[1] = (xy - wz) * sy
-    out.flat[5] = (1 - (xx + zz)) * sy
-    out.flat[9] = (yz + wx) * sy
-    out.flat[13] = 0
-
-    out.flat[2] = (xz + wy) * sz
-    out.flat[6] = (yz - wx) * sz
-    out.flat[10] = (1 - (xx + yy)) * sz
-    out.flat[14] = 0
-
-    out.flat[3] = translation[0]
-    out.flat[7] = translation[1]
-    out.flat[11] = translation[2]
-    out.flat[15] = 1
+    out[0, 0] = (1 - (yy + zz)) * sx
+    out[1, 0] = (xy + wz) * sx
+    out[2, 0] = (xz - wy) * sx
+    out[3, 0:3] = 0
+
+    out[0, 1] = (xy - wz) * sy
+    out[1, 1] = (1 - (xx + zz)) * sy
+    out[2, 1] = (yz + wx) * sy
+
+    out[0, 2] = (xz + wy) * sz
+    out[1, 2] = (yz - wx) * sz
+    out[2, 2] = (1 - (xx + yy)) * sz
+
+    out[0:3, 3] = translation
+    out[3, 3] = 1
 
     return out
 

tests/test_matrix.py

@@ -282,12 +282,25 @@ def naive_mat_compose(translation, rotation, scaling, /, *, out=None, dtype=None
 
 def test_mat_compose_naive():
     """Compare the direct composition with the naive composition."""
-    npt.assert_almost_equal(
+    npt.assert_equal(
         la.mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], [1, -2, 9]),
         naive_mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], [1, -2, 9]),
     )
 
 
+def test_mat_compose_scalar_scaling():
+    """Check that a scaler scaling argument is supported in mat_compose."""
+    npt.assert_equal(
+        la.mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], [1.25, 1.25, 1.25]),
+        la.mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], 1.25),
+    )
+
+    npt.assert_equal(
+        la.mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], [1.25, 1.25, 1.25]),
+        la.mat_compose([1, 2, 3], [np.pi, np.pi / 4, 0, 1], [1.25]),
+    )
+
+
 def test_mat_perspective():
     a = la.mat_perspective(-1, 1, -1, 1, 1, 100)
     npt.assert_array_almost_equal(