我有一个3D数组（矢量的2D数组），我想用旋转矩阵变换每个矢量。旋转位于两个独立的2D弧度角值数组中，分别称为cols和rows。
我已经能够让NumPy为我计算角度，而无需Python循环。现在，我正在寻找一种方法，使NumPy也生成旋转矩阵，希望可以大大提高性能。
size = img.shape[:2]
# Create an array that assigns each pixel the percentage of
# the correction (value between -1 and 1, distributed linearly).
cols = np.array([np.arange(size[1]) for __ in range(size[0])]) / (size[1] - 1) * 2 - 1
rows = np.array([np.arange(size[0]) for __ in range(size[1])]).T / (size[0] - 1) * 2 - 1
# Atan distribution based on F-number and Sensor size.
cols = np.arctan(sh * cols / (2 * f))
rows = np.arctan(sv * rows / (2 * f))
### This is the loop that I would like to remove and find a
### clever way to make NumPy do the same operation natively.
for i in range(size[0]):
for j in range(size[1]):
ah = cols[i,j]
av = rows[i,j]
# Y-rotation.
mat = np.matrix([
[ np.cos(ah), 0, np.sin(ah)],
[0, 1, 0],
[-np.sin(ah), 0, np.cos(ah)]
])
# X-rotation.
mat *= np.matrix([
[1, 0, 0],
[0, np.cos(av), -np.sin(av)],
[0, np.sin(av), np.cos(av)]
])
img[i,j] = img[i,j] * mat
return img
有什么聪明的方法可以重写NumPy操作中的循环吗？
最佳答案
（假设img的形状为(a, b, 3)。）
首先，cols和rows不需要完全扩展为(a, b)（您可以写cols[j]代替cols[i,j]）。使用np.linspace可以轻松生成它们：
cols = np.linspace(-1, 1, size[1]) # shape: (b,)
rows = np.linspace(-1, 1, size[0]) # shape: (a,)
cols = np.arctan(sh * cols / (2*f))
rows = np.arctan(sv * rows / (2*f))
然后，我们可以预先计算矩阵的成分。
# shape: (b,)
cos_ah = np.cos(cols)
sin_ah = np.sin(cols)
zeros_ah = np.zeros_like(cols)
ones_ah = np.ones_like(cols)
# shape: (a,)
cos_av = np.cos(rows)
sin_av = np.sin(rows)
zeros_av = np.zeros_like(rows)
ones_av = np.ones_like(rows)
然后构造旋转矩阵：
# shape: (3, 3, b)
y_mat = np.array([
[cos_ah, zeros_ah, sin_ah],
[zeros_ah, ones_ah, zeros_ah],
[-sin_ah, zeros_ah, cos_ah],
])
# shape: (3, 3, a)
x_mat = np.array([
[ones_av, zeros_av, zeros_av],
[zeros_av, cos_av, -sin_av],
[zeros_av, sin_av, cos_av],
])
现在让我们看看。如果有一个循环，我们将编写：
for i in range(size[0]):
for j in range(size[1]):
img[i, j, :] = img[i, j, :] @ y_mat[:, :, j] @ x_mat[:, :, i]
或者，如果我们展开矩阵乘法：
使用
np.einsum可以很好地处理此问题（请注意，i，j，k，m，n完全类似于上面的公式）：img = np.einsum('ijk,kmj,mni->ijn', img, y_mat, x_mat)
总结一下：
size = img.shape[:2]
cols = np.linspace(-1, 1, size[1]) # shape: (b,)
rows = np.linspace(-1, 1, size[0]) # shape: (a,)
cols = np.arctan(sh * cols / (2*f))
rows = np.arctan(sv * rows / (2*f))
cos_ah = np.cos(cols)
sin_ah = np.sin(cols)
zeros_ah = np.zeros_like(cols)
ones_ah = np.ones_like(cols)
cos_av = np.cos(rows)
sin_av = np.sin(rows)
zeros_av = np.zeros_like(rows)
ones_av = np.ones_like(rows)
y_mat = np.array([
[cos_ah, zeros_ah, sin_ah],
[zeros_ah, ones_ah, zeros_ah],
[-sin_ah, zeros_ah, cos_ah],
])
x_mat = np.array([
[ones_av, zeros_av, zeros_av],
[zeros_av, cos_av, -sin_av],
[zeros_av, sin_av, cos_av],
])
return np.einsum('ijk,kmj,mni->ijn', img, y_mat, x_mat)