python - 如何使用 numpy 使 for 循环更快

Question

如果输入一个 numpy 二维数组，例如

                           100   100   100   100   100
                           100    0     0     0    100
                           100    0     0     0    100
                           100    0     0     0    100
                           100   100   100   100   100

应该得到这样的输出

                            100   100   100   100   100
                            100    50    25    50   100
                            100    25     0    25   100
                            100    50    25    50   100
                            100   100   100   100   100

其中除边框外的每个数字都成为其相邻数字的平均值。

我当前的代码有效，但我需要在没有 for 循环的情况下使用它，并使用 numpy 对其进行矢量化。

我当前的代码:

import numpy as np
def evolve_heat_slow(u):      
    u2 = np.copy(u)
    x=u2.shape[0]
    y=u2.shape[1]
    for i in range(1,x-1):
        for s in range(1,y-1):
            u2[i,s]=(u[i-1,s]+u[i+1,s]+u[i,s+1]+u[i,s-1])/4
    return u2

Answer 1

这几乎就是二维卷积的定义。 scipy 为您提供帮助。我复制 a 以保留边框； valid 模式下的卷积将生成一个较小的数组(没有边框)，然后我将其粘贴到准备好的“框架”中。

import numpy as np
from scipy.signal import convolve2d

a = np.array([[100, 100, 100, 100, 100], [100, 0, 0, 0, 100],  [100, 0, 0, 0, 100], [100, 0, 0, 0, 100], [100, 100, 100, 100, 100]])
b = np.array([[0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0]])
r = np.copy(a)
r[1:-1, 1:-1] = convolve2d(a, b, mode='valid')
r
# => array([[100, 100, 100, 100, 100],
#           [100,  50,  25,  50, 100],
#           [100,  25,   0,  25, 100],
#           [100,  50,  25,  50, 100],
#           [100, 100, 100, 100, 100]])