Speed up numpy summation between big arrays
<p>I have a code running operations on <code>numpy</code> arrays.
While linear algebra operations seem fast, I now am finding a bottleneck in a different issue: the summation of two distinct arrays. 
In the example below <code>WE3</code> and <code>T1</code> are two <code>1000X1000X1000</code> arrays.
First I calculate <code>WE3</code> using a numpy operation, then I sum those arrays.</p>

<pre><code>import numpy as np
import scipy as sp
import time
N = 100
n = 1000
X = np.random.uniform(size = (N,n))
wE = np.mean(X,0)
wE3 = np.einsum('i,j,k-&gt;ijk', wE, wE, wE) #22 secs
T1 = np.random.uniform(size = (n,n,n))
a = wE3 + T1 #115 secs
</code></pre>

<p>The calculation of <code>wE3</code> takes like 22 seconds, while the addition between <code>WE3</code> and <code>T1</code> takes 115 seconds.</p>

<p>Is there any known reason why the summation of those arrays is such slower than the calculation of WE3? They should have more or less the same complexity..</p>

<p>Is there a way to speed up that code? </p>

<p>That <code>np.einsum('i,j,k-&gt;ijk', wE, wE, wE)</code> part isn't doing any sum-reduction and is essentially just broadcasted elementwise multiplication. So, we can replace that with something like this -</p>

<pre><code>wE[:,None,None] * wE[:,None] * wE
</code></pre>

<p>Runtime test -</p>

<pre><code>In [9]: # Setup inputs at 1/5th of original dataset sizes
   ...: N = 20
   ...: n = 200
   ...: X = np.random.uniform(size = (N,n))
   ...: wE = np.mean(X,0)
   ...: 

In [10]: %timeit np.einsum('i,j,k-&gt;ijk', wE, wE, wE)
10 loops, best of 3: 45.7 ms per loop

In [11]: %timeit wE[:,None,None] * wE[:,None] * wE
10 loops, best of 3: 26.1 ms per loop
</code></pre>

<hr>

<p>Next up, we have <code>wE3 + T1</code>, where <code>T1 = np.random.uniform(size = (n,n,n))</code>  doesn't look like could  be helped in a big way as we have to create <code>T1</code> anyway and then it's just element-wise addition. It seems we can use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.add.html" rel="nofollow"><code>np.add</code></a> that lets us write back the results to one of the arrays : <code>wE3</code> or <code>T1</code>. Let's say we choose <code>T1</code>, if  that's okay to be modified. I guess this would bring slight memory efficiency as we won't be adding another variable into workspace.</p>

<p>Thus, we could do -</p>

<pre><code>np.add(wE3,T1,out=T1)
</code></pre>

<p>Runtime test -</p>

<pre><code>In [58]: def func1(wE3):
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     return wE3 + T1
    ...: 
    ...: def func2(wE3):
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     np.add(wE3,T1,out=T1)
    ...:     return T1
    ...: 

In [59]: # Setup inputs at 1/4th of original dataset sizes
    ...: N = 25
    ...: n = 250
    ...: X = np.random.uniform(size = (N,n))
    ...: wE = np.mean(X,0)
    ...: wE3 = np.einsum('i,j,k-&gt;ijk', wE, wE, wE)
    ...: 

In [60]: %timeit func1(wE3)
1 loops, best of 3: 390 ms per loop

In [61]: %timeit func2(wE3)
1 loops, best of 3: 363 ms per loop
</code></pre>

<p>Using <a href="http://stackoverflow.com/questions/39923085/speed-up-numpy-summation-between-big-arrays#comment67131799_39923433"><code>@Aaron's suggestion</code></a>, we can use a loop and assuming that writing back the results into <code>wE3</code> is okay, we could do -</p>

<pre><code>wE3 = wE[:,None,None] * wE[:,None] * wE
for x in wE3: 
    np.add(x, np.random.uniform(size = (n,n)), out=x)
</code></pre>

<p><strong>Final results</strong></p>

<p>Thus, putting back all the suggested improvements, finally the runtime test results were -</p>

<pre><code>In [97]: def func1(wE):
    ...:     wE3 = np.einsum('i,j,k-&gt;ijk', wE, wE, wE)
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     return wE3 + T1
    ...: 
    ...: def func2(wE):
    ...:     wE3 = wE[:,None,None] * wE[:,None] * wE
    ...:     for x in wE3: 
    ...:         np.add(x, np.random.uniform(size = (n,n)), out=x)
    ...:     return wE3
    ...: 

In [98]: # Setup inputs at 1/3rd of original dataset sizes
    ...: N = 33
    ...: n = 330
    ...: X = np.random.uniform(size = (N,n))
    ...: wE = np.mean(X,0)
    ...: 

In [99]: %timeit func1(wE)
1 loops, best of 3: 1.09 s per loop

In [100]: %timeit func2(wE)
1 loops, best of 3: 879 ms per loop
</code></pre>

python arrays numpy