Can one vectorize the summing of a large data set onto record-specific elements of an array?

user_1895996 0
Viewed 32

I have a very large data set (order billions of records) that I need to sum on a 2D array. For each value, there are indices specifying which element of the array the value should be added to:

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]

with the result

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

My question: Is there a way to vectorize the above operation so as to eliminate the for-loop?

vectorization numpy arrays python
2 Answers

Your index arrays lend themselves well to fancy indexing. In the simplest case, you can do

A[I, J] += X

If you have any duplicate indices, i.e., you want to increment some location in A more than once, a more robust approach would be

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

Here's a fully-vectorized solution:

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

output:

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