Numpy: Single loop vectorized code slow compared to two loop iteration
<p>The following codes iterates over each element of two array to compute pairwise euclidean distance. </p>

<pre><code>def compute_distances_two_loops(X, Y):
    num_test = X.shape[0]
    num_train = Y.shape[0]
    dists = np.zeros((num_test, num_train))
    for i in range(num_test):
        for j in range(num_train):
            dists[i][j] = np.sqrt(np.sum((X[i] - Y[j])**2))
    return dists
</code></pre>

<p>The following code serves the same purpose but with single loop.</p>

<pre><code>def compute_distances_one_loop(X, Y):
    num_test = X.shape[0]
    num_train = Y.shape[0]
    dists = np.zeros((num_test, num_train))
    for i in range(num_test):
        dists[i, :] = np.sqrt(np.sum((Y - X[i])**2, axis=1))
    return dists
</code></pre>

<p>Below are time comparison for both.</p>

<pre><code>two_loop_time = time_function(compute_distances_two_loops, X, Y)
print ('Two loop version took %f seconds' % two_loop_time)

&gt;&gt; Two loop version took 20.3 seconds

one_loop_time = time_function(compute_distances_one_loop, X, Y)
print ('One loop version took %f seconds' % one_loop_time)

&gt;&gt; One loop version took 80.9 seconds
</code></pre>

<p>Both X and Y are numpy.ndarray with shape - </p>

<p>X: (500, 3000)
Y: (5000, 3000)</p>

<p>Out of intuition the results are not correct, the single loop should run at least with same speed. What am I missing here ? </p>

<p>PS: The result is not from a single run. I ran the code number of times, on different machines, the results are similar.</p>

<p>The reason is size of arrays within the loop body.
In the two loop variant works on two arrays of 3000 elements. This easily fits into at least the level 2 cache of a cpu which is much faster than the main memory but it is also large enough that computing the distance is much slower than the python loop iteration overhead.</p>

<p>The second case the loop body works on on 5000 * 3000 elements. This is so much that the data needs go to main memory in each computation step (first the Y-X[i] subtraction into a temporary array, squaring the temporary into another temporary and then read it back to sum it).
The main memory is much slower than the cpu for the simple operations involved so it takes much longer despite removing a loop.
You could speed it up a bit by using inplace operations writing into preallocated temporary array, but it will still be slower than the two loop variant for these array sizes.</p>

<p>Note that <code>scipy.spatial.distance.cdist(X, Y)</code> is probably going to be fastest as it does not need any temporaries at all</p>

python performance numpy time nearest-neighbor