algorithm - 用于半稀疏数组中最近邻计算的 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 将存储在晶格中的配置映射到一个标量值。不确定如何在此处显示 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)