Simple animation for Conway's Game of Life with FuncAnimation

Question

I have made a simple Conway's Game of Life program in Python and I need help making an animation with matplotlib because tbh I'm very lost and I can't seem to get my head around how's it done.

My code looks like this:

import matplotlib.pyplot as plt
import numpy as np


def initialize(size):
    grid = np.random.choice([0, 1], size*size, p=[0.8, 0.2]).reshape(size, size)
    plt.imshow(grid)
    plt.show(block=False)
    plt.pause(0.2)
    return grid


def conway_step(grid, size):
    new_grid = np.zeros_like(grid)
    for x in range(size):
        for y in range(size):
            total = sum([grid[(x+i) % size, (y+j) % size] for i in range(-1, 2) for j in range(-1, 2)])
            if grid[x, y] == 1 and total-1 in (2, 3):
                new_grid[x, y] = 1
            elif grid[x, y] == 0 and total == 3:
                new_grid[x, y] = 1
            else:
                new_grid[x, y] = 0
    grid = np.copy(new_grid)
    return grid


def conway(random=True, size=100):
    grid = initialize(size)
    for i in range(30):
        grid = conway_step(grid, size)
        plt.imshow(grid)
        plt.show(block=False)
        plt.pause(0.2)
    return


if __name__ == "__main__":
    conway(size=100)

This works fine but I would like to implement this as an animation and possibly get an mp4 file out. I've tried something like this:

def conway(size):
    grid = initialize(size)
    fig, ax = plt.subplots()
    img = ax.imshow(grid)
    ani = animation.FuncAnimation(fig, conway_step, fargs=(grid, size))
    plt.show()

But it doesn't work. Any help?

Answer 1

The step function in FuncAnimation normally updates a graphical object, which is also should return (the return statement should end with a comma, as it needs to be part of a list or tuple).

The grid and the graphical object (img_plot in the example code) need to be global variables. If you want to save the animation, FuncAnimation needs a frames= parameter to avoid it would run indefinitely.

import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.colors import ListedColormap
import numpy as np

grid, grid_size, img_plot = None, None, None

def initialize(size):
    grid = np.random.choice([0, 1], size * size, p=[0.8, 0.2]).reshape(size, size)
    return grid

def conway_step(frame):
    global grid, grid_size, img_plot
    if frame < 3:   # no movement for the first few steps
        new_grid = grid
    else:
        new_grid = np.zeros_like(grid)
        for x in range(grid_size):
            for y in range(grid_size):
                total = sum(
                    [grid[(x + i) % grid_size, (y + j) % grid_size] for i in range(-1, 2) for j in range(-1, 2)])
                if grid[x, y] == 1 and total - 1 in (2, 3):
                    new_grid[x, y] = 1
                elif grid[x, y] == 0 and total == 3:
                    new_grid[x, y] = 1
                else:
                    new_grid[x, y] = 0
        grid = new_grid
    img_plot.set_data(new_grid)
    return img_plot,

def conway(random=True, size=100):
    global grid, grid_size, img_plot
    grid_size = size
    grid = initialize(size)
    fig, ax = plt.subplots(figsize=(10, 10))
    img_plot = ax.imshow(grid, interpolation='nearest', cmap=ListedColormap(['darkturquoise', 'yellow']))
    ax.set_xticks([])
    ax.set_yticks([])
    ani = animation.FuncAnimation(fig, frames=100, func=conway_step, interval=100)
    plt.tight_layout()
    ani.save('testconway.gif')
    plt.show()
    return ani

conway(size=100)