代码之家 › 专栏 › 技术社区 › Physicist

在距离和距离计算中优化numpy矢量化np.总和

numpy-broadcasting vectorization numpy python

0

Physicist · 技术社区 · 6 年前

我有以下代码:

# 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公司 质点求作用在质点上的总力 我 粒子。这只是代码的一部分,但在分析了一些代码之后,这部分是最关键的一步。

所以我的问题是如何优化上面的代码。我已经尽可能地把它矢量化了,我不确定是否还有改进的余地。分析结果显示 {method 'reduce' of 'numpy.ufunc' objects} , fromnumeric.py:1778(sum) 和 linalg.py:2103(norm) 用最长的时间跑。第一个死的是阵列广播吗?如何优化这三个函数调用?

2 回复 | 直到 6 年前

1

Divakar 6 年前

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

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)

2

1

user2653663 6 年前

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