Here is the code as it stands right now:

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from random import randint

arraySize = 50
Z = np.array([[randint(0, 1) for x in range(arraySize)] for y in range(arraySize)])

def computeNeighbours(Z):
    rows, cols = len(Z), len(Z[0])
    N = np.zeros(np.shape(Z))

    for x in range(rows):
        for y in range(cols):
            Q = [q for q in [x-1, x, x+1] if ((q >= 0) and (q < cols))]
            R = [r for r in [y-1, y, y+1] if ((r >= 0) and (r < rows))]
            S = [Z[q][r] for q in Q for r in R if (q, r) != (x, y)]
            N[x][y] = sum(S)

    return N

def iterate(Z):
    Zprime = Z.copy()
    rows, cols = len(Zprime), len(Zprime[0])
    N = computeNeighbours(Zprime)

    for x in range(rows):
        for y in range(cols):
            if Zprime[x][y] == 1:
                if (N[x][y] < 2) or (N[x][y] > 3):
                    Zprime[x][y] = 0
                if (N[x][y] == 3):
                    Zprime[x][y] = 1

    return Zprime

fig = plt.figure()

Zs = [Z]
ims = []

for i in range(0, 100):
    im = plt.imshow(Zs[len(Zs)-1], interpolation = 'nearest', cmap='binary')

ani = animation.ArtistAnimation(fig, ims, interval=250, blit=True)

I am interested in knowing what sequence of optimizations one would perform for this sort of an application, so that I can get a handle on how to use NumPy's power for my current project, which is simply a (perhaps) three dimensional, cellular automaton with many rules.


1 Answer 1


I'll just be cheeky and post a slightly modified version of my SO answer here.

So first things first, you want to get rid of the loops. They are slow to execute.

The first loop:

for x in range(rows):
        for y in range(cols):
            if Z[x][y] == 1:
                if (N[x][y] < 2) or (N[x][y] > 3):
                    Z[x][y] = 0
                if (N[x][y] == 3):
                    Z[x][y] = 1

could be replaced by:

set_zero_idxs = (Z==1) & ((N<2) | (N>3))
set_one_idxs = (Z!=1) & (N==3)
Z[set_zero_idxs] = 0
Z[set_one_idxs] = 1

Here you are generating boolean arrays that indicate which values satisfy your criteria and then you are setting those values to either 0 or 1.

Benchmarking this we see a significant improvement in speed (x 180):

# new version without loop
In [49]: %timeit no_loop(z,n)
1000 loops, best of 3: 177 us per loop

# original version with loop 
In [50]: %timeit loop(z,n)
10 loops, best of 3: 31.2 ms per loop

The second major loop:

for x in range(rows):
        for y in range(cols):
            Q = [q for q in [x-1, x, x+1] if ((q >= 0) and (q < cols))]
            R = [r for r in [y-1, y, y+1] if ((r >= 0) and (r < rows))]
            S = [Z[q][r] for q in Q for r in R if (q, r) != (x, y)]
            N[x][y] = sum(S)

could be replaced by:

N = np.roll(Z,1,axis=1) + np.roll(Z,-1,axis=1) + np.roll(Z,1,axis=0) + np.roll(Z,-1,axis=0)

[Note: this is actually wrong for the game of life, as the code should check for diagonal neighbours as well. This is simply a refactoring of the original for loop construct.]

Here there is an implicit assumption that the array does not have bounds and that x[-1] is next to x[0]. If this is a problem, you could add a buffer of zeros around your array with:

shape = Z.shape
new_shape = (shape[0]+2,shape[1]+2)
# b_z is a new array which will be our buffer
b_z = np.zeros(new_shape)

# set the middle of the array equal to the original `Z`
b_z[1:-1,1:-1] = Z

# do our rolls on the buffered array so that we don't have boundary isssues
b_n = np.roll(b_z,1,axis=1) + np.roll(b_z,-1,axis=1) + np.roll(b_z,1,axis=0) + np.roll(b_z,-1,axis=0)

# write back the part of the array that is of interest to us
N = b_n[1:-1,1:-1]

and for a benchmark:

# original function with loops
In [4]: %timeit computeNeighbours(z)
10 loops, best of 3: 140 ms per loop 

# new function without a buffer
In [5]: %timeit noloop_computeNeighbours(z)
10000 loops, best of 3: 133 us per loop

# new function with a buffer to remove boundary counts
In [6]: %timeit noloop_with_buffer_computeNeighbours(z)
10000 loops, best of 3: 170 us per loop

So just a small improvement of a factor of x 1052. Hooray for Numpy!


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.