numpy: efficiently add rows of a matrix
<p>I have a matrix.  </p>

<pre><code>mat = array([
   [ 0,  1,  2,  3],
   [ 4,  5,  6,  7],
   [ 8,  9, 10, 11]
   ])
</code></pre>

<p>I'd like to get the sum of the rows at certain indices: eg.</p>

<pre><code>ixs = np.array([0,2,0,0,0,1,1])
</code></pre>

<p>I know I can compute the answer as:</p>

<pre><code>mat[ixs].sum(axis=0)
&gt; array([16, 23, 30, 37])
</code></pre>

<p>The problem is ixs may be very long, and I don't want to use all the memory to create the intermediate product mat[ixs], only to reduce it again with the sum.</p>

<p>I also know I could simply count up the indices and use multiplication instead. </p>

<pre><code>np.bincount(ixs, minlength=mat.shape[0).dot(mat)
&gt; array([16, 23, 30, 37])
</code></pre>

<p>But that will be expensive if my ixs are sparse. </p>

<p>I know about scipy's sparse matrices, and I suppose I could use them, but I'd prefer a pure numpy solution as sparse matrices are limited in various ways (such as only being 2-d)</p>

<p>So, is there a pure numpy way to merge the indexing and sum-reduction in this case?</p>

<h1>Conclusions:</h1>

<p>Thanks you Divakar and hpaulj for your very thorough responses.  By "sparse" I meant that most of the values in <code>range(w.shape[0])</code> are not in ixs.  Using that new definition (and with more realisitic data size, I re-ran Divakar tests, with some new funcitona dded :</p>

<pre><code>rng = np.random.RandomState(1234)
mat = rng.randn(1000, 500)
ixs = rng.choice(rng.randint(mat.shape[0], size=mat.shape[0]/10), size=1000)

# Divakar's solutions
In[42]: %timeit org_indexing_app(mat, ixs)
1000 loops, best of 3: 1.82 ms per loop
In[43]: %timeit org_bincount_app(mat, ixs)
The slowest run took 4.07 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 177 Âµs per loop
In[44]: %timeit indexing_modified_app(mat, ixs)
1000 loops, best of 3: 1.81 ms per loop
In[45]: %timeit bincount_modified_app(mat, ixs)
1000 loops, best of 3: 258 Âµs per loop
In[46]: %timeit simply_indexing_app(mat, ixs)
1000 loops, best of 3: 1.86 ms per loop
In[47]: %timeit take_app(mat, ixs)
1000 loops, best of 3: 1.82 ms per loop
In[48]: %timeit unq_mask_einsum_app(mat, ixs)
10 loops, best of 3: 58.2 ms per loop 
# hpaulj's solutions
In[53]: %timeit hpauljs_sparse_solution(mat, ixs)
The slowest run took 9.34 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 524 Âµs per loop
%timeit hpauljs_second_sparse_solution(mat, ixs)
100 loops, best of 3: 9.91 ms per loop
# Sparse version of original bincount solution (see below):
In[60]: %timeit sparse_bincount(mat, ixs)
10000 loops, best of 3: 71.7 Âµs per loop
</code></pre>

<p><strong>The winner</strong> in this case is the sparse version of the bincount solution.</p>

<pre><code>def sparse_bincount(mat, ixs):
    x = np.bincount(ixs)
    nonzeros, = np.nonzero(x)
    x[nonzeros].dot(mat[nonzeros])
</code></pre>

<p>Since we are assuming that <code>ixs</code> could be <em>sparsey</em>, we could modify the strategy to get the summations of rows from the <code>zero-th</code> row and rest of the rows separately based on the given row indices. So, we could use the <code>bincount</code> method for the <code>non-zero-th</code> indexed rows summation and add it with the <code>(zero-th row x no. of zeros</code> in <code>ixs</code>).</p>

<p>Thus, the second approach could be modified, like so -</p>

<pre><code>nzmask = ixs!=0
nzsum = np.bincount(ixs[nzmask]-1, minlength=mat.shape[0]-1).dot(mat[1:])
row0_sum = mat[0]*(len(ixs) - np.count_nonzero(nzmask))
out = nzsum + row0_sum
</code></pre>

<p>We could extend this strategy to the first approach as well, like so -</p>

<pre><code>out = mat[0]*(len(ixs) - len(nzidx)) + mat[ixs[nzidx]].sum(axis=0)
</code></pre>

<p>If we are working with lots of non-zero indices that are repeated, we could alternatively make use of <code>np.take</code> with focus on performance. Thus, <code>mat[ixs[nzidx]]</code> could be replaced by <code>np.take(mat,ixs[nzidx],axis=0)</code> and similarly <code>mat[ixs]</code> by <code>np.take(mat,ixs,axis=0)</code>. With such repeated indices based indexing <code>np.take</code> brings out some noticeable speedup as compared to simply indexing.</p>

<p>Finally, we could use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html" rel="nofollow"><code>np.einsum</code></a> to perform these row ID based selection and summing, like so -</p>

<pre><code>nzmask = ixs!=0
unq,tags = np.unique(ixs[nzmask],return_inverse=1)
nzsum = np.einsum('ji,jk-&gt;k',np.arange(len(unq))[:,None] == tags,mat[unq])
out = mat[0]*(len(ixs) - np.count_nonzero(nzmask)) + nzsum
</code></pre>

<h2>Benchmarking</h2>

<p>Let's list out all the five approaches posted thus far in this post and also include the two approaches posted in the question for some runtime testing as functions -</p>

<pre><code>def org_indexing_app(mat,ixs):
    return mat[ixs].sum(axis=0)

def org_bincount_app(mat,ixs):
    return np.bincount(ixs, minlength=mat.shape[0]).dot(mat)

def indexing_modified_app(mat,ixs):
    return np.take(mat,ixs,axis=0).sum(axis=0)

def bincount_modified_app(mat,ixs):
    nzmask = ixs!=0
    nzsum = np.bincount(ixs[nzmask]-1, minlength=mat.shape[0]-1).dot(mat[1:])
    row0_sum = mat[0]*(len(ixs) - np.count_nonzero(nzmask))
    return nzsum + row0_sum

def simply_indexing_app(mat,ixs):
    nzmask = ixs!=0
    nzsum = mat[ixs[nzmask]].sum(axis=0)
    return mat[0]*(len(ixs) - np.count_nonzero(nzmask)) + nzsum

def take_app(mat,ixs):
    nzmask = ixs!=0
    nzsum = np.take(mat,ixs[nzmask],axis=0).sum(axis=0)
    return mat[0]*(len(ixs) - np.count_nonzero(nzmask)) + nzsum

def unq_mask_einsum_app(mat,ixs):
    nzmask = ixs!=0
    unq,tags = np.unique(ixs[nzmask],return_inverse=1)
    nzsum = np.einsum('ji,jk-&gt;k',np.arange(len(unq))[:,None] == tags,mat[unq])
    return mat[0]*(len(ixs) - np.count_nonzero(nzmask)) + nzsum
</code></pre>

<p><strong>Timings</strong></p>

<p>Case #1 (<code>ixs</code> is 95% sparsey) :</p>

<pre><code>In [301]: # Setup input
     ...: mat = np.random.rand(20,4)
     ...: ixs = np.random.randint(0,10,(100000))
     ...: ixs[np.random.rand(ixs.size)&lt;0.95] = 0 # Make it approx 95% sparsey
     ...: 

In [302]: # Timings
     ...: %timeit org_indexing_app(mat,ixs)
     ...: %timeit org_bincount_app(mat,ixs)
     ...: %timeit indexing_modified_app(mat,ixs)
     ...: %timeit bincount_modified_app(mat,ixs)
     ...: %timeit simply_indexing_app(mat,ixs)
     ...: %timeit take_app(mat,ixs)
     ...: %timeit unq_mask_einsum_app(mat,ixs)
     ...: 
100 loops, best of 3: 4.89 ms per loop
1000 loops, best of 3: 428 Âµs per loop
100 loops, best of 3: 3.29 ms per loop
1000 loops, best of 3: 329 Âµs per loop
1000 loops, best of 3: 537 Âµs per loop
1000 loops, best of 3: 462 Âµs per loop
1000 loops, best of 3: 1.07 ms per loop
</code></pre>

<p>Case #2 (<code>ixs</code> is 98% sparsey) :</p>

<pre><code>In [303]: # Setup input
     ...: mat = np.random.rand(20,4)
     ...: ixs = np.random.randint(0,10,(100000))
     ...: ixs[np.random.rand(ixs.size)&lt;0.98] = 0 # Make it approx 98% sparsey
     ...: 

In [304]: # Timings
     ...: %timeit org_indexing_app(mat,ixs)
     ...: %timeit org_bincount_app(mat,ixs)
     ...: %timeit indexing_modified_app(mat,ixs)
     ...: %timeit bincount_modified_app(mat,ixs)
     ...: %timeit simply_indexing_app(mat,ixs)
     ...: %timeit take_app(mat,ixs)
     ...: %timeit unq_mask_einsum_app(mat,ixs)
     ...: 
100 loops, best of 3: 4.86 ms per loop
1000 loops, best of 3: 438 Âµs per loop
100 loops, best of 3: 3.5 ms per loop
1000 loops, best of 3: 260 Âµs per loop
1000 loops, best of 3: 318 Âµs per loop
1000 loops, best of 3: 288 Âµs per loop
1000 loops, best of 3: 694 Âµs per loop
</code></pre>

python numpy indexing