如何向量化包含索引作为变量的函数

Question

NumPy (1.10) 是否可以对包含索引作为变量的函数进行矢量化？ 我有一个 bool 数组 A(i,j,k) 和第二个数组 B（bool 和 3 d，但更大），其中元素应该根据某些条件获得其状态，简化：

if A[i][j][k] == True:
    B[i+1][j][k] = True
    B[i][j+1][k] = True
    B[i][j][k+1] = True

（这样对于 A 中的每个为 True 的元素，B 中在各个方向上距离原始位置 +1 步的元素都设置为 true）

我觉得我无法避免这里的嵌套循环，但我想（数组很大，执行时间很重要），我会很高兴提示如何实现这一点。 （如果问题是noobish，请多多包涵，我是初学者。）

Answer 1

我怀疑你把你的问题简化得太多了，但正如它所说，它可以通过在所有维度上生成移位来解决，然后或所有东西一起：

import numpy as np

# Test data
A = np.zeros([3,5,2], dtype=bool)
A[0,2,1] = True
A[1,4,1] = True

# Generate bigger result matrices
i,j,k = A.shape
B0 = np.zeros([i+1,j+1,k+1], dtype=bool)
B1 = np.zeros([i+1,j+1,k+1], dtype=bool)
B2 = np.zeros([i+1,j+1,k+1], dtype=bool)

# Get partial answers by offsetting A
B0[1:,:-1,:-1] = A
B1[:-1,1:,:-1] = A
B2[:-1,:-1,1:] = A

# Combine answers by oring together parts
B = B0 | B1 | B2

不过，这恐怕不是你想要的。

Answer 2

这适用于二维矩阵：

import numpy as np
from scipy import signal

A = (np.random.normal(0, 1, 64) > 1).reshape(8, 8)
k = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0]])
B = np.minimum(signal.convolve2d(A, k, mode='same'), 1)

给予：

>>> A
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 1, 0, 0, 0],
       [0, 0, 0, 1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1, 0, 0],
       [1, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 1, 1, 0]])

>>> B
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 1, 0, 0],
       [0, 0, 0, 0, 1, 0, 0, 0],
       [0, 0, 1, 1, 0, 0, 1, 0],
       [0, 1, 0, 1, 0, 1, 1, 0],
       [0, 1, 1, 0, 0, 1, 0, 0],
       [1, 0, 0, 0, 0, 0, 1, 1]])

---

可以类似地将其扩展到 N-D 矩阵。

from scipy import ndimage

A = (np.random.normal(0, 1, 64) > 1).reshape(8, 8)
k = np.array([
    [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
    [[0, 1, 0], [1, 0, 0], [0, 0, 0]],
    [[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
B = np.minimum(ndimage.correlate(A, k, mode='constant'), 1)

给予：

>>> A
array([[[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 0, 1],
        [0, 0, 0, 0]],

       [[0, 1, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 1, 0, 0]],

       [[1, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 1, 1, 0]]])

>>> B
array([[[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 0, 1]],

       [[0, 0, 1, 0],
        [1, 1, 1, 0],
        [0, 1, 0, 1],
        [0, 0, 1, 0]],

       [[0, 1, 0, 0],
        [1, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 1, 1, 1]]])

完成！

---

验证上述代码正确性的一个简单方法是在这样的矩阵上运行它：

A = np.zeros((4, 4, 4), int)
A[0, 1, 2] = 1

这会产生预期的结果：

>>> A
array([[[0, 0, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]]])

>>> B
array([[[0, 0, 0, 0],
        [0, 0, 0, 1],
        [0, 0, 1, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]]])