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

У меня есть следующий код:

# 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) выполняются дольше всего. Первый умрет от массивного вещания? Как я могу оптимизировать эти три вызова функций?


python numpy vectorization array-broadcasting
person Physicist    schedule 10.09.2018    source источник
comment
Этот ответ весьма интересен. Существует также функция cdist. в scipy   -  person xdze2    schedule 10.09.2018
comment
Каковы типичные значения N и d?   -  person Divakar    schedule 10.09.2018
comment
N обычно составляет от 50 до 100, а d — от 2 до 50, но было бы здорово, если бы оптимизацию можно было применить, даже если N достигает ~2000, а d достигает ~100.   -  person Physicist    schedule 11.09.2018

Ответы (2)


arrow_upward
1
arrow_downward

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

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)
person Divakar    schedule 10.09.2018
comment
должно быть s = rands*(mass[:K]/posd[:,:K])? С этим изменением он дает те же результаты и действительно быстрее. Спасибо. - person Physicist; 11.09.2018

arrow_upward
1
arrow_downward

Мне пришлось внести некоторые коррективы, так как в вашем коде отсутствовало несколько частей. Но первой оптимизацией будет избавление от цикла 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.

person user2653663    schedule 10.09.2018