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[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 = 0) and (r  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]  3):
                    Z[x][y] = 0
            else:
                if (N[x][y] == 3):
                    Z[x][y] = 1

could be replaced by:

set_zero_idxs = (Z==1) & ((N3))
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 = 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!

StackExchange Code Review Q#46011, answer score: 13