Calculation/manipulation of numpy array
<p>Looking to make the this calculation as quickly as possible. I have X as n x m numpy array. I want to define Y to be the following:</p>

<p><code>Y_11 = 1 / (exp(X_11-X_11) + exp(X_11-X_12) + ... exp(X_11 - X_1N) ).</code></p>

<p>or for Y_00</p>

<p><code>1/np.sum(np.exp(X[0,0]-X[0,:]))</code></p>

<p>So basically, Y is also n x m where the i,j element is the 1 / sum_j' exp(X_ij - X_ij') </p>

<p>Any tips would be great! Thanks.</p>

<p><em>Sample code as requested:</em></p>

<pre><code>np.random.seed(111)
J,K = 110,120
X = np.random.rand(J,K)
Y = np.zeros((J,K))
for j in range(J):
    for k in range(K):
        Y[j,k] = 1/np.sum(np.exp(X[j,k]-X[j,:]))

# note each row will sum to 1 under this operation
np.sum(Y,axis=1)
</code></pre>

<p>Here are few fully vectorized approaches -</p>

<pre><code>def vectorized_app1(X):
    return 1/np.exp(X[:,None] - X[...,None]).sum(1)

def vectorized_app2(X):
    exp_vals = np.exp(X)
    return 1/(exp_vals[:,None]/exp_vals[...,None]).sum(1)

def vectorized_app3(X):
    exp_vals = np.exp(X)
    return 1/np.einsum('ij,ik-&gt;ij',exp_vals,1/exp_vals)
</code></pre>

<p><strong>Tricks used and lessons learnt</strong></p>

<ul>
<li><p>Extend dimensions with <code>None/np.newaxis</code> to bring in broadcasting and perform everything in a vectorized manner.</p></li>
<li><p><code>np.exp</code> is an expensive operation. So, using it on broadcasted huge array would be costly. So, the trick used is : <code>exp(A-B) = exp(A)/exp(B)</code>. Thus, we perform <code>np.exp(X)</code> upfront and then perform broadcasted divisions.</p></li>
<li><p>Those sum-reductions could be implemented with <code>np.einsum</code>. This brings in memory-efficiency as we won't have to create huge broadcasted arrays and this effects in massive performance boost.</p></li>
</ul>

<p>Runtime test -</p>

<pre><code>In [111]: # Setup input, X
     ...: np.random.seed(111)
     ...: J,K = 110,120
     ...: X = np.random.rand(J,K)
     ...: 

In [112]: %timeit original_approach(X)
1 loop, best of 3: 322 ms per loop

In [113]: %timeit vectorized_app1(X)
10 loops, best of 3: 124 ms per loop

In [114]: %timeit vectorized_app2(X)
100 loops, best of 3: 14.6 ms per loop

In [115]: %timeit vectorized_app3(X)
100 loops, best of 3: 3.01 ms per loop
</code></pre>

<p>Looks like <code>einsum</code> showed its magic again with <strong><code>100x+</code></strong> speedup!</p>

python performance numpy sum numpy-einsum