Fix compose

pylinalg/matrix.py

@@ -323,27 +323,26 @@ def mat_compose(translation, rotation, scaling, /, *, out=None, dtype=None):
     wy = w * y2
     wz = w * z2
 
+    scaling = np.asarray(scaling)
+    if scaling.size == 1:
+        scaling = np.broadcast_to(scaling, (3,))
     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

@@ -268,6 +268,39 @@ def test_mat_compose_validation():
         la.mat_decompose(matrix, scaling_signs=signs)
 
 
+def naive_mat_compose(translation, rotation, scaling, /, *, out=None, dtype=None):
+    return la.mat_combine(
+        [
+            la.mat_from_translation(translation),
+            la.mat_from_quat(rotation),
+            la.mat_from_scale(scaling),
+        ],
+        out=out,
+        dtype=dtype,
+    )
+
+
+def test_mat_compose_naive():
+    """Compare the direct composition with the naive composition."""
+    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(