用于半稀疏数组中最近邻计算的 Numpy 矢量化

Question

我有一个关于 numpy 向量化的问题。

我有一个半稀疏矩阵，其中非零元素是 +1 或 –1，并且在 numpy 数组（我们称之为 lattice）中随机排序，因此 lattice[i,j] 处的“站点”具有lattice[i+1,j+1]、lattice[i+1,j-1]、lattice[i-1,j+1] 和 lattice[i-1,j-1] 处的最近邻：

lattice = np.array([[1,  0, -1,  0],
                    [0,  1,  0, -1],
                    [1,  0,  1,  0],
                    [0,  1,  0, -1]]) 

有一个函数 F 将存储在 lattice 中的配置映射到一个标量值。不知道如何在此处显示 mathjax / latex，但我会尽力输入。

F(lattice) = a * sum( lattice[i,j] ) + b * sum( lattice[i,j] * lattice[i±1, j±1] )

第一个表达式是按 a 缩放的所有格点的总和，第二个表达式是所有格点及其与每个站点的最近邻集的乘积，按 b 缩放。

为了简化所有这些，当我生成我的格子时，我创建了一个字典，其中每个键对应于格子中每个站点的索引元组，并且存储在那里的值是指向的元组列表最近的邻居，即

sites = { (i,j) : [ (i+1, j+1), (i+1, j-1), (i-1,j+1), (i-1,j-1) ]}

我绝对可以利用矢量化来计算第一个总和，但我很难找到矢量化第二个总和的方法，因为最近邻居的总和取决于所考虑的格点的特定索引集.

def F(a,b, lattice, sites):
    """
    Parameters
    ----------
    a: float
    b: float 
    lattice: np.ndarray, shape [nx, ny]
    sites: dict
         Keys are tuples of the lattice site indices, returns a list 
         of tuples of the site's nearest neighbors

    Returns
    -------
    c: float
    """

    #-- Compute the first sum
    c = a * np.sum( (np.ma.masked_equal(lattice,0) + 1) / 2. )

    #-- Compute the second sum
    c += np.array([ [ b*lattice[i,j]*lattice[x,y] for (x,y) in sites[(i,j)] ] for (i,j) in sites.keys() ]).sum()

    return c

这提供了 ok 性能，但是当我考虑大格时，使用列表推导计算第二个总和确实开始影响我的性能。

你们知道我如何对第二个和进行矢量化吗？

我正在考虑以某种方式使用 numpy 数组的步幅来计算这一点，但如果 (a) 可行且 (b) 值得的话，我不是。

Answer 1

如果你像这样创建一个辅助数组：

lattice_neighbor_sum = np.zeros_like(lattice)
lattice_neighbor_sum[:, :-1] += lattice[:, 1:]
lattice_neighbor_sum[:, 1:] += lattice[:, :-1]
lattice_neighbor_sum[:-1] += lattice[1:]
lattice_neighbor_sum[1:] += lattice[:-1]

您的示例输入的结果将是：

>>> lattice_neighbor_sum
array([[ 0,  1,  0, -1,  0],
       [ 3,  0,  0,  0, -1],
       [ 0,  4,  0, -2,  0],
       [ 2,  0,  1,  0, -2]])

稍微应用分配属性，您的计算可以重写为：

a * np.sum(lattice) + b * np.sum(lattice * lattice_neighbor_sum)

Answer 2

@Jaime，我尝试了您提供的答案，并提出了这种稍微修改过的方法。另外，我很抱歉，但我在最初的晶格定义中犯了一个错误（假设是偶数行和列），所以我在这里使用的与我提供给您的略有不同。如果我们采用以下格：

lattice = np.array([[1,  0, -1,  0],
                    [0,  1,  0, -1],
                    [1,  0,  1,  0],
                    [0,  1,  0, -1]])

并以稍微不同的方式创建辅助数组：

neigh_sum = np.zeros_like(lattice) 
#-- neighbors at (i-1, j+1)
neigh_sum += np.pad(lattice, ((1,0),(1,0)), 'constant')[:-1,:-1]
#-- neighbors at (i+1, j+1)
neigh_sum += np.pad(lattice, ((1,0),(0,1)), 'constant')[:-1,1:]
#-- neighbors at (i+1, j-1)
neigh_sum += np.pad(lattice, ((0,1),(0,1)), 'constant')[1:,1:]
#-- neighbors at (i-1,j-1)
neigh_sum += np.pad(lattice, ((0,1),(1,0)), 'constant')[1:,:-1]

使用这种方法，我得到：

neigh_sum = array([[ 1,  0,  0,  0],
                   [ 0,  2,  0,  0],
                   [ 2,  0,  0,  0],
                   [ 0,  2,  0,  1]])

然后我可以使用您建议的行来计算总和：

c = a * np.sum(lattice) + b * np.sum(lattice * neigh_sum)