Векторизованная версия for-loop + numpy.where

Question

Фоновая информация: У меня есть большое количество (N) частиц в 3D. Для всех частиц пары [i, j], которые имеют определенные свойства, я вычисляю геометрический коэффициент c [i, j]. Затем я хочу суммировать вклад всех пар [i, j] для фиксированного i и называть это c [i] (и повторить эту процедуру для всех частиц i).

Как правило, количество соответствующих пар намного меньше, чем N^2, поэтому наличие (N, N) -мерного массива C с соответствующей информацией в позициях [i, j] и множестве нулей в другом месте справедливо быстрый с numpy, но также очень неэффективный с точки зрения использования памяти. Итак, теперь я просто храню C [i, j] для соответствующих пар и частиц, образующих пары в 1D массивах.

Это, вероятно, лучше всего иллюстрируется в примере: Скажем, у меня две пары, состоящие из частиц (3,5) и (3,10). Схематически мой переменные выглядеть следующим образом (двойной счет, предназначенный):

i = [3,3,5,10] #list of particles i that form a pair 
j = [5,10,3,3] #corresponding particles j (not used in the later example) 
cij = [c35,c310,-c35,-c310] #(with actual numbers in reality) 

Теперь это действительно сводится к нахождению эффективного векторизованному способа переписать следующий цикл:

part_list=np.arange(N) 
for a in part_list: 
    cond = np.where(i == a) 
    ci[a] = np.sum(cij[cond]) 

Других решения I подумал, но хотел бы избежать:

a) Параллелизировать цикл for: Невозможно b/c, это встроено во внешний цикл, который я хочу распараллелить в какой-то момент.

b) Напишите for-loop в C и заверните его в Python: Кажется, это излишняя проблема для этой (надеюсь) довольно простой проблемы.

Answer 1

Вы можете получить то, что хотите, с np.bincount. Если частицы, где последовательно пронумерованных от 0 до, вы можете просто сделать:

ci = np.bincount(i, weights= cij) 

Чтобы увидеть, что это делает:

>>> i = [3, 3, 5, 10] 
>>> j = [5, 10, 3, 3] 
>>> cij = [0.1, 0.2, -0.1, -0.2] 
>>> np.bincount(i, weights= cij) 
array([ 0. , 0. , 0. , 0.3, 0. , -0.1, 0. , 0. , 0. , 0. , -0.2]) 

Если вы не хотите, чтобы все эти дополнительные нули, вы могли бы сделать что-то как:

>>> unq_i, inv_i = np.unique(i, return_inverse=True) 
>>> unq_ci = np.bincount(inv_i, weights=cij) 
>>> unq_i 
array([ 3, 5, 10]) 
>>> unq_ci 
array([ 0.3, -0.1, -0.2]) 

И вы можете позже присвоить эти уникальные ценности, делая что-то вроде:

ci[unq_i] = unq_ci