Computation of inter-grid distance array for 2D lattice I have a 2D lattice with U × U number of grids. My goal is to create an array – of shape U^2 × U^2 – where each entry denotes the least-distance between each grid and any other grid.  

Currently, for U=60, my Python script takes ~5 minutes to complete. As a novice programmer, I suspect there must be a more efficient approach – I've read about performance issues related to numpy loops. My objective is to handle U=1000.

import numpy as np
np.random.seed(13)
import matplotlib.pyplot as plt
import time

start_time = time.time()

U = 60

def FromGridToList(i_, j_):
    return i_*U + j_

def FromListToGrid(l_):
    i = np.floor(l_/U)
    j = l_ - i*U
    return np.array((i,j))

Ulist = range(U**2)

dist_array = []
for l in Ulist:
    print l, 'of', U**2
    di = np.array([np.abs(FromListToGrid(l)[0]-FromListToGrid(i)[0]) for i, x in enumerate(Ulist)])
    di = np.minimum(di, U-di)

    dj = np.array([np.abs(FromListToGrid(l)[1]-FromListToGrid(i)[1]) for i, x in enumerate(Ulist)])
    dj = np.minimum(dj, U-dj)

    d = np.sqrt(di**2+dj**2)
    dist_array.append(d)

dist_array = np.vstack(dist_array)
print time.time() - start_time, 'seconds'
1. Introduction

Veedrac is totally correct: you can't expect to compute a \$ 1000^2 \$ by \$ 1000^2 \$ array of distances: this would have \$ 10^{12} \$ elements, and if each element is a double-precision floating-point number then that would require 8 terabytes of memory.

So you need to reconsider your approach to whatever it is you are trying to do. You're going to have to compute these distances when you use them, instead of precomputing an array. Unfortunately, you haven't explained what you are trying to do, so we can't help. The only thing we can do, given what you've presented, is show you how to compute your result a bit faster and a bit tidier.

2. Review


There are no docstrings or comments. What do your functions compute? What is the meaning of variables like U? It's very hard to provide advice about code like this when the purpose is unclear.
The use of the print statement means that the code is not portable to Python 3. Better to use the print function.
Python has a built-in module timeit for timing execution of code. It's better to use this instead of writing your own.
Code at the top level of a module makes it hard to test a program or time its execution. It would be better to put this code into a function, for example:

def dist_array(U=60):
    Ulist = range(U**2)
    a = []
    ...
    return np.vstack(a)


Then in the interactive interpreter you could run:

&gt;&gt;&gt; from timeit import timeit
&gt;&gt;&gt; timeit(dist_array, number=1)
554.5194303740282



3. Performance

As always when speeding up NumPy code, you should scrutinize every Python loop—that is, every for, while, generator and comprehension—to see if the computation can be vectorized, that is, transformed into a sequence of NumPy whole-array operations. The usual approach is to start by removing the innermost loops and work outwards.

So let's start with the group of lines:

di = np.array([np.abs(FromListToGrid(l)[0]-FromListToGrid(i)[0])
               for i, x in enumerate(Ulist)])
di = np.minimum(di, U-di)
dj = np.array([np.abs(FromListToGrid(l)[1]-FromListToGrid(i)[1])
               for i, x in enumerate(Ulist)])
dj = np.minimum(dj, U-dj)
d = np.sqrt(di**2+dj**2)


I'm going to give speedups in terms of percentages of the runtime for the original version of the code.


Note that x is unused, so we don't need the enumerate, we can just write for i in range(len(Ulist)) which is the same as for i in Ulist. This saves about 2%, bringing the runtime down to 98%.
Expand the function FromListToGrid inline:

di = np.array([np.abs(np.floor(l / U) - np.floor(i / U))
               for i in Ulist])
di = np.minimum(di, U - di)
dj = np.array([np.abs(l - np.floor(l / U) * U - i + np.floor(i / U) * U)
               for i in Ulist])
dj = np.minimum(dj, U - dj)
d = np.sqrt(di ** 2 + dj ** 2)


This brings the runtime down to 32%.
Note that l does not depend on i so some of the computation can be moved out of the loops:

zi = np.floor(l / U)
di = np.array([np.abs(zi - np.floor(i / U)) for i in Ulist])
di = np.minimum(di, U - di)
zj = l - zi * U
dj = np.array([np.abs(zj - i + np.floor(i / U) * U) for i in Ulist])
dj = np.minimum(dj, U - dj)
d = np.sqrt(di**2+dj**2)


This brings the runtime down to 22%.
Avoid the list comprehensions altogether using NumPy whole-array operations:

Uarray = np.arange(U ** 2)
zi = np.floor(l / U)
di = np.abs(zi - np.floor(Uarray / U))
di = np.minimum(di, U - di)
zj = l - zi * U
dj = np.abs(zj - Uarray + np.floor(Uarray / U) * U)
dj = np.minimum(dj, U - dj)
d = np.sqrt(di ** 2 + dj ** 2)


This brings the runtime down to 0.4%.
Move the computation of anything that does not depend on l out of the loop over l. Also, use numpy.hypot instead of np.sqrt(di ** 2 + dj ** 2):

a = []
Uarray = np.arange(U ** 2)
Ui = np.floor(Uarray / U)
Uj = Ui * U - Uarray
for l in Uarray:
    zi = np.floor(l / U)
    di = np.abs(zi - Ui)
    di = np.minimum(di, U - di)
    zj = l - zi * U
    dj = np.abs(zj + Uj)
    dj = np.minimum(dj, U - dj)
    d = np.hypot(di, dj)
    a.append(d)


This brings the runtime down to 0.2%.
Now we can start to remove the loop over l. This can be done in stages, lifting one line at a time out of the loop (and checking that the results are correct as we go). To show how this works, let's lift the computation of di out of the loop. We have:

for l in Uarray:
    zi = np.floor(l / U)
    di = np.abs(zi - Ui)
    di = np.minimum(di, U - di)
    # etc.


We can easily compute all the values of zi at once, like this:

Zi = np.floor(Uarray / U)
for l in Uarray:
    di = np.abs(Zi[l] - Ui)
    di = np.minimum(di, U - di)
    # etc.


Now we observe that Zi is just the same as Ui, so that becomes:

for l in Uarray:
    di = np.abs(Ui[l] - Ui)
    di = np.minimum(di, U - di)
    # etc.


Next we lift Ui[l] - Ui out of the loop. This expression computes the pairwise differences between Ui[l] and all elements in Ui. NumPy's "broadcasting" mechanism means that we can do this in one step, for all values of l, just by indexing Ui in two different ways, like this:

Ui.reshape(-1, 1) - Ui.reshape(1, -1)


In fact, broadcasting means that we don't even need the second reshape; this is the same as:

Ui.reshape(-1, 1) - Ui


Putting that together, we now have:

di = np.abs(Ui.reshape(-1, 1) - Ui)
di = np.minimum(di, U - di)
for l in Uarray:
    # etc.

We can keep applying this process to lift all the computation out of the loop (I won't go into the details), resulting in this:

def dist_array(U=60):
    Uarray = np.arange(U ** 2)
    Ui = np.floor(Uarray / U)
    di = np.abs(Ui.reshape(-1, 1) - Ui)
    di = np.minimum(di, U - di)
    Uj = Ui * U - Uarray
    dj = np.abs(Uarray.reshape(-1, 1) - Ui.reshape(-1, 1) * U + Uj)
    dj = np.minimum(dj, U - dj)
    return np.hypot(di, dj)



There's probably a little bit more time to be squeezed out by using NumPy's out= parameter to reduce the memory usage, but I'll stop here.

This is more than 500 times as fast as the original code, and computes the U=60 case in about a second on my laptop.

4. Simplifying


The computation of dj now looks like this:

    Uj = Ui * U - Uarray
    dj = np.abs(Uarray.reshape(-1, 1) - Ui.reshape(-1, 1) * U + Uj)


If we do the multiplication before the reshape, this becomes:

    dj = np.abs((Uarray - Ui * U).reshape(-1, 1) + Uj)


which is:

    dj = np.abs(-Uj.reshape(-1, 1) + Uj)


and since we're taking the absolute value, that's:

    dj = np.abs(Uj.reshape(-1, 1) - Uj)


which is symmetric with the computation of di.
We have:

    Uarray = np.arange(U ** 2)
    Ui = np.floor(Uarray / U)
    Uj = Ui * U - Uarray


but this is the same as:

    Uarray = np.arange(U ** 2)
    Ui, Uj = divmod(Uarray, U)


and this is the same as:

    Ui, Uj = np.mgrid[:U, :U]
    Ui = Ui.reshape(-1)
    Uj = Uj.reshape(-1)


(see numpy.mgrid).
Now that the computation of di and dj is so similar, it makes sense to extract the common code into a function, like this:

def dist_array(U=60):
    def d(i):
        i = np.abs(i.reshape(-1, 1) - i.reshape(-1))
        return np.minimum(i, U - i)
    i, j = np.mgrid[:U, :U]
    return np.hypot(d(i), d(j))