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

Question

我想使用奇特的索引、矢量化和/或广播来加速 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 循环(或对每个索引使用 np.sum 的单个 for 循环)...但如果可能的话，我想使用广播和/或花式索引。对于 Numpy 语法来说，这可能是一个太复杂的问题，但我觉得它应该是可能的。

本质上，它可以归结为:如何在广播操作期间引用当前索引？

Answer 1

从一维示例开始:

x = np.arange(10)

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

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

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

如果您希望保留边缘，可以预先分配输出缓冲区并直接添加到其中:

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

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

希望您能够了解如何将其推广到 2D 情况。您不是在 -1、0、1 之间移动单个索引，而是在两者之间的 -1、0、1 的每种可能组合中拥有两个索引。

附录

以下是获得 no-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

此实现有些初级，因为它不会检查没有维度或形状 < 2 的输入，但它确实适用于任意数量的维度。

请注意，对于 1D 情况，循环将恰好运行 3 次，对于 2D 情况，循环将运行 9 次，对于 ND 情况，循环将运行 3N 次。在这种情况下，我发现显式 for 循环适合 numpy。与在大型数组上完成的工作相比，循环非常小，对于小型数组来说足够快，并且肯定比为 3D 情况手动写出所有 27 种可能性更好。

还有一点需要注意的是连续索引是如何生成的。在 Python 中，带冒号的索引(例如 x[1:2:3])会转换为相对未知的 slice对象:切片(1,2,3)。由于(几乎)所有带逗号的内容都会被解释为元组，因此表达式 x[1:2,::-1, :2] 中的索引完全等同于 (slice( 1, 2), 切片(无, 无, -1), 切片(无, 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保证迭代的顺序，因此两个压缩迭代器将保持同步。