How to covert replace numpy loop with faster broadcast operation

Viewed 140

I am trying to use broadcasting to speed up my numpy code. the real code has much larger arrays and loops through multiple times, but I think this snippet illustrates the issue.

import numpy as np
row    = np.array([0,0,1,1,4])
dl_ddk = np.array([0,8,29,112,11])
change1 = np.zeros(5)
change2 = np.zeros(5)
for k in range(0, row.shape[0]):
   i          = row[k]
   change1[i] += dl_ddk[k]
change2[row] += dl_ddk
print(change1)
print(change2)

change1 = [8, 141, 0, 0 11] change2 = [8, 112, 0, 0 11]

I thought these two change arrays would be equals however, it seems that the broadcast operations += is overwriting rather than adding values. Is there a way to vectorize a loop in np with matrix referencing like this that will give the same results as change1?

2 Answers

You can use np.bincount() and use dl_ddk as the weights:

import numpy as np

row    = np.array([0,0,1,1,4])
dl_ddk = np.array([0,8,29,112,11])

change1 = np.bincount(row, weights=dl_ddk)
print(change1)
# [  8. 141.   0.   0.  11.]

The bit in the docs show using it in a way almost exactly like your problem:

If weights is specified the input array is weighted by it, i.e. if a value n is found at position i, out[n] += weight[i] instead of out[n] += 1.

In [1]: row    = np.array([0,0,1,1,4]) 
   ...: dl_ddk = np.array([0,8,29,112,11]) 
   ...: change1 = np.zeros(5) 
   ...: change2 = np.zeros(5) 
   ...: for k in range(0, row.shape[0]): 
   ...:    i          = row[k] 
   ...:    change1[i] += dl_ddk[k] 
   ...: change2[row] += dl_ddk                                                                       

change2 does not match because of buffering. ufunc has added a at method to address this:

Performs unbuffered in place operation on operand 'a' for elements specified by 'indices'.           
                                                           
In [3]: change3 = np.zeros(5)                                                                        
In [4]: np.add.at(change3, row, dl_ddk)                                                              
In [5]: change1                                                                                      
Out[5]: array([  8., 141.,   0.,   0.,  11.])
In [6]: change2                                                                                      
Out[6]: array([  8., 112.,   0.,   0.,  11.])
In [7]: change3                                                                                      
Out[7]: array([  8., 141.,   0.,   0.,  11.])
Related