Python script for creating a Langton's Ant cellular automata

langton.py

#!/usr/bin/python
# Langton's Ant
# Abraham Hmiel, 2017
# CC 3.0 attribution sharealike license

import numpy as np
import re
import sys
import matplotlib.pyplot as plt
import matplotlib as mpl

# extensions (?)
# import matplotlib.animation as animation
# from optparse import OptionParser

"""
Global parameters

The plan is to have these as command line arguments. But for now this will do.
"""

dimen_x = 128
dimen_y = 128
maxiter = 11000
num_ants = 1
face_direction = 0
rule = 'rl'
xpos = int(dimen_x/2)
ypos = int(dimen_y/2)
nstates = len(rule)

# possiby extend to toroidal geometry, 2-sphere, Klein bottle, etc.
geometry = 'finite'

def checkrule(string, search=re.compile(r'[^lnru]').search):

    return not bool(search(string))


def main(argv):

    # parse command-line arguments, but not quite yet.
    # parser = OptionParser()
    # parser.add_option("-r", "--rule", dest="rule",
    #                  help="Turning rule. Must be a string containing 'l', 'r', 'n', and 'u' only",
    #                  default='lr')
    # (options, args) = parser.parse_args()

    if not checkrule(rule):
        sys.exit('Invalid rule. Exiting.')

    # create grid
    grid = Grid(dimen_x, dimen_y, geometry)
    # create ants
    ants = []
    for i in range(num_ants):
        ants.append(Ant(face_direction, xpos, ypos, rule, nstates))

    for i in np.arange(maxiter):
        # loop over ants (in case there's more than one)
        grid, ants = update(grid, ants, i)

    grid.final_plot(ants, maxiter-1)


def update(grid, ants, i):

    for ant in ants:
        try:
            grid.board = ant.move(grid.board)
        except:
            # general error handling: just quit, write debug msg
            grid.final_plot(ants, i)
            sys.exit("Something weird is going on")

        # check geometry to make sure we didn't fall off a cliff, quit if we did
        # for other topologies, apply boundary conditions
        if not grid.check_geometry(ant.position):
            # end the simulation if the ant hits the wall
            grid.final_plot(ants, i)
            sys.exit("Ant fell off the map at timestep = %d!" % i)

    return grid, ants


class Grid:

    """
    Create arrays as numpy arrays since indexing will be a little faster,
    and the total number of iterations may be extremely large
    """

    def __init__(self, dimen_x, dimen_y, geometry):

        self.dimen = (dimen_x, dimen_y)
        self.geometry = geometry
        self.board = np.zeros((self.dimen[0], self.dimen[1]), dtype=np.int)

    def final_plot(self, ants, step):

        # plot the board state and ants
        # use a mask to make the ant array transparent and overlay only
        # the ants' positions onto the final result
        y = np.zeros((self.dimen[0], self.dimen[1]))
        for ant in ants:
            y[ant.position[0], ant.position[1]] = 1
        y = np.ma.masked_where(y == 0, y)

        # use imshow to print matrix elements using gray colormap. Ants are red.
        plt.imshow(self.board, cmap=plt.get_cmap('gray_r'), interpolation='none')
        plt.imshow(y, cmap=mpl.cm.jet_r, interpolation='none')
        plt.savefig(rule+'-'+str(num_ants)+'ants'+'-'+str(step+1)+'steps'+'.png')
        plt.show()

    def check_geometry(self, antpos):

        # return true if in valid geometry, flase if ant has fallen off the map
        # also, for non-finite, but bounded geometries, adjust ant position
        check = True
        if self.geometry == 'finite' and (antpos[0] < 0 or antpos[0] > self.dimen[0] or
                                          antpos[1] < 0 or antpos[1] > self.dimen[0]):
            check = False

        return check

class Ant:

    """
    Facing Direction:

         Up              [0,-1]  1
    Left    Right  2 [-1,0]  [1,0]  0
        Down          3  [0,1]

    dirs = [[1,0],[0,-1],[-1,0],[0,1]]
    index of dirs is the face_direction
    Right turn applies cyclic shift in negative direction
    Left turn applies cyclic shift in positive direction
    """

    def __init__(self, face_direction, xpos, ypos, rule, nstates):

        self.face_direction = face_direction
        self.position = [xpos, ypos]
        self.rule = rule
        self.nstates = nstates
        self.possiblefacings = ((1, 0), (0, -1), (-1, 0), (0, 1))
        self.geometry = geometry

    def move(self, board):

        # get state of board and current direction
        state = board[self.position[0], self.position[1]]
        directive = self.rule[state]

        # change the ant's direction
        self.face_direction = self.cycle_dir(directive)

        # cyclically increment the state of the board
        board[self.position[0], self.position[1]] = (state + 1) % self.nstates

        # apply motion based on new direction
        self.position[0] = self.position[0] + self.possiblefacings[self.face_direction][0]
        self.position[1] = self.position[1] + self.possiblefacings[self.face_direction][1]

        return board

    def cycle_dir(self, directive):

        new_face_direction = None
        # perform a cyclic permutation on the possible facing
        # directions with respect to the movement directive
        if directive == 'r':
            new_face_direction = (self.face_direction - 1) % len(self.possiblefacings)
        elif directive == 'l':
            new_face_direction = (self.face_direction + 1) % len(self.possiblefacings)
        elif directive == 'u':
            new_face_direction = (self.face_direction + 2) % len(self.possiblefacings)
        elif directive == 'n':
            new_face_direction = self.face_direction

        return new_face_direction


#pretend there's command-line arguments for now
if __name__ == "__main__":
    main(sys.argv[1:])

Langton's Ant cellular automata.
For details, see: https://en.wikipedia.org/wiki/Langton's_ant