Numpy：制作四元数乘法的批处理版本

Question

我改造了下面的函数

def quaternion_multiply(quaternion0, quaternion1):
    """Return multiplication of two quaternions.

    >>> q = quaternion_multiply([1, -2, 3, 4], [-5, 6, 7, 8])
    >>> numpy.allclose(q, [-44, -14, 48, 28])
    True

    """
    x0, y0, z0, w0 = quaternion0
    x1, y1, z1, w1 = quaternion1
    return numpy.array((
         x1*w0 + y1*z0 - z1*y0 + w1*x0,
        -x1*z0 + y1*w0 + z1*x0 + w1*y0,
         x1*y0 - y1*x0 + z1*w0 + w1*z0,
        -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=numpy.float64)

到批处理版本

def quat_multiply(self, quaternion0, quaternion1):
    x0, y0, z0, w0 = np.split(quaternion0, 4, 1)
    x1, y1, z1, w1 = np.split(quaternion1, 4, 1)

    result = np.array((
         x1*w0 + y1*z0 - z1*y0 + w1*x0,
        -x1*z0 + y1*w0 + z1*x0 + w1*y0,
         x1*y0 - y1*x0 + z1*w0 + w1*z0,
        -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=np.float64)
    return np.transpose(np.squeeze(result))

此函数处理形状为 (?,4) 的 quaternion1 和 quaternion0。现在我希望该函数可以处理任意数量的维度，例如 (?,?,4)。如何做到这一点？

Answer 1

您可以利用 np.rollaxis 将最后一个轴放在前面，帮助我们切出 4 个数组而不实际拆分它们。我们执行所需的操作，最后 send 将第一个轴返回到末尾，以保持输出数组形状与输入相同。因此，我们将有一个通用的 n 维 ndarrays 的解决方案，就像这样 -

def quat_multiply_ndim(quaternion0, quaternion1):
    x0, y0, z0, w0 = np.rollaxis(quaternion0, -1, 0)
    x1, y1, z1, w1 = np.rollaxis(quaternion1, -1, 0)
    result = np.array((
         x1*w0 + y1*z0 - z1*y0 + w1*x0,
        -x1*z0 + y1*w0 + z1*x0 + w1*y0,
         x1*y0 - y1*x0 + z1*w0 + w1*z0,
        -x1*x0 - y1*y0 - z1*z0 + w1*w0), dtype=np.float64)
    return np.rollaxis(result,0, result.ndim)

样本运行-

In [107]: # N-dim arrays
     ...: a1 = np.random.randint(0,9,(2,3,2,4))
     ...: b1 = np.random.randint(0,9,(2,3,2,4))
     ...: 

In [108]: quat_multiply_ndim(a1,b1) # New ndim approach
Out[108]: 
array([[[[ 154.,   48.,   55.,  -57.],
         [  31.,   81.,   29.,  -95.]],

        [[  31.,   14.,   88.,   12.],
         [   3.,   30.,   20.,  -51.]],

        [[ 104.,   61.,  102.,  -39.],
         [   0.,   14.,   14.,  -56.]]],


       [[[ -28.,   36.,   24.,   -8.],
         [  11.,   76.,   -7.,  -36.]],

        [[  54.,    3.,   -2.,  -19.],
         [  52.,   62.,   15.,  -55.]],

        [[  76.,   28.,   28.,  -60.],    <--------|
         [  14.,   54.,   13.,    5.]]]])          |
                                                   |
In [109]: quat_multiply(a1[1,2],b1[1,2]) # Old 2D approach
Out[109]:                                          |
array([[ 76.,  28.,  28., -60.], ------------------|
       [ 14.,  54.,  13.,   5.]])

Answer 2

你就快完成了！您只需要注意如何拆分和连接数组即可：

def quat_multiply(quaternion0, quaternion1):
    x0, y0, z0, w0 = np.split(quaternion0, 4, axis=-1)
    x1, y1, z1, w1 = np.split(quaternion1, 4, axis=-1)

    return np.squeeze(np.stack((
         x1*w0 + y1*z0 - z1*y0 + w1*x0,
        -x1*z0 + y1*w0 + z1*x0 + w1*y0,
         x1*y0 - y1*x0 + z1*w0 + w1*z0,
        -x1*x0 - y1*y0 - z1*z0 + w1*w0), axis=-1), axis=-2)

在这里，我们两次都使用 axis=-1 沿最后一个轴拆分，然后沿最后一个轴连接回去。最后，正如您正确注意到的那样，我们挤出倒数第二个轴。并向您展示它是有效的：

>>> q0 = np.array([-5, 6, 7, 8])
>>> q1 = np.array([1, -2, 3, 4])
>>> q0 = np.tile(q1, (2, 2, 1))
>>> q0
array([[[-5,  6,  7,  8],
        [-5,  6,  7,  8]],
       [[-5,  6,  7,  8],
        [-5,  6,  7,  8]]])
>>> q1 = np.tile(q2, (2, 2, 1))
>>> q = quat_multiply(q0, q1)
array([[[-44, -14,  48,  28],
        [-44, -14,  48,  28]],
       [[-44, -14,  48,  28],
        [-44, -14,  48,  28]]])
>>> q.shape
(2, 2, 4)

希望这就是您所需要的！这应该适用于任意维度和任意数量的维度。

注意： np.split 似乎不适用于列表。所以你只能将数组传递给你的新函数，就像我上面所做的那样。如果你想传递列表，你可以调用

 np.split(np.asarray(quaternion0), 4, -1)

在你的函数中。

此外，您的测试用例似乎有误。我想你已经交换了 quaternion0 和 quaternion1 的位置：我在上面测试 q0 和 q1 时把它们交换回来了。

Answer 3

您只需将 axis-=-1 传递给 np.split 以沿最后一个轴拆分即可获得您想要的行为。

并且由于您的数组具有烦人的大小为 1 的尾随维度，而不是沿着新的维度堆叠，然后将那个维度挤出，您可以简单地连接它们，再次沿着（最后）axis=-1:

def quat_multiply(self, quaternion0, quaternion1):
    x0, y0, z0, w0 = np.split(quaternion0, 4, axis=-1)
    x1, y1, z1, w1 = np.split(quaternion1, 4, axis=-1)
    return np.concatenate(
        (x1*w0 + y1*z0 - z1*y0 + w1*x0,
         -x1*z0 + y1*w0 + z1*x0 + w1*y0,
         x1*y0 - y1*x0 + z1*w0 + w1*z0,
         -x1*x0 - y1*y0 - z1*z0 + w1*w0),
        axis=-1)

请注意，使用这种方法，您不仅可以乘以任意维数的相同形状的四元数堆栈：

>>> a = np.random.rand(6, 5, 4)
>>> b = np.random.rand(6, 5, 4)
>>> quat_multiply(None, a, b).shape
(6, 5, 4)

但是你也得到了很好的广播，它允许你将一堆四元数与一个四元数相乘，而不必 fiddle 尺寸：

>>> a = np.random.rand(6, 5, 4)
>>> b = np.random.rand(4)
>>> quat_multiply(None, a, b).shape
(6, 5, 4)

或者用最少的操作在一行中完成两个堆栈之间的所有叉积：

>>> a = np.random.rand(6, 4)
>>> b = np.random.rand(5, 4)
>>> quat_multiply(None, a[:, None], b).shape
(6, 5, 4)

Numpy：制作四元数乘法的批处理版本

Numpy: make batched version of quaternion multiplication

python

numpy

linear-algebra