在Python中对网格进行矢量化计算

Question

我有矩阵： f = np.zeros((M, N))充满了一些浮点数。 我想用其邻居的平均值替换每个内部点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

只需使用卷积即可 。 有多种开源实现，但是scipy的作品：

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