# Optimizing very simple piece of “Game of Life” code by taking advantage of NumPy's functionality

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)
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)
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
else:
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')
ims.append([im])
Zs.append(iterate(Zs[len(Zs)-1]))

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


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.

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
else:
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 : %timeit no_loop(z,n)
1000 loops, best of 3: 177 us per loop

# original version with loop
In : %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. If this is a problem, you could add a buffer of zeros around your array with:

shape = Z.shape
new_shape = (shape+2,shape+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 : %timeit computeNeighbours(z)
10 loops, best of 3: 140 ms per loop

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

# new function with a buffer to remove boundary counts
In : %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!