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