Wa-Tor World on a torus

from wator import World, EMPTY, FISH, SHARK
from torus import Torus
import random
import numpy as np

# Save every SAVE_EVERYth chronon iteration.
SAVE_EVERY = 5
# Run the simulation for MAX_CHRONONS chronons (time intervals).
MAX_CHRONONS = 2000

# PRNG seed.
SEED = 15
random.seed(SEED)

# The major and minor radius of the torus.
c, a = 2.5, 1.5
# Image dimensions and scaling factors from torus units to image units.
width, height = 800, 600
scalex, scaley = 100, 100
# Tait-Bryan angles for intrinsic rotation of the figure.
alpha, beta, gamma = 90, 35, 235
# Camera position, C, and projection plane position, E (relative to C).
cx, cy, cz = 0, 6, 0
ex, ey, ez = 0, 3, 0
C, E = np.array((cx,cy,cz)), np.array((ex,ey,ez))

def preamble(fo):
    """The SVG preamble and styles."""

    print('<?xml version="1.0" encoding="utf-8"?>\n'

    '<svg xmlns="http://www.w3.org/2000/svg"\n' + ' '*5 +
       'xmlns:xlink="http://www.w3.org/1999/xlink" width="{}" height="{}" >'
            .format(width, height), file=fo)

    print("""
        <defs>
        <style type="text/css"><![CDATA[""", file=fo)

    print('path {stroke-width: 1px; stroke: #000;}', file=fo)

    print("""]]></style>
    </defs>""", file=fo)

world_width = nphi = 100
world_height = ntheta = int(world_width * c/a)

def get_colour(i, u):
    """Get a colour for quad i from the colourmap."""
    cm = {EMPTY: '#00008b', FISH: '#ff69b4', SHARK: '#ffd700'}
    return cm[u[i]]

def draw_torus(itr, u):
    """Draw a torus on iteration itr of the diffusion with array u."""
    torus = Torus(c, a, ntheta, nphi)
    torus.setup_torus()
    torus.rotate_xyz(alpha, beta, gamma)

    torus.get_perspective_view(C, E)
    torus.get_quads()

    # Draw the torus as a SVG image and save.
    with open('frames/wator-{:04d}.svg'.format(itr), 'w') as fo:
        preamble(fo)
        for i in torus.idy:
            quad = torus.quads[i] * (scalex, scaley) + (width/2, height/2)
            colour = get_colour(i, u.ravel())
            print('<path d="M{},{} L{},{} L{},{} L{},{} Z" fill="{}"/>'.format(
                *quad.ravel(), colour), file=fo)
        print('</svg>', file=fo)

world = World(world_width, world_height)
world.populate_world()
for chronon in range(MAX_CHRONONS):
    if not chronon % SAVE_EVERY:
        print('{}/{}: {}'.format(chronon, MAX_CHRONONS, len(world.creatures)))
        draw_torus(chronon, np.array(world.get_world_image_array()))
        alpha += 1
        beta += 1
    world.evolve_world()

initial_energies = {FISH: 15, SHARK: 5}
fertility_thresholds = {FISH: 3, SHARK: 6}