使用 numpy 的逐行矩阵乘法

Question

I want to implement a "row wise" matrix multiplication.

More specifically speaking, I want to plot a set of arrows whose directions range from (-pi, pi). The following code is how I implemented it.

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()

In my code, "rot" ends up being an array of (450, 2, 2) meaning for each arrow, I have created a corresponding rotation matrix to rotate it. I have 450 points stored in "points" (450, 2) that I want to draw arrows with. (Here the arrows are all initialized with [0, 1]. However, it can be initialized with different values which is why I want to have 450 individual points instead of just rotating a single point by 450 different angles)

The way I did is using a for-loop, ie for each arrow, I transform it individually.

points_rot = []

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

points_rot = np.array(points_rot)

However, I wonder if there's any nicer and easy way to do this completely through numpy, such as some operations that can perform matrix multiplication row-wise. Any idea will be grateful, thanks in advance!

Answer 1

This is a nice use-case for 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)

So that each row of cc is the product of corresponding rows of aa and bb :

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

Explanation: the Einstein notation ijk,ik->ij is saying:

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

Ie, all variables that do not appear in the right-hand side are summed away.