是否可以在 Numpy 矢量化广播操作期间访问当前索引？

Posted 2023-04-18

【中文标题】是否可以在 Numpy 矢量化广播操作期间访问当前索引？

【英文标题】：Is it possible to access the current indices during a Numpy vectorized broadcasting operation?

【发布时间】：2019-02-15 04:05:16

【问题描述】：

我想在 Numpy 中使用花哨的索引、矢量化和/或广播来加速单个数组上的函数。对于数组中的每个值，我需要进行涉及相邻值的计算。因此，在我的矢量化操作中，我需要访问当前索引，以便获取它周围的索引。考虑以下简单的数组操作：

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
y[:] = x + 1

我想使用类似的语法，但不是简单的增量，我想做一些事情，比如将相邻索引处的所有值添加到矢量化循环中的当前值。例如，如果索引 [i, j] == 7 周围的区域看起来像

3 2 5
2 7 6
5 5 5

我希望[i, j] 的计算值为3 + 2 + 5 + 2 + 7 + 6 + 5 + 5 + 5，并且我想对所有索引[i, j] 这样做。

这是一个简单的嵌套 for 循环（或单个 for 循环，对每个索引使用 np.sum）...但如果可能，我想使用广播和/或花式索引。对于 Numpy 语法来说，这可能是一个太复杂的问题，但我觉得应该是可能的。

本质上归结为：在广播操作期间如何引用当前索引？

【问题讨论】：

虽然可以访问索引 - 只需构建索引数组并将它们合并到您的计算中 - 这样做并不是使用相邻元素值的好方法。

@user2357112 你能详细说明为什么这不是一个好方法吗？

一旦你得到了索引，你想象你会用它们做什么来获取相邻元素的值？我能想象到的大多数事情要么是不可行的，要么需要显式的 Python 循环，要么会产生数组格式的相邻值，而无需使用显式索引就可以更有效地构建。

简短的回答是，这不是你在 numpy 中做类似事情的方式。

@user2357112 “可以在不使用显式索引的情况下更有效地构建。”这就是我正在寻找的方法——我不在乎使用花哨的索引，除非它能让我达到目标。本质上，我想做我所描述的，但使用 Numpy 的快速编译库而不是慢速 Python 循环。我不接受使用索引数组的想法。

【参考方案1】：

从一维示例开始：

x = np.arange(10)

您必须做出选择：是否丢弃边缘，因为它们没有两个邻居？如果这样做，您基本上可以一步创建输出数组：

result = x[:-2] + x[1:-1] + x[2:]

请注意，所有三个加数都是视图，因为它们使用简单的索引。您希望尽可能避免花哨的索引，因为它通常涉及复制。

如果您更喜欢保留边，可以预先分配输出缓冲区并直接添加到其中：

result = x.copy()
result[:-1] += x[1:]
result[1:] += x[:-1]

这两种情况的基本思想是要将操作应用于所有个相邻元素，只需将数组移动 +/-1。你不需要知道任何索引，或者做任何花哨的事情。越简单越好。

希望您能了解如何将其推广到 2D 案例。与在 -1、0、1 之间移动的单个索引不同，您在它们两者之间的 -1、0、1 的每个可能组合中都有两个索引。

附录

以下是无 egde 结果的通用方法：

from itertools import product
def sum_shifted(a):
    result = np.zeros(tuple(x - 2 for x in a.shape), dtype=a.dtype)
    for index in product([slice(0, -2), slice(1, -1), slice(2, None)], repeat=a.ndim):
        result += a[index]
    return result

这个实现有点初级，因为它不检查没有尺寸或形状

请注意，对于 1D 情况，循环将运行 3 次，2D 为 9 次，ND 为 3N。这是我发现一个明确的 for 循环适合 numpy 的一种情况。与在大型阵列上完成的工作相比，循环非常小，对于小型阵列来说足够快，并且肯定比为 3D 案例手动编写所有 27 种可能性要好。

还有一点需要注意的是连续索引是如何生成的。在 Python 中，带有冒号的索引（如 x[1:2:3]）被转换为相对未知的 slice 对象：slice(1, 2, 3)。因为（几乎）所有带逗号的东西都被解释为一个元组，所以表达式x[1:2, ::-1, :2] 中的索引完全等同于(slice(1, 2), slice(None, None, -1), slice(None, 2))。该循环恰好生成这样一个表达式，每个维度都有一个元素。所以结果实际上是跨所有维度的简单索引。

如果您想保留边缘，可以使用类似的方法。唯一显着的区别是您需要同时索引输入和输出数组：

from itertools import product
def sum_shifted(a):
    result = np.zeros_like(a)
    for r_index, a_index in zip(product([slice(0, -1), slice(None), slice(1, None)], repeat=a.ndim),
                                product([slice(1, None), slice(None), slice(0, -1)], repeat=a.ndim)):
        result[r_index] += a[a_index]
    return result

这是因为itertools.product 保证了迭代的顺序，所以两个压缩的迭代器将保持同步。

【讨论】：

是的，我想这会让我到达那里。但是，现在我想起来了，我应该把代码写成一个循环，然后用 Numba 编译。 “读代码比写代码多，”伟大的 Guido 说——并且向后弯腰编写晦涩难懂的代码来做简单的概念是非常反 Python 的。 （但我发现自己经常这样做！）试图避免使用 Numba 或 Cython。

@physmom：如果你有 Numba，那么对许多高级 NumPy 问题的最佳答案就是“只需编写 for 循环并将其扔给 Numba”。

这是一个使用快速整体阵列构建块的示例。实际上，这是非常 Pythonic 的。列表和字典是我们尽可能将其用作整个对象的构建块。 numpy 更进一步。 MATLAB 过去也需要这种思维方式，但现在它的 JIT 可以让您恢复到迭代思维，而不会造成太多性能损失。

【参考方案2】：

试试这个：

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        if i>0 and i<x.shape[0]-1 and j>0 and j<x.shape[1]-1:
            y[i,j]=x[i,j]+x[i-1,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i,j+1]+x[i+1,j+1]+x[i-1,j+1]+x[i+1,j-1]
        if j==0:
            if i==0:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j+1]+x[i-1,j+1]+x[i-1,j]
            else:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]+x[i-1,j]+x[i-1,j+1]

        if j==x.shape[1]-1:
            if i==0:
                y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j] 
            else:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i-1,j]+x[i+1,j-1]
        if i==0 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]+x[i+1,j+1]+x[i,j+1]
        if i==x.shape[0]-1 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j]+x[i-1,j+1]+x[i,j+1]
print(y)