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

时间:2019-02-21 17:18:48

标签: python vectorization

我有矩阵: 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中是否有更优雅的(矢量化)方法?谢谢。

1 个答案:

答案 0 :(得分:0)

只需使用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