Simple program to look at elementary cellular automata

ca.py

#! ./venv/bin/python3

# Q to quit
# R to randomize
# S to make a simple initial state
# + or - to change rule number by 1 (watch the title bar)
# Digits to enter the rule number; you may need to prefix it with some 0s.
# Space to pause / resume

import random
import sys

import pygame

WIDTH=800
HEIGHT=600
STARTRULE=153


def random_row():
    rv = [0] * WIDTH
    randbits = random.getrandbits(WIDTH)
    for j in range(WIDTH):
        rv[j] = randbits % 2
        randbits >>= 1
    return rv


def make_rulearray(rule):
    return [(rule >> n) % 2 for n in range(8)]


def main():
    rule = STARTRULE

    pygame.init()
    clock = pygame.time.Clock()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Cellular Automaton %i" % (rule,))
    drawsurf = pygame.Surface(screen.get_size())
    colors = [pygame.Color("black"), pygame.Color("white")]

    initial = random_row()
    rulearray = make_rulearray(rule)

    running = True

    while True:
        clock.tick(30)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return 0
            elif event.type == pygame.KEYDOWN:
                if event.unicode == "q":
                    return 0
                elif event.unicode == " ":
                    running = not running
                elif event.unicode == "r":
                    initial = random_row()
                elif event.unicode == "s":
                    initial = [0] * WIDTH
                elif event.unicode == "+" or event.unicode == "=":
                    rule += 1
                elif event.unicode == "-":
                    rule -= 1
                elif "0" <= event.unicode <= "9":
                    rule = (rule * 10 + int(event.unicode))
                    if rule > 255:
                        rule = int(event.unicode)
                rulearray = make_rulearray(rule)
                pygame.display.set_caption("Cellular Automaton %i" % (rule,))

        if not running:
            continue

        bitToFlip = random.randrange(WIDTH)
        initial[bitToFlip] ^= 1
        newrow = initial

        pxarray = pygame.PixelArray(drawsurf)
        for i in range(HEIGHT):
            # The next line takes about 1/10 of the runtime.
            rowimg = [colors[v] for v in newrow]
            pxarray[:, i] = rowimg
            oldrow = newrow
            # The next line takes about 2/3 of the runtime.
            newrow = [rulearray[4 * oldrow[j-1] +
                                2 * oldrow[j] +
                                1 * oldrow[j+1]]
                       for j in range(WIDTH-1)]
            newrow.append(rulearray[4 * oldrow[WIDTH-2] +
                                    2 * oldrow[WIDTH-1] +
                                    1 * oldrow[0]])
        pxarray = None
        screen.blit(drawsurf, (0, 0))
        pygame.display.flip()


if __name__ == "__main__":
    sys.exit(main())