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

Posted 2023-04-15

【中文标题】在计算距离和 np.sum 时优化 numpy 矢量化

【英文标题】：optimizing numpy vectorization on calculating distances and np.sum

【发布时间】：2018-09-10 13:31:29

【问题描述】：

我有以下代码：

# 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 粒子的“力”。然后，我需要对K 粒子求和以找到作用在粒子i 上的总力，并对所有N 粒子重复。它只是代码的一部分，但经过一些分析后，这部分是最关键的步骤。

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

【问题讨论】：

这个answer 很有趣。 scipy中还有cdist函数

N和d的典型值是多少？

N 通常在 50-100 之间，d 在 2-50 左右，但如果即使 N 达到 ~2000 并且 @ 987654341@ 达到 ~100。

【参考方案1】：

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

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)

【讨论】：

应该是s = rands*(mass[:K]/posd[:,:K])吗？通过这种更改，它会给出相同的结果并且确实更快。谢谢。

【参考方案2】：

由于您的代码缺少一些部分，我不得不进行一些调整。但是第一个优化是去掉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