[英]Is it possible to access the current indices during a Numpy vectorized broadcasting operation?
我想使用花式索引,向量化和/或广播来加速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语法,这可能是一个太复杂的问题,但是我认为应该可以实现。
从本质上讲,可以归结为:如何在广播操作期间引用当前索引?
从一维示例开始:
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的每种可能组合中都有两个索引。
附录
这是获得无分胜负结果的通用方法:
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为3D。 在这种情况下,我发现显式的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
,是因为itertools.product
保证了迭代的顺序,所以两个压缩的迭代器将保持同步。
尝试这个:
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)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.