Можно ли векторизовать суммирование большого набора данных на c элементов массива, специфичных для записи?

Question

У меня очень большой набор данных (порядка миллиардов записей), который мне нужно суммировать в 2D-массиве. Для каждого значения есть индексы, указывающие, к какому элементу массива следует добавить значение:

import numpy as np

I = [0, 2, 1, 2, 1]
J = [1, 2, 1, 2, 1]
X = [2., 5., 0., 6., 4.]

A = np.zeros((3,3), dtype = 'f')

for i in range(len(I)) :
    A[I[i], J[i]] += X[i]

с результатом

> print(A)
[[ 0.  2.  0.]
 [ 0.  4.  0.]
 [ 0.  0. 11.]]

Мой вопрос: есть ли способ векторизовать вышеуказанная операция, чтобы устранить for-l oop?

Answer 1

Ваши индексные массивы хорошо подходят для причудливой индексации. В простейшем случае вы можете сделать

A[I, J] += X

. Если у вас есть какие-либо повторяющиеся индексы, т.е. вы хотите увеличить какое-то место в A более одного раза, более надежным подходом будет

np.add.at(A, (I, J), X)

Answer 2

Вот полностью векторизованное решение:

X = [\
     [(0,1), 2. ], \
     [(2,2), 5. ], \
     [(1,1), 0. ], \
     [(2,2), 6. ], \
     [(1,1), 4. ]  \
     ]
# create a dataframe with x, y, and val. I'm not doing it very. 
# efficiently here - but since you control the data structure 
# you can just start from this kind of dataframe. 
records = [(r[0], r[1], t) for r,t in X]
df = pd.DataFrame.from_records(records, columns=["x", "y", "val"])

A = np.zeros((3,3), dtype = 'float64')

df = df.groupby(["x", "y"], as_index = False).sum()
A[df.x, df.y] = df.val

вывод:

array([[ 0.,  2.,  0.],
       [ 0.,  4.,  0.],
       [ 0.,  0., 11.]])