numpy 基于索引的点积

Question

我正在尝试优化转换问题，让 numpy 尽可能多地完成繁重的工作。

在我的例子中，我有一系列坐标集，每个坐标集都必须点缀有相应的索引 roll/pitch/yaw 值。

程序上看起来像这样：

In [1]: import numpy as np 
   ...: from ds.tools.math import rotation_array 
   ...: from math import pi 
   ...:  
   ...: rpy1 = rotation_array(pi, 0.001232, 1.234243) 
   ...: rpy2 = rotation_array(pi/1, 1.325, 0.5674543)
In [2]: rpy1                                                                                                           
Out[2]: 
array([[ 3.30235500e-01,  9.43897768e-01, -1.23199969e-03],
       [ 9.43898485e-01, -3.30235750e-01,  1.22464587e-16],
       [-4.06850342e-04, -1.16288264e-03, -9.99999241e-01]])

In [3]: rpy2                                                                                                           
Out[3]: 
array([[ 2.05192356e-01,  1.30786082e-01, -9.69943863e-01],
       [ 5.37487075e-01, -8.43271987e-01,  2.97991829e-17],
       [-8.17926489e-01, -5.21332290e-01, -2.43328794e-01]])
   ...:  
   ...: a1 = np.array([[-9.64996132, -5.42488639, -3.08443], 
   ...:                [-8.08814188, -4.56431952, -3.01381]]) 
   ...:  
   ...: a2 = np.array([[-6.91346292, -3.91137259, -2.82621], 
   ...:                [-4.34534536, -2.34535546, -4.87692]])

然后我在a1中用rpy1点坐标，在a2中用rpy2点坐标

In [4]: a1.dot(rpy1)                                                                                                   
Out[4]: 
array([[-8.30604694, -7.31349869,  3.09631641],
       [-6.97801968, -6.12357288,  3.0237723 ]])

In [5]: a2.dot(rpy2)                                                                                                   
Out[5]: 
array([[-1.20926993,  3.86756074,  7.3933692 ],
       [ 1.83673215,  3.95195774,  5.40143613]])

我不想遍历 a's 和 rpy's 的列表，而是想在一个操作中完成所有事情。所以我希望通过以下代码实现这种效果，以便 a12 中的每组坐标都点缀有来自 rpy_a.

的相应索引数组

但很明显，从输出中我得到的比我希望的要多：

In [6]: rpy_a = np.array([rpy1, rpy2]) 
   ...:  
   ...: a12 = np.array([a1, a2]) 

In [7]: a12.dot(rpy_a)                                                                                                 
Out[7]: 
array([[[[-8.30604694, -7.31349869,  3.09631641],
         [-2.37306761,  4.92058705, 10.1104514 ]],

        [[-6.97801968, -6.12357288,  3.0237723 ],
         [-1.6478126 ,  4.36234287,  8.57839034]]],


       [[[-5.9738597 , -5.23064061,  2.83472524],
         [-1.20926993,  3.86756074,  7.3933692 ]],

        [[-3.64678058, -3.32137028,  4.88226976],
         [ 1.83673215,  3.95195774,  5.40143613]]]])

我需要的是：

array([[[-8.30604694, -7.31349869,  3.09631641],
        [-6.97801968, -6.12357288,  3.0237723 ]],

       [[-1.20926993,  3.86756074,  7.3933692 ],
        [ 1.83673215,  3.95195774,  5.40143613]]])

谁能告诉我如何做到这一点？

编辑：可运行示例：

import numpy as np

rpy1 = np.array([[ 3.30235500e-01,  9.43897768e-01, -1.23199969e-03],
                 [ 9.43898485e-01, -3.30235750e-01,  1.22464587e-16],
                 [-4.06850342e-04, -1.16288264e-03, -9.99999241e-01]])

rpy2 = np.array([[ 2.05192356e-01,  1.30786082e-01, -9.69943863e-01],
                 [ 5.37487075e-01, -8.43271987e-01,  2.97991829e-17],
                 [-8.17926489e-01, -5.21332290e-01, -2.43328794e-01]])

a1 = np.array([[-9.64996132, -5.42488639, -3.08443],
               [-8.08814188, -4.56431952, -3.01381]])

a2 = np.array([[-6.91346292, -3.91137259, -2.82621],
               [-4.34534536, -2.34535546, -4.87692]])

print(a1.dot(rpy1))
# array([[-8.30604694, -7.31349869,  3.09631641],
#        [-6.97801968, -6.12357288,  3.0237723 ]])
print(a2.dot(rpy2))
# array([[-1.20926993,  3.86756074,  7.3933692 ],
#        [ 1.83673215,  3.95195774,  5.40143613]])

rpy_a = np.array([rpy1, rpy2])
a12 = np.array([a1, a2])

print(a12.dot(rpy_a))
# Result:
# array([[[[-8.30604694, -7.31349869,  3.09631641],
#          [-2.37306761,  4.92058705, 10.1104514 ]],
#         [[-6.97801968, -6.12357288,  3.0237723 ],
#          [-1.6478126 ,  4.36234287,  8.57839034]]],
#        [[[-5.9738597 , -5.23064061,  2.83472524],
#          [-1.20926993,  3.86756074,  7.3933692 ]],
#         [[-3.64678058, -3.32137028,  4.88226976],
#          [ 1.83673215,  3.95195774,  5.40143613]]]])

# Need:
# array([[[-8.30604694, -7.31349869,  3.09631641],
#         [-6.97801968, -6.12357288,  3.0237723 ]],
#        [[-1.20926993,  3.86756074,  7.3933692 ],
#         [ 1.83673215,  3.95195774,  5.40143613]]])

Answer 1

假设你想处理任意数量的数组 rpy1, rpy2, ..., rpyn 和 a1, a2, ..., an，我建议使用显式广播进行显式第一轴连接，仅仅是因为 "显式是比隐式更好:

a12 = np.concatenate([_a[None, ...] for _a in (a1, a2)], axis=0)
rpy_a = np.concatenate([_a[None, ...] for _a in (rpy1, rpy2)], axis=0)

这等于：

a12 = np.array([a1, a2])
rpy_a = np.array([rpy1, rpy2])

np.array 需要更少的代码，也比我的显式方法更快，但我只是喜欢显式定义轴，这样每个阅读代码的人都可以在不执行它的情况下猜出结果形状。

无论您选择什么路径，重要的部分如下：

np.einsum('jnk,jkm->jnm', a12, rpy_a)
# Out:
array([[[-8.30604694, -7.3134987 ,  3.09631641],
        [-6.97801969, -6.12357288,  3.0237723 ]],

       [[-1.20926993,  3.86756074,  7.3933692 ],
        [ 1.83673215,  3.95195774,  5.40143613]]])

使用爱因斯坦求和约定，您可以定义要沿特定轴执行的 np.matmul（对于二维数组等于 np.dot）。
在这种情况下，我们将连接轴 j（或第一个 dim. 或轴 0）定义为共享轴，沿着该轴操作 'nk,km->nm'（等于 np.matmul，参见签名out 参数中的描述）被执行。

也可以简单地使用 np.matmul（或 python 运算符 @）来获得相同的结果：

np.matmul(a12, rpy_a)
a12 @ rpy_a

但同样：对于一般情况，连接轴或形状可能会发生变化，更明确的 np.einsum 更可取。如果您知道不会对形状等进行任何更改，则应该首选 np.matmul（代码更少，速度更快）。

numpy 基于索引的点积

Index based dot product with numpy

numpy

coordinate-transformation