ants sorting woodchips in space

ants.py

#!/usr/bin/env python

""" This is a simulation of ants moving woodchips around
using the pyscape classes.  Hopefully it demonstrates the
construction of central clusters using only decentralized
algorithms followed by autonomous agents. 

We begin with a scattering of woodchips across the grid.
Ants randomly walk around, following a simple rule:
    if you see a woodchip:
        if you have a woodchip:
            drop yours and walk off
        else:
            take the woodchip
            
Someone told me this would work and I want to see it for myself.
"""

import random
import time
import curses
from pyscape import *

# my own color constants
CP_WHITE = 0
CP_RED = 1
CP_BLACK = 8
CP_BLACKWHITE = 9

class AntWorld(World):
    def __init__(self, size, numants, win, statwin):
        World.__init__(self, size)

        self.win = win  # ncurses window object
        self.statwin = statwin

        # populate the grid with our cells
        cx, cy = size
        for y in range(cy):
            for x in range(cx):
                self.grid[y][x] = AntCell((x, y))
        # and some ants
        self.creatures = [ AntCreature(self.RandPos()) for x in range(numants) ]
        # and scatter some woodchips
        for i in range(cx * cy / 30):
            cell = self.CellFromPos(self.RandPos())
            cell.PutChip(True)
            cell.Commit()

    def CalcStats(self):
        """ Returns a tuple:  (numstacks, avgstackheight, maxstackheight, heldbyants, antcount, iteration) """
        cx, cy = self.size
        accum, numstacks, maxstack = 0, 0, 0

        for y in range(cy):
            for x in range(cx):
                count = self.grid[y][x].ChipCount()
                if count:
                    if count > maxstack:
                        maxstack = count
                    numstacks += 1
                    accum += count
        heldbyants = 0
        for ant in self.creatures:
            heldbyants += ant.ChipCount()
            
        return (numstacks, float(accum) / numstacks, maxstack, heldbyants, len(self.creatures), self.iteration)
                
            
    def Draw(self):
        # draw the state of the system on self.win
        self.win.erase()
        self.win.box()
        self.statwin.erase()
        self.statwin.box()
        
        cx, cy = self.size
        ox, oy = 1, 1 # drawing offsets
        
        # BUGBUG fix these limits
        try:
            for y in range(cy):
                for x in range(cx-1):
                    cell = self.grid[y][x]
                    # cells can now have multiple chips
                    if cell.HasChip():
                        chips = cell.ChipCount()
                        
                        if chips < 10:
                            charnum = ord('0') + chips
                            attr = curses.color_pair((chips-1) % 8)
                        elif 10 <= chips <= 36:
                            charnum = ord('A') + chips - 10
                            attr = curses.color_pair((chips-1) % 8) | curses.A_BOLD
                        else:
                            charnum = ord('@') # stands for "a lot"
                            attr = curses.color_pair(1)
                    else:
                        charnum = ord(' ')
                        attr = curses.color_pair(CP_WHITE)
                    self.win.addch(y+oy, x+ox, charnum, attr)

            for creature in self.creatures:
                x, y = creature.state['pos']
                if 0 <= x < cx and 0 <= y < cy:
                    
                    if creature.HasChip():
                        charnum = ord('*')
                        attr = curses.A_BOLD | curses.color_pair(CP_RED)
                    else:
                        charnum = ord('#')
                        attr = curses.A_BOLD | curses.color_pair(CP_WHITE)
                    self.win.addch(y+oy, x+ox, charnum, attr)
            
            stats = self.CalcStats()
            self.statwin.addnstr(1, 1, 'count: %4d  avg stack: %4.4f  max stack: %4d' % stats[0:3], cx)
            self.statwin.addnstr(2, 1, 'held:  %4d  ants:        %4d  iter: %9d' % stats[3:], cx)
            self.statwin.refresh()
            self.win.refresh()
        except curses.error, v:
            raise v
            pass # swallow bugs
        

# Multiple inheritance is fun and profitable.
class ChipMixIn:
    """ All code dealing with the storage of woodchips lives here. """
    
    # originally, an entity could only hold one chip at at time.
    # now cells can have stacks of chips.
    
    def HasChip(self):
        return self.state['woodchip'] > 0
    
    def ChipCount(self):
        return self.state['woodchip']

    def PutChip(self, chip = True):
        if chip:
            self.nextstate['woodchip'] += 1
        else:
            # can't take away what's not there.
            if not self.HasChip():
                return  
            self.nextstate['woodchip'] -= 1
    
    def ClearChips(self):
        self.nextstate['woodchip'] = 0

class AntCell(Cell, ChipMixIn):
    def __init__(self, pos):
        Cell.__init__(self, pos)
        self.ClearChips()
        self.Commit()
        # ...
    # well.  nothing much to do right now.
    # the environment is just a container here.

        
class AntCreature(Creature, ChipMixIn):
    def __init__(self, pos):
        Creature.__init__(self, pos)
        self.ClearChips()
        self.Commit()
        
    # the real guts of the simulation are here.
    def Tick(self, world):
        pos = self.state['pos']
        cell = world.CellFromPos(pos)
        assert cell # make sure we're really on the grid
        
        # have we encountered a chip?
        if cell.HasChip():
            if not self.HasChip():
                # we have no chip but we are on one - take it.
                cell.PutChip(False)
                self.PutChip(True)
            else: # we have a chip, so does the cell - drop it
                cell.PutChip(True)
                self.PutChip(False)

        # now randomly walk
        x, y = pos
        self.nextstate['pos'] = (random.choice((max(0, x-1), x, min(world.size[0]-1, x+1))), \
                                 random.choice((max(0, y-1), y, min(world.size[1]-1, y+1))))

def define_colors():
    # fixup colors -- this should be totally unnecessary
    for i in range(1, 8):
        curses.init_pair(i, i, curses.COLOR_BLACK)
    # the banner color pair
    curses.init_pair(CP_BLACKWHITE, curses.COLOR_BLACK, curses.COLOR_WHITE)
    curses.init_pair(CP_BLACK, curses.COLOR_BLACK, curses.COLOR_BLACK)

def draw_banner(win):
    # draw the banner prettily.
    title = '* ants.py -- Decentralized ants sorting small woodchips in space.  Woo.'
    scy, scx = win.getmaxyx()
    win.addnstr(0, 0, title + ' '*(scx-len(title)), scx, curses.color_pair(CP_BLACKWHITE))
    win.refresh()

# drawing courtesy of curses
def main(stdscr):
    nodelay = 1
    stdscr.nodelay(nodelay)
    scy, scx = stdscr.getmaxyx()
    
    # try to hide the cursor
    try: curses.curs_set(0)
    except curses.error: pass
    
    # set up the grid with margins
    cx, cy = scx-2, scy - 10
    
    define_colors()
    worldwin = stdscr.subwin(scy-5, scx, 1, 0)
    statuswin = stdscr.subwin(4, scx, cy + 6, 0)
    world = AntWorld((cx, cy), 20, worldwin, statuswin)
    draw_banner(stdscr)    
    
    # dirty hard loop for now
    while 1:
        # advance the world state
        world.Tick()
        world.Draw()

        c = stdscr.getch()
        if c in [ord('q'), ord('Q'), curses.KEY_EXIT]:
            break
        elif c in [ord('s'), ord('S')]:
            nodelay = not nodelay
            stdscr.nodelay(nodelay)
        
    stdscr.erase()
    stdscr.refresh()

   
if __name__ == '__main__':
    curses.wrapper(main)

pyscape.py

"""
This module contains various generic architecture functionality
for the pyscape system.  Specific simulations must import this library
and derive classes from Cell, Creature, and World to obtain useful
modelling behavior.
"""

# This will get more useful as I use it more and add more utility functions.

import random

class Entity:
    """ Abstract base for all interactive objects in the grid.
    
    Entities have a transactional state system - changes in state
    are actually made to a separate copy of the state (.nextstate)
    and then all entities switch to the next state in unison at the
    end of a turn.
    """
    
    def __init__(self):
        self.state = {} # grab bag to be filled by user
        self.nextstate = {}
        
    def Tick(self, world):
        """ called once for each entity for each turn in the 
        system.  used to update the values in self.nextstate from
        self.state and information found in the world. """
        pass
    
    def Commit(self):
        """ Called at the end of a turn to commit the tentative
        changes in self.nextstate to self.state.  Note that a 
        SHALLOW copy is done of nextstate.  Override if deeper
        copying must be done. """
        
        self.state = self.nextstate.copy()

        
class Cell(Entity):
    """ Represents a cell in the grid.  Cells can have any
    arbitrary behavior just as creatures do, but cells have
    a fixed position. """
    
    def __init__(self, pos):
        Entity.__init__(self)
        # pos must never change, so it is not stored in the statebag.
        self.pos = pos

        
class Creature(Entity):
    """ Represents a creature that can move about the grid
    and interact with cells and/or other creatures. """
    
    def __init__(self, pos):
        Entity.__init__(self)
        # pos can change, so it is in the statebag.
        self.nextstate['pos'] = pos
        self.Commit()
        
class World:
    """ Represents the World, which contains the grid matrix
    and all of the creatures. """
    
    def __init__(self, size):
        self.size = size
        cx, cy = size
        
        self.creatures = []
        self.grid = [[None]*cx for y in range(cy)]
        self.iteration = 0
        # derived class must populate grid with creatures/cells.
        
    def Tick(self):
        """ This is the master world update routine which calculates
        the next state of every entity and the commits to that state.
        Called once per iteration of the simulation. """
        
        cx, cy = self.size
        
        # update the cells
        for y in range(cy):
            for x in range(cx):
                if self.grid[y][x]:
                    self.grid[y][x].Tick(self)
        # now for the creatures
        for creature in self.creatures:
            creature.Tick(self)
        
        # And once again to commit to the next state.
        for y in range(cy):
            for x in range(cx):
                if self.grid[y][x]:
                    self.grid[y][x].Commit()
        for creature in self.creatures:
            creature.Commit()

        
        self.iteration += 1

    def RandPos(self):
        """ Returns a random tuple representing valid grid coordinates. """
        cx, cy = self.size
        return (random.randint(0, cx-1), random.randint(0, cy-1))

    def CellFromPos(self, pos):
        x, y = pos
        cx, cy = self.size
        if x < 0 or x >= cx or y < 0 or y >= cy:
            return None
        else:
            return self.grid[y][x]
        
    def Neighbors(self, cell):
        """ Return a list of all neighbors for this cell.  This will
        be a tuple of 8 Cell objects:
                    0 1 2
                    3 * 4
                    5 6 7
        Any cells that would be off the edge of the grid are represented
        by a None. """
        
        nabes = []
        px, py = cell.pos
        cx, cy = self.size
        for y in range(py - 1, py + 2):
            for x in range(px - 1, px + 2):
                # don't include this cell
                if (x, y) == (px, py):
                    continue
                nabes.append(self.CellFromPos((x, y)))
        
        return tuple(nabes) # tuplify the list.