Batches the quat to vec, rotation, and mat compose functions for upst…

pylinalg/matrix.py

@@ -295,17 +295,20 @@ def mat_compose(
     translation, rotation, scaling, /, *, out=None, dtype=None
 ) -> np.ndarray:
     """
-    Compose a transformation matrix given a translation vector, a
-    quaternion and a scaling vector.
+    Compose transformation matrices given translation vectors, quaternions,
+    and scaling vectors.
 
     Parameters
     ----------
-    translation : number or ndarray, [3]
-        translation vector
-    rotation : ndarray, [4]
-        quaternion
-    scaling : number or ndarray, [3]
-        scaling factor(s)
+    translation : ndarray, [3] or [num_vectors, 3]
+    rotation : ndarray, [4] or [num_vectors, 4]
+    scaling : ndarray, [3] or [num_vectors, 3]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 freshly-allocated array is returned. A tuple must have
+        length equal to the number of outputs.
+    dtype : data-type, optional
+        Overrides the data type of the result.
     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
@@ -316,48 +319,56 @@ def mat_compose(
 
     Returns
     -------
-    ndarray, [4, 4]
-        Transformation matrix
+    ndarray, [num_vectors, 4, 4] or [4, 4]
     """
+    rotation = np.asarray(rotation)
+    translation = np.asarray(translation)
+    scaling = np.asarray(scaling)
+
+    if rotation.ndim == 1:
+        rotation = rotation[None, :]
+    if translation.ndim == 1:
+        translation = translation[None, :]
+    if scaling.ndim == 0:
+        scaling = np.full((1, 3), scaling)
+    elif scaling.ndim == 1 and scaling.shape[0] == 3:
+        scaling = scaling[None, :]
+    elif scaling.ndim == 1:
+        scaling = scaling[:, None] * np.ones(3)
+
+    num_vectors = max(rotation.shape[0], translation.shape[0], scaling.shape[0])
+
     if out is None:
-        out = np.empty((4, 4), dtype=dtype)
-
-    x, y, z, w = rotation
-    x2 = x + x
-    y2 = y + y
-    z2 = z + z
-    xx = x * x2
-    xy = x * y2
-    xz = x * z2
-    yy = y * y2
-    yz = y * z2
-    zz = z * z2
-    wx = w * x2
-    wy = w * y2
-    wz = w * z2
+        out = np.empty((num_vectors, 4, 4), dtype=dtype)
+    else:
+        out[..., :, :] = 0
 
-    scaling = np.asarray(scaling)
-    if scaling.size == 1:
-        scaling = np.broadcast_to(scaling, (3,))
-    sx, sy, sz = scaling
+    x, y, z, w = rotation[:, 0], rotation[:, 1], rotation[:, 2], rotation[:, 3]
 
-    out[0, 0] = (1 - (yy + zz)) * sx
-    out[1, 0] = (xy + wz) * sx
-    out[2, 0] = (xz - wy) * sx
-    out[3, 0:3] = 0
+    x2, y2, z2 = x + x, y + y, z + z
+    xx, xy, xz = x * x2, x * y2, x * z2
+    yy, yz, zz = y * y2, y * z2, z * z2
+    wx, wy, wz = w * x2, w * y2, w * z2
 
-    out[0, 1] = (xy - wz) * sy
-    out[1, 1] = (1 - (xx + zz)) * sy
-    out[2, 1] = (yz + wx) * sy
+    sx, sy, sz = scaling[:, 0], scaling[:, 1], scaling[:, 2]
 
-    out[0, 2] = (xz + wy) * sz
-    out[1, 2] = (yz - wx) * sz
-    out[2, 2] = (1 - (xx + yy)) * sz
+    out[:, 0, 0] = (1 - (yy + zz)) * sx
+    out[:, 1, 0] = (xy + wz) * sx
+    out[:, 2, 0] = (xz - wy) * sx
 
-    out[0:3, 3] = translation
-    out[3, 3] = 1
+    out[:, 0, 1] = (xy - wz) * sy
+    out[:, 1, 1] = (1 - (xx + zz)) * sy
+    out[:, 2, 1] = (yz + wx) * sy
 
-    return out
+    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
+    out[:, 3, :3] = 0  # only unassigned entries
+
+    return out.squeeze(0) if out.shape[0] == 1 else out
 
 
 def mat_decompose(

pylinalg/quaternion.py

@@ -113,15 +113,15 @@ def quat_mul(a, b, /, *, out=None, dtype=None) -> np.ndarray:
 
 
 def quat_from_vecs(source, target, /, *, out=None, dtype=None) -> np.ndarray:
-    """Rotate one vector onto another.
+    """Rotate one vector onto one or more other vectors.
 
-    Create a quaternion that rotates ``source`` onto ``target``.
+    Create quaternion(s) that rotates ``source`` onto ``target``.
 
     Parameters
     ----------
-    source : ndarray, [3]
+    source : ndarray, [3] or [1, 3]
         The vector that should be rotated.
-    target : ndarray, [3]
+    target : ndarray, [3] or [num_vectors, 3]
         The vector that will be rotated onto.
     out : ndarray, optional
         A location into which the result is stored. If provided, it
@@ -133,7 +133,7 @@ def quat_from_vecs(source, target, /, *, out=None, dtype=None) -> np.ndarray:
 
     Returns
     -------
-    ndarray, [4]
+    ndarray, [4] or [num_vectors, 4]
         Quaternion.
 
     Notes
@@ -148,39 +148,41 @@ def quat_from_vecs(source, target, /, *, out=None, dtype=None) -> np.ndarray:
     ``source`` in the direction of ``target``.
 
     """
-
     source = np.asarray(source, dtype=float)
+    if source.ndim == 1:
+        source = source[None, :]
     target = np.asarray(target, dtype=float)
+    if target.ndim == 1:
+        target = target[None, :]
 
+    num_vecs = target.shape[0]
+    result_shape = (num_vecs, 4)
     if out is None:
-        result_shape = np.broadcast_shapes(source.shape, target.shape)[:-1]
-        out = np.empty((*result_shape, 4), dtype=dtype)
+        out = np.empty(result_shape, dtype=dtype)
 
-    axis = np.cross(source, target)
-    angle = np.arctan2(np.linalg.norm(axis), np.dot(source, target))
+    axis = np.cross(source, target)  # (num_pts, 3)
+    axis_norm = np.linalg.norm(axis, axis=-1)  # (num_pts,)
+    angle = np.arctan2(axis_norm, (target @ source.T).squeeze(1))  # (num_pts,)
 
-    # if source and target are parallel, axis will be 0. In this case, we
-    # need to choose a replacement axis, which is any vector that is orthogonal
-    # to source (and/or target).
-    use_fallback = np.linalg.norm(axis, axis=-1) == 0
+    # Handle degenerate case: source and target are parallel (axis is zero vector).
+    # Pick any axis orthogonal to source as a replacement.
+    use_fallback = axis_norm == 0
     if np.any(use_fallback):
-        fallback = np.empty((*use_fallback.shape, 3), dtype=float)
-        fallback = np.atleast_2d(fallback)
+        t = np.broadcast_to(source, (num_vecs, 3))[use_fallback]
 
-        template = source[use_fallback]
-        y_zero = template[..., 1] == 0
-        z_zero = template[..., 2] == 0
-        both_nonzero = ~(y_zero | z_zero)
+        # Better case split:
+        y_zero = t[:, 1] == 0
+        z_zero = t[:, 2] == 0
+        neither_zero = ~y_zero & ~z_zero
 
-        # if any axis is zero, we can use that axis
-        fallback[y_zero, :] = (0, 1, 0)
-        fallback[z_zero, :] = (0, 0, 1)
+        fb = np.empty((y_zero.shape[0], 3), dtype=float)
+        fb[y_zero] = (0.0, 1.0, 0.0)
+        fb[~y_zero & z_zero] = (0.0, 0.0, 1.0)
+        fb[neither_zero, 0] = 0.0
+        fb[neither_zero, 1] = -t[neither_zero, 2]
+        fb[neither_zero, 2] = t[neither_zero, 1]
 
-        # if two axes are non-zero we can use those
-        if np.any(both_nonzero):
-            fallback[both_nonzero, :] = (0, -1, 1) * template[both_nonzero, [0, 2, 1]]
-
-        axis[use_fallback] = np.squeeze(fallback)
+        axis[use_fallback] = fb
 
     return quat_from_axis_angle(axis, angle, out=out)
 
@@ -220,14 +222,14 @@ def quat_inv(quaternion, /, *, out=None, dtype=None) -> np.ndarray:
 def quat_from_axis_angle(axis, angle, /, *, out=None, dtype=None) -> np.ndarray:
     """Quaternion from axis-angle pair.
 
-    Create a quaternion representing the rotation of an given angle
+    Create a quaternion representing the rotation of a given angle
     about a given unit vector
 
     Parameters
     ----------
-    axis : ndarray, [3]
+    axis : ndarray, [num_vectors, 3] or [3]
         Unit vector
-    angle : number
+    angle : number or np.ndarray of shape [num_pts,]
         The angle (in radians) to rotate about axis
     out : ndarray, optional
         A location into which the result is stored. If provided, it
@@ -239,7 +241,7 @@ def quat_from_axis_angle(axis, angle, /, *, out=None, dtype=None) -> np.ndarray:
 
     Returns
     -------
-    ndarray, [4]
+    ndarray, [num_pts, 4] or [4]
         Quaternion.
     """
 
@@ -257,7 +259,7 @@ def quat_from_axis_angle(axis, angle, /, *, out=None, dtype=None) -> np.ndarray:
     out[..., :3] = axis * np.sin(angle / 2).reshape(lengths_shape)
     out[..., 3] = np.cos(angle / 2)
 
-    return out
+    return out.squeeze(0) if out.shape[0] == 1 else out
 
 
 def quat_from_euler(angles, /, *, order="xyz", out=None, dtype=None) -> np.ndarray:

tests/test_matrix.py

@@ -177,6 +177,32 @@ def test_mat_compose():
         ],
     )
 
+def test_mat_compose_batch():
+    """Test that the matrices are composed correctly in SRT order."""
+    # non-uniform scaling such that the test would fail if rotation/scaling are
+    # applied in the incorrect order
+    num_repeats = 5
+    scaling = [1, 2, 1]
+    # quaternion corresponding to 90 degree rotation about z-axis
+    rotation_single = np.array([0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2])
+    rotations = np.zeros((num_repeats, rotation_single.shape[0]))
+    rotations[:, :] = rotation_single[None, :] #Shape (num_repeats, 5)
+    translation = [2, 2, 2]
+    # compose the transform
+    result = la.mat_compose(translation, rotations, scaling)
+
+    for k in range(num_repeats):
+        npt.assert_array_almost_equal(
+            result[k],
+            [
+                [0, -2, 0, 2],
+                [1, 0, 0, 2],
+                [0, 0, 1, 2],
+                [0, 0, 0, 1],
+            ],
+        )
+
+
 
 def test_mat_decompose():
     """Test that the matrices are decomposed correctly."""

tests/test_quaternion.py

@@ -94,6 +94,28 @@ def test_quaternion_from_unit_vectors(
     assert np.allclose(actual, target_direction)
 
 
+@given(ct.test_unit_vector, ct.test_unit_vector, ct.legal_positive_number)
+def test_quaternion_from_unit_vectors_broadcasting(
+    source_direction, target_direction, source_length
+):
+    assume(abs(source_length) > 1e-8)
+
+    num_repeats = 5
+    # Note: the length of the cross product of two large vectors can overflow
+    # and become Inf. to avoid this, we only scale source.
+    source = source_length * source_direction
+    target = np.zeros((num_repeats, np.prod(target_direction.shape)))
+    target[:] = target_direction[None, :]
+
+    rotation = la.quat_from_vecs(source, target)
+
+    for ind in range(num_repeats):
+        actual = la.vec_transform_quat(source_direction, rotation[ind, :])
+        assert np.allclose(actual, target_direction)
+
+
+
+
 def test_quat_inv():
     a = np.array([0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2])
     ai = la.quat_inv(a)