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粒子而产生的“力”。然后，我需要对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矢量化

2 个答案: