Векторизация вычислений на сетке в 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'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