使用 numpy 的逐行矩阵乘法

Question

我想实现“逐行”矩阵乘法。

更具体地说，我想绘制一组箭头，其方向范围为 (-pi, pi)。以下代码是我的实现方式。

scan_phi = np.linspace(-np.pi*0.5, np.pi*0.5, 450)

points = np.ones((450, 2), dtype=np.float)
points[..., 0] = 0.0

n_pts = len(points)

sin = np.sin(scan_phi)
cos = np.cos(scan_phi)

rot = np.append(np.expand_dims(np.vstack([cos, -sin]).T, axis=1), 
                np.expand_dims(np.vstack([sin, cos]).T, axis=1), 
                axis=1)

points_rot = []

for idx, p in enumerate(points):
    points_rot.append(np.matmul(rot[idx], p.T))

points_rot = np.array(points_rot)
sample = points_rot[::10]

ax = plt.axes()
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)

for idx, p in enumerate(sample):
    if idx == 0:
        ax.arrow(0, 0, p[0], p[1], head_width=0.05, head_length=0.1, color='red')
    else:
        ax.arrow(0, 0, p[0], p[1], head_width=0.05, head_length=0.1, fc='k', ec='k')

plt.show()

在我的代码中，“rot”最终是一个包含 (450, 2, 2) 的数组，表示每个箭头，我创建了一个相应的旋转矩阵来旋转它。我有 450 个点存储在“点”(450, 2) 中，我想用这些点来绘制箭头。 （这里的箭头都是用 [0, 1] 初始化的。但是，它可以用不同的值初始化，这就是为什么我想要有 450 个单独的点，而不是仅仅将单个点旋转 450 个不同的角度）

我的做法是使用 for 循环，即对于每个箭头，我单独对其进行转换。

points_rot = []

for idx, p in enumerate(points):
    points_rot.append(np.matmul(rot[idx], p.T))

points_rot = np.array(points_rot)

但是，我想知道是否有更好更简单的方法可以完全通过 numpy 完成此操作，例如一些可以逐行执行矩阵乘法的操作。任何想法将不胜感激，在此先感谢！

Answer 1

这是np.einsum 的一个很好的用例：

aa = np.random.normal(size=(450, 2, 2))
bb = np.random.normal(size=(450, 2))

cc = np.einsum('ijk,ik->ij', aa, bb)

所以cc的每一行都是aa和bb对应行的乘积：

np.allclose(aa[3].dot(bb[3]), cc)  # returns True

解释：爱因斯坦符号ijk,ik->ij 表示：

cc[i,j] = sum(aa[i,j,k] * bb[i,k] for k in range(2))

即，所有未出现在右侧的变量都被相加。