python - 使用 Numpy 向量化矩阵中最近邻居的和

Question

给定一个 N 乘 N 矩阵 A，我想构造另一个 N 乘 N 矩阵 B，其条目由下式给出:

B_{i,j} = A_{i+1, j} + A_{i-1, j} + A_{i, j+1} + A_{i, j-1}

注意系统使用周期性边界条件所以

A_{N, j} = A_{0, j}

A_{i, -1} = A{i, N-1}

等等。

如何在 numpy 中将其作为矢量化操作来执行？

Answer 1

您描述的操作是 convolution .你可以使用 scipy.ndimage.convolve :

import numpy as np
from scipy.ndimage import convolve, generate_binary_structure

gen = np.random.RandomState(0)
A = gen.random_integers(0, 3, (6, 6))   # input array
k = generate_binary_structure(2, 1)     # kernel
B = convolve(A, k, mode='wrap')         # use mode='wrap' for periodic boundaries

print(A)
# [[0 3 1 0 3 3]
#  [3 3 1 3 1 2]
#  [0 3 2 0 0 0]
#  [2 1 2 3 3 2]
#  [0 1 1 1 1 0]
#  [1 0 3 0 3 1]]

print(k)
# [[False  True False]
#  [ True  True  True]
#  [False  True False]]

print(B)
# [[10  7  8  7 10  9]
#  [ 8 13 10  5  9  9]
#  [ 8  9  8  8  4  4]
#  [ 5  9  9  9  9  7]
#  [ 4  3  8  6  8  4]
#  [ 2  8  5  7  8  8]]