向量化 numpy 代码

Question

我正在尝试向量化以下代码 -

for i in range(imL.shape[0]):
    for j in range(imL.shape[1]):
        if j - disp[i][j] >= 0:
            imR[i, j - disp[i][j], :] = imL[i, j, :]

本质上，我想这样做 -

A[i, j - B[i, j], k] = C[i, j, k] for all (i, j, k)

我研究了布尔索引并提出了这个 -

tmp = np.arange(imR.shape[1])
tmp = np.repeat(tmp[np.newaxis, :], imR.shape[0], axis=0)
tmp = np.repeat(tmp[:, :, np.newaxis], imR.shape[2], axis=2)

imR[(tmp - disp) >= 0] = imL

但是，这会引发错误 -

TypeError: NumPy boolean array indexing assignment requires a 0 or 1-dimensional input, input has 3 dimensions

我相信整数索引是解决方案，但是，我想不出任何办法。有没有办法有效地对其进行矢量化？

注意：imL 和 imR 是 N x M x 3 个 ndarray。 disp 是一个 N x M ndarray.

Answer 1

您可以使用 np.indices:

简化临时索引数组的创建

i_indices, j_indices = np.indices(disp.shape)

然后您可以使用自定义公式创建更新的索引数组：

# j - B[i, k]
new_j_ind = j_indices - disp

将更新后的 j 值替换为原始 j 条件：

# if j - disp[i][j] < 0, keep original j index
new_j_ind[new_j_ind < 0] = j_indices[new_j_ind < 0]

并写入数组。

# A[i, j - B[i, j] if j - B[i, j] > 0 else j, k] = C[i, j, k] for all (i, j, k)
imR[i_indices, new_j_ind] = imL

不过要小心：如果 [i, j - B[i, j]] 曾经映射到任何 (i, k) 的相同坐标，则定义不明确。在 for 循环中它定义明确（最后写入的值获胜），在 numpy 向量化代码中它不是。

向量化 numpy 代码

Vectorizing numpy code

numpy

vectorization

numpy-ndarray