Conway's game of life in NumPy

game_of_life.py

"""NumPy/SciPy implementation of Conway's Game of Life.

https://bitstorm.org/gameoflife/
"""

import argparse
import curses
import getopt
import time
from itertools import repeat
import sys

import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np
from scipy import signal


# Convolutional kernel that counts diagonals and top/bottom/sides as neighbors
diag_kernel = np.ones((3, 3), dtype=np.int8)
diag_kernel[1, 1] = 0

# Convolutional kernel that doesn't count diagonals as neighbors
# We give user option to choose either one.
cross_kernel = np.zeros((3, 3), dtype=np.uint8)
rows = np.array([[0, 2], [1, 1]])
cols = np.array([[1, 1], [0, 2]])
cross_kernel[rows, cols] = 1


class Board(np.ndarray):

    def __new__(cls, size=15, diag=True):
        """Initialize as a square array of zeros, shape (`size`, `size`)."""
        obj = np.zeros((size, size), dtype=np.uint8).view(cls)
        if diag:
            obj.kernel = diag_kernel
        else:
            obj.kernel = cross_kernel
        return obj

    def clear(self, out=True):
        """Clear (reset) the board to be empty."""
        self[:] = 0
        if out:
            return self

    def start(self, n, out=False):
        """Place a length-n row of ones somewhere randomly on the board."""
        if n >= self.shape[1]:
            raise ValueError('Piece must fit on board.')
        self.clear(out=False)

        row = np.random.randint(0, len(self))
        col = np.random.randint(0, self.shape[1] - n + 1)
        self[row, col:col+n] = np.ones((1, n), dtype=np.uint8)
        if out:
            return self

    @property
    def neighbors(self):
        """Elementwise number of neighbors."""
        # TODO: Borders fluid?  (We're currently treating them as such)
        return signal.convolve(self, self.kernel, mode='same')

    @property
    def is_empty(self):
        """True if entire board is empty.  'Stop' condition."""
        return np.count_nonzero(self) == 0

    def _evolve(self, out=True):
        """Make one turn."""
        populated = self.astype(np.bool_)
        empty_3n = np.logical_and(~populated, self.neighbors >= 3)
        pop_14n = np.logical_and(populated, np.logical_or(
            self.neighbors <= 1, self.neighbors >= 4))
        pop_23n = np.logical_and(populated, np.logical_or(
            self.neighbors == 2, self.neighbors == 3))
        self[:] = np.where(pop_14n, 0,
                           np.where(pop_23n, 1,
                                    np.where(empty_3n, 1, self)))
        if out:
            return self

    def evolve(self, out=True, turns=1, illustrate=False, sleep=1,
               xmarker=False):
        """Make multiple turns."""
        if illustrate:
            for _ in repeat(None, turns):
                if xmarker:
                    stdscr.addstr(0, 0, str(
                        np.where(self._evolve(out=True), 'x', 'o')))
                else:
                    stdscr.addstr(0, 0, str(self._evolve(out=True)))
                time.sleep(sleep)
                stdscr.refresh()

                if self.is_empty:
                    break
        else:
            for _ in repeat(None, turns):
                self._evolve(out=True)
                if self.is_empty:
                    break
            if out:
                return self


def plot_kernel(kernel: np.ndarray, figsize=(4, 4), fontsize=14):
    """Visualize convolutional kernels.

    kernel: square array of 1s/0s.
    """

    cmap = colors.ListedColormap(['green', 'white'])
    norm = colors.BoundaryNorm([0., 0.5, 1.], cmap.N)

    fig, ax = plt.subplots(figsize=figsize)
    ax.imshow(kernel, cmap=cmap.reversed(), norm=norm)
    ax.set_xticks([0.5, 1.5])
    ax.set_yticks([0.5, 1.5])
    ax.grid(which='major', axis='both', linestyle='-', color='white',
            linewidth=2)
    ax.tick_params(axis='both', which='both', bottom='off', top='off',
                   left='off', right='off', labelbottom='off', labeltop='off',
                   labelleft='off', labelright='off')
    for (i, j), z in np.ndenumerate(kernel):
        label_kwargs = dict(ha='center', va='center', fontsize=fontsize)
        if z == 1:
            ax.text(j, i, '{:d}'.format(z), color='white', **label_kwargs)
        else:
            ax.text(j, i, '{:d}'.format(z), color='black', **label_kwargs)
    return ax


if __name__ == '__main__':
    """Main game visualizer.

    Example: $ python3 game_of_life.py --size=14 --n=5 --turns=5
    """

    opts, _ = getopt.getopt(sys.argv[1:], shortopts='',
                            longopts=['size=', 'n=', 'turns=', 'diag=',
                                           'xmarker='])
    opts = dict(opts)
    size, n, turns = (int(opts.get(i, d)) for i, d in zip(
        ('--size', '--n', '--turns'), (15, 7, 10)))
    diag = True if opts.get('--diag', 'true').lower() == 'true' else False
    xmarker = True if opts.get('--xmarker', 'true').lower() == 'true' else False

    stdscr = curses.initscr()
    curses.noecho()
    curses.cbreak()
    board = Board(size=size, diag=diag)
    board.start(n=n)

    try:
        board.evolve(turns=turns, illustrate=True, xmarker=xmarker, out=False)
    finally:
        curses.echo()
        curses.nocbreak()
        curses.endwin()