在计算距离和 np.sum 时优化 numpy 向量化

Question

我有以下代码：

# positions: np.ndarray of shape(N,d) 
# fitness: np.ndarray of shape(N,)
# mass: np.ndarray of shape(N,)

iteration = 1
while iteration <= maxiter:
    K = round((iteration-maxiter)*(N-1)/(1-maxiter) + 1)

    for i in range(N):
        displacement = positions[:K]-positions[i]
        dist = np.linalg.norm(displacement, axis=-1)
        if i<K:
            dist[i] = 1.0       # prevent 1/0

        force_i = (mass[:K]/dist)[:,np.newaxis]*displacement
        rand = np.random.rand(K,1)
        force[i] = np.sum(np.multiply(rand,force_i), axis=0)

所以我有一个数组存储 N 个粒子在 d 维度中的坐标。我需要先计算粒子 i 和第一个 K 粒子之间的欧氏距离，然后计算每个 K 粒子的 'force'。然后，我需要对 K 个粒子求和以找到作用在粒子 i 上的总力，并对所有 N 个粒子重复。这只是部分代码，但经过一些分析后，这部分是时间最关键的步骤。

所以我的问题是如何优化上面的代码。我已经尝试尽可能地对其进行矢量化，我不确定是否还有改进的余地。分析结果显示 {method 'reduce' of 'numpy.ufunc' objects}、fromnumeric.py:1778(sum) 和 linalg.py:2103(norm) 花费最长的时间到运行。第一个死于阵列广播？如何优化这三个函数调用？

Answer 1

由于您的代码缺少一些部分，我不得不进行一些调整。但第一个优化是摆脱 for i in range(N) 循环：

import numpy as np

np.random.seed(42)

N = 10
d = 3
maxiter = 50

positions = np.random.random((N, d))
force = np.random.random((N, d))
fitness = np.random.random(N)
mass = np.random.random(N)

iteration = 1
while iteration <= maxiter:
    K = round((iteration-maxiter)*(N-1)/(1-maxiter) + 1)

    displacement = positions[:K, None]-positions[None, :]
    dist = np.linalg.norm(displacement, axis=-1)
    dist[dist == 0] = 1

    force = np.sum((mass[:K, None, None]/dist[:,:,None])*displacement * np.random.rand(K,N,1), axis=0)
    iteration += 1

其他改进是尝试更快地实施规范，例如 scipy.cdist 或 numpy.einsum

Answer 2

我们会保留循环，但尝试通过预先计算某些东西来优化 -

from scipy.spatial.distance import cdist

iteration = 1
while iteration <= maxiter:
    K = round((iteration-maxiter)*(N-1)/(1-maxiter) + 1)

    posd = cdist(positions,positions)
    np.fill_diagonal(posd,1)
    rands = np.random.rand(N,K)
    s = rands*(mass[:K]/posd[:,:K])
    for i in range(N):
        displacement = positions[:K]-positions[i]
        force[i] = s[i].dot(displacement)

在计算距离和 np.sum 时优化 numpy 向量化

optimizing numpy vectorization on calculating distances and np.sum

python

numpy

vectorization

array-broadcasting