Vectorize iterative addition in NumPy arrays

Question

For each element in a randomized array of 2D indices (with potential duplicates), I want to "+=1" to the corresponding grid in a 2D zero array. However, I don't know how to optimize the computation. Using the standard for loop, as shown here,

def interadd():
    U = 100 
    input = np.random.random(size=(5000,2)) * U
    idx = np.floor(input).astype(np.int) 

    grids = np.zeros((U,U))      
    for i in range(len(input)):
        grids[idx[i,0],idx[i,1]] += 1
    return grids

The runtime can be quite significant:

>> timeit(interadd, number=5000)
43.69953393936157

Is there a way to vectorize this iterative process?

Answer 1

You could speed it up a little by using np.add.at, which correctly handles the case of duplicate indices:

def interadd(U, idx):
    grids = np.zeros((U,U))      
    for i in range(len(idx)):
        grids[idx[i,0],idx[i,1]] += 1
    return grids

def interadd2(U, idx):
    grids = np.zeros((U,U))
    np.add.at(grids, idx.T.tolist(), 1)
    return grids

def interadd3(U, idx):
    # YXD suggestion
    grids = np.zeros((U,U))
    np.add.at(grids, (idx[:,0], idx[:,1]), 1)
    return grids

which gives

>>> U = 100
>>> idx = np.floor(np.random.random(size=(5000,2))*U).astype(np.int)
>>> (interadd(U, idx) == interadd2(U, idx)).all()
True
>>> %timeit interadd(U, idx)
100 loops, best of 3: 8.48 ms per loop
>>> %timeit interadd2(U, idx)
100 loops, best of 3: 2.62 ms per loop

---

And YXD's suggestion:

>>> (interadd(U, idx) == interadd3(U, idx)).all()
True
>>> %timeit interadd3(U, idx)
1000 loops, best of 3: 1.09 ms per loop

Answer 2

Divakar's answer lead me to try the following, which looks to be the fastest method yet:

lin_idx = idx[:,0]*U + idx[:,1]
grids = np.bincount(lin_idx, minlength=U**2).reshape(U, U)

Timings:

In [184]: U = 100 
     ...: input = np.random.random(size=(5000,2)) * U
     ...: idx = np.floor(input).astype(np.int)

In [185]: %timeit interadd3(U, idx)  # By DSM / XYD
1000 loops, best of 3: 1.68 ms per loop

In [186]: %timeit unique_counts(U, idx)  # By Divakar
1000 loops, best of 3: 676 µs per loop

In [187]: %%timeit
     ...: lin_idx = idx[:,0]*U + idx[:,1]
     ...: grids = np.bincount(lin_idx, minlength=U*U).reshape(U, U)
     ...: 
10000 loops, best of 3: 97.5 µs per loop

Answer 3

You can convert R,C indices from idx to linear indices, then find out the unique ones alongwith their counts and finally store them in the output grids as the final output. Here's the implementation to achieve the same -

# Calculate linear indices corressponding to idx
lin_idx = idx[:,0]*U + idx[:,1]

# Get unique linear indices and their counts
unq_lin_idx,idx_counts = np.unique(lin_idx,return_counts=True)

# Setup output array and store index counts in raveled/flattened version
grids = np.zeros((U,U))  
grids.ravel()[unq_lin_idx] = idx_counts

Runtime tests -

Here are the runtime tests covering all the approaches (including @DSM's approaches) and using the same definitions as listed in that solution -

In [63]: U = 100
    ...: idx = np.floor(np.random.random(size=(5000,2))*U).astype(np.int)
    ...: 

In [64]: %timeit interadd(U, idx)
100 loops, best of 3: 7.57 ms per loop

In [65]: %timeit interadd2(U, idx)
100 loops, best of 3: 2.59 ms per loop

In [66]: %timeit interadd3(U, idx)
1000 loops, best of 3: 1.24 ms per loop

In [67]: def unique_counts(U, idx):
    ...:     lin_idx = idx[:,0]*U + idx[:,1]
    ...:     unq_lin_idx,idx_counts = np.unique(lin_idx,return_counts=True)
    ...:     grids = np.zeros((U,U))  
    ...:     grids.ravel()[unq_lin_idx] = idx_counts
    ...:     return grids
    ...: 

In [68]: %timeit unique_counts(U, idx)
1000 loops, best of 3: 595 µs per loop

Runtimes suggest that the proposed np.unique based approach is more than 50% faster than the second fastest approach.