Custom Functions - Databricks

Custom Functions

Conway’s Game of Life is an interesting computer science simulation that takes place on board with many cells in a square configuration, much like a chess board. The simulation takes place with specific times steps, and each cell on the board can be either 1 (alive) or 0 (dead). After a specific time step, each cell is either alive or dead following these rules:

  • If the cell is alive, but has one or zero neighbours, it “dies” through “under-population.
  • If the cell is alive and has two or three neighbours, it stays alive.
  • If the cell has more than three neighbours it dies through over-population.
  • Any dead cell with three neighbours regenerates.

While the rules seem quite morbid, the actual simulation is quite straight-forward but creates very interesting patterns. We are going to create a TensorFlow program that rules Conway’s Game of Life, and in the process learn about custom py_func functions, and generates an animation like below:


Try it yourself »

First, let’s generate the board. This is quite basic, as it is just a matrix of 0s and 1s. We generate the initial board randomly, which will provide a different board each time this is run:

import tensorflow as tf
from matplotlib import pyplot as plt

shape = (50, 50)
initial_board = tf.random_uniform(shape, minval=0, maxval=2, dtype=tf.int32)

with tf.Session() as session:
    X =

fig = plt.figure()
plot = plt.imshow(X, cmap='Greys',  interpolation='nearest')

We generate an initial_board of randomly chosen 0s and 1s, and then run it to get the values. We then use matplotlib.pyplot to show it, using the imshow function, which basically just plots out the values from a matrix according to some cmap colour scheme. In this case, using ‘Greys’ results in a black and white matrix, and a single initial starting point for our Game of Life:

Updating the state of the board

As the board state of the Game of Life is represented as a matrix, it makes sense to use matrix operators to update it. This should provide a fast method to update the state at a given point in time.

The very talented Jake VanderPlas has done some excellent work in using SciPy and NumPy to update a given state in the Game of Life. His write up is worth reading and is available here. If you are interested in how the below code works, I recommend that you read Jake’s description. In a nutshell, though, the convolve2d line identifies how many neighbours each cell has (this is a common operator in computer vision). I’ve updated the code slightly to reduce the number of lines, see the updated function below:

import numpy as np
from scipy.signal import convolve2d

def update_board(X):
    # Check out the details at:
    # Compute number of neighbours,
    N = convolve2d(X, np.ones((3, 3)), mode='same', boundary='wrap') - X
    # Apply rules of the game
    X = (N == 3) | (X & (N == 2))
    return X

The update_board function is a function on NumPy arrays. It won’t work on Tensors, and to date, there isn’t a good way to do this in TensorFlow (you can definitely write it yourself though using the existing tools, it just isn’t straight-forward).

In version 0.7 of TensorFlow, a new feature py_func was added that allows us to take a python function and turn it into a node in TensorFlow.

At the time of writing (March 22nd), 0.6 was the official release, and it does not have py_func in it. I recommend following the instructions at TensorFlow’s Github page to install a nightly build for your system. As an example, for Ubuntu users, you download the relevant wheel file (a python installation file) and install it with:

python -m wheel install --force ~/Downloads/tensorflow-0.7.1-cp34-cp34m-linux_x86_64.whl

Keep in mind that you will need to have your TensorFlow source activated properly (if you are so inclined).

The end result should be that you have version 0.7 or later of TensorFlow installed. You can check this by running this code in your terminal:

python -c "import tensorflow as tf; print(tf.__version__)"

The result will be the version number, which at time of writing is 0.7.1.

On to the code:

board = tf.placeholder(tf.int32, shape=shape, name='board')
board_update = tf.py_func(update_board, [board], [tf.int32])

From here, you can run the initial board as per normal for a Tensor Op node (i.e. board_update). A slight thing to keep in mind is that the results from running board_update are a list of matrices, even though our function defined only a single return value. We compensate by adding [0] at the end of the line to get just the first result, our updated board that is stored in X.

with tf.Session() as session:
    initial_board_values =
    X =, feed_dict={board: initial_board_values})[0]

The resulting value, X is the updated board after the initial configuration. It looks much like an initial random board, but we were never shown the original (although you can update the code to plot both values out).


Here is where things get really interesting, although from a TensorFlow perspective, we have already done the hard work for this lesson. We can use matplotlib to display and animation, and therefore show the simulation state as is progresses through time steps, like our original GIF. The intricacies of using matplotlib animations is a little tricky, but you create a function that updates and returns the plot, and you call the animation code with that function:

import matplotlib.animation as animation
def game_of_life(*args):
    X =, feed_dict={board: X})[0]
    return plot,

ani = animation.FuncAnimation(fig, game_of_life, interval=200, blit=True)

Hint: you will need to remove the from the earlier code to make this run!

I’ll leave the problem of putting the pieces of the puzzle as an exercise to the reader, but the end result will be a window appearing with the game state updating every 200 milliseconds.

Send us a message if you got it working!

1) Get the full code example working, resulting in an animated Game of Life using  matplotlib and TensorFlow
2) Conway’s Game of Life has been extensively studied, and has lots of interesting patterns. Create a function that loads patterns from a file, and uses those instead of the random board. I recommend starting with the Gosper Glider Gun.
3) One issue (feature?) with Game of Life is that boards can repeat, resulting in loops that never quite stop. Write some code that keeps track of previous game states, and stops looping when a game state is repeated.