оптимизация векторизации NumPy при расчете расстояний и np.sum

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

Мы бы сохраняли циклы, но пытались оптимизировать, предварительно вычисляя определенные вещи -

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)

Answer 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

.