python - Python 中网格的矢量化计算

Question

我有矩阵: f = np.zeros((M, N))填充一些 float 。我想替换每个内部点 f[i, j]与其邻居的平均值: f_new[i,j] = (f[i-1, j] + f[i+1, j] + f[i, j+1] + f[i, j-1])/4 。有一个明显的方法可以做到这一点:

f_new = f.copy()
for i in range(1, M-1):
    for j in range(1, N-1):
       f_new[i,j] = (f[i-1, j] + f[i+1, j] + f[i, j+1] + f[i, j-1])/4
f = f_new

在 Python 中是否有更优雅(矢量化)的方法来执行此操作？谢谢。

Answer 1

只需使用 convolution 。有多种开源实现，但是 scipy's作品:

import numpy as np
import scipy.signal as signal
import time

f = np.random.random((1000, 1000))

# -------- For loop
start = time.perf_counter()
f_pad = np.pad(f, ((1, 1), (1, 1)), 'constant', constant_values=0) # pad the array

f_new = f_pad.copy()
for i in range(1, f.shape[0]+1):
    for j in range(1, f.shape[1]+1):
        f_new[i,j] = (f_pad[i-1, j] + f_pad[i+1, j] + f_pad[i, j+1] + f_pad[i, j-1])/4.
f_new = f_new[1:-1, 1:-1]
print("For loop time:", str(time.perf_counter() - start))

# -------- Convolution
start = time.perf_counter()
f_newer = signal.convolve2d(f, 
                            np.array([[0, 0.25, 0],
                                      [0.25, 0, 0.25],
                                      [0, 0.25, 0]]),
                            mode='same',
                            boundary='fill',
                            fillvalue=0)
print("Convolution time:", str(time.perf_counter() - start))

# >> For loop time: 1.4474979060469195
# >> Convolution time: 0.0972418460296467