如何用numpy中的矢量化替换for循环？

Question

这是我的函数，我知道 for 循环很慢 在 numpy 中必须有更快/更好的方法来做到这一点

import numpy as np

def heat(u):
    my_u = np.empty_like(u)
    np.copyto(my_u, u)
    for row in range(1,u.shape[0]-1):
        for col in range(1,u.shape[1]-1):
            my_u[row, col] = (u[row-1][col] + u[row+1][col] 
            + u[row][col-1] + u[row][col+1])/4
    return my_u

print(heat(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]])))

Answer 1

您确定您只想添加作为十字架的 4 个元素，而不是所有邻居吗？在 SO 上使用 'sum neighbors numpy' 的搜索可以通过您如何掩盖它，为您提供许多答案，无论是交叉还是完整的邻居集。

如果您想稍微提高速度，请放弃数组的双重查找，这样会快 30% 左右。

my_u[row, col] = (u[row-1, col] + u[row+1, col] 
                 + u[row, col-1] + u[row, col+1])/4