在Numpy中为值对生成矩阵

Question

我有一个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操作中有没有巧妙的方法来重写循环？

Answer 1

（我们假设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)