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.

1. 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:

>>> from timeit import timeit
>>> timeit(dist_array, number=1)
554.5194303740282

1. 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 difference

StackExchange Code Review Q#85282, answer score: 3