# Initialize import pygame import random import time from pygame.locals import * from pygame.color import THECOLORS # constants EVAPORATE = 0.999 DIFFUSION = 1 ANTCOUNT = 1 DROPRATE = 0.9 INITPHER = 20 PHERBOUND = 0.01 PROBRANDOM = 0.1 SLEEPTIME = 0 OFFSET = 30 SIZE = 10 XMAX = 40 YMAX = 40 VONNEUMANN = 0 MOORE = 1 EDGE = 0 TORUS = 1 class Cell(pygame.sprite.Sprite): def __init__(self, pos, size, state): pygame.sprite.Sprite.__init__(self) self.x = pos[0] self.y = pos[1] self.size = size self.newstate = state self.image = pygame.Surface((self.size * 2, self.size * 2)).convert() self.image.fill(THECOLORS["green"]) self.image.set_colorkey(THECOLORS["green"]) self.rect = pygame.rect.Rect(self.x - self.size, self.y - self.size, self.size * 2, self.size * 2) self.update() def changeState(self): otherP = 0 for n in self.neighbors: otherP += n.state otherP += self.state otherP /= len(self.neighbors) + 1 self.newstate = EVAPORATE * (self.state + (DIFFUSION * (otherP - self.state))) if self.newstate < PHERBOUND: self.newstate = 0 def update(self): self.state = self.newstate if self.state > 0: pygame.draw.rect(self.image, (min(((int(self.state) / 32) + 1) * 32, 255), 0, 0), pygame.rect.Rect(0, 0, self.size * 2, self.size * 2)) else: pygame.draw.rect(self.image, THECOLORS["black"], pygame.rect.Rect(0, 0, self.size * 2, self.size * 2)) class Ant(pygame.sprite.Sprite): states = {0:THECOLORS["blue"], 1:THECOLORS["green"]} def __init__(self, pos, size, state, forest, food, edge): pygame.sprite.Sprite.__init__(self) self.food = food self.forest = forest self.edge = edge self.cx = pos[0] self.cy = pos[1] self.hx = self.cx self.hy = self.cy self.x = pos[0] * SIZE + OFFSET self.y = pos[1] * SIZE + OFFSET self.size = size self.newstate = state self.image = pygame.Surface((self.size * 2, self.size * 2)).convert() self.image.fill(THECOLORS["white"]) self.image.set_colorkey(THECOLORS["white"]) self.update() def changeState(self): if self.state == 0: found = False for f in self.food: if f.newstate > 0 and f.cx == self.cx and f.cy == self.cy: f.newstate = f.newstate - 1 self.newstate = 1 self.drop = INITPHER found = True break if not found: places = [] if self.cx > 0: places.append((-1, 0)) if self.cy > 0: places.append((0, -1)) if self.cx < XMAX - 1: places.append((1, 0)) if self.cy < YMAX - 1: places.append((0, 1)) # Make a random move # Follow Pheremone Gradient count = 0 i = 0 max = 0 if self.edge == EDGE: bound = 1 else: bound = 0 for tg in places: if self.forest.g2d[self.cx + bound + tg[0]][self.cy + bound + tg[1]].state > max: max = self.forest.g2d[self.cx + bound + tg[0]][self.cy + bound + tg[1]].state count = i i += 1 if max == 0 or random.random() < PROBRANDOM: togo = random.choice(places) else: togo = places[count] self.cx += togo[0] self.cy += togo[1] self.x += togo[0] * SIZE self.y += togo[1] * SIZE if self.state == 1: if self.cx == self.hx and self.cy == self.hy: self.newstate = 0 else: if self.edge == EDGE: bound = 1 else: bound = 0 self.forest.g2d[self.cx + bound][self.cy + bound].newstate += self.drop self.drop *= DROPRATE xd = abs(self.cx - self.hx) yd = abs(self.cy - self.hy) if xd > yd: togo = (-(self.cx - self.hx)/abs(self.cx - self.hx), 0) else: togo = (0, -(self.cy - self.hy)/abs(self.cy - self.hy)) self.cx += togo[0] self.cy += togo[1] self.x += togo[0] * SIZE self.y += togo[1] * SIZE def update(self): self.state = self.newstate self.rect = pygame.rect.Rect(self.x - self.size, self.y - self.size, self.size * 2, self.size * 2) pygame.draw.circle(self.image, Ant.states[self.state], (self.size, self.size), self.size) class Food(pygame.sprite.Sprite): def __init__(self, pos, size, state): pygame.sprite.Sprite.__init__(self) self.cx = pos[0] self.cy = pos[1] self.x = pos[0] * SIZE + OFFSET self.y = pos[1] * SIZE + OFFSET self.size = size self.newstate = state self.image = pygame.Surface((self.size * 2, self.size * 2)).convert() self.image.fill(THECOLORS["white"]) self.image.set_colorkey(THECOLORS["white"]) self.update() def changeState(self): pass def update(self): self.state = self.newstate self.rect = pygame.rect.Rect(self.x - self.size, self.y - self.size, self.size * 2, self.size * 2) if self.state > 0: color = THECOLORS["orange"] else: color = THECOLORS["grey"] pygame.draw.circle(self.image, color, (self.size, self.size), self.size) class Forest(): def __init__(self, x, y, offset, neighborhood, edge): self.x = x self.y = y self.offset = offset self.neighborhood = neighborhood self.edge = edge self.sprites = pygame.sprite.Group() d = Cell((1,1), 1, 0) if self.edge == EDGE: self.g2d = [[d] * (self.x + 2)] else: self.g2d = [] for i in range(self.x): if self.edge == EDGE: grid = [d] else: grid = [] for j in range(self.y): state = 0 c = Cell((i*SIZE + self.offset, j*SIZE + self.offset), SIZE/2, state) grid.append(c) self.sprites.add(c) if self.edge == EDGE: grid.append(d) else: pass self.g2d.append(grid) if self.edge == EDGE: self.g2d.append([d] * (self.x + 2)) self.setNeighborhood() def setNeighborhood(self): # run through all cells and change state if necessary if self.edge == EDGE: bound = 1 else: bound = 0 modx = self.x + bound * 2 mody = self.y + bound * 2 for i in range(bound, self.x + bound): for j in range(bound, self.y + bound): n = [self.g2d[(i-1) % modx][j % mody], self.g2d[(i+1) % modx][j % mody], self.g2d[i % modx][(j-1) % mody], self.g2d[i % modx][(j+1) % mody]] if self.neighborhood == MOORE: n.extend([self.g2d[(i-1) % modx][(j-1) % mody], self.g2d[(i-1) % modx][(j+1) % mody], self.g2d[(i+1) % modx][(j+1) % mody], self.g2d[(i+1) % modx][(j-1) % mody]]) self.g2d[i][j].neighbors = n def update(self): # run through all cells and change state if necessary if self.edge == EDGE: bound = 1 else: bound = 0 for i in range(bound, self.x + bound): for j in range(bound, self.y + bound): self.g2d[i][j].changeState() def getCounts(self): # run through all cells and change state if necessary if self.edge == EDGE: bound = 1 else: bound = 0 c = {} for i in range(bound, self.x + bound): for j in range(bound, self.y + bound): state = self.g2d[i][j].state if state in c: c[state] += 1 else: c[state] = 1 return c def main(): pygame.init() # Display screen = pygame.display.set_mode((640, 480)) # Background background = pygame.Surface(screen.get_size()).convert() background.fill(THECOLORS["white"]) # Entities t = 0 f = Forest(XMAX, YMAX, OFFSET, MOORE, EDGE) food = pygame.sprite.Group() food.add(Food((10, 10), SIZE/2, 255)) food.add(Food((35, 20), SIZE/2, 255)) food.add(Food((35, 5), SIZE/2, 255)) food.add(Food((7, 12), SIZE/2, 255)) # HOME food.add(Food((20, 20), SIZE, 0)) colony = pygame.sprite.Group() for i in range(ANTCOUNT): colony.add(Ant((XMAX/2, YMAX/2), SIZE/2, 0, f, food, EDGE)) done = False while not done: t += 1 time.sleep(SLEEPTIME) events = pygame.event.get() for e in events: if (e.type == KEYDOWN): if (e.key == K_ESCAPE): done = True break screen.blit(background, (0,0)) print t, zeros = True for fc in food: print fc.state, if fc.state != 0: zeroes = False print if zeroes: done = True f.update() for ant in colony: ant.changeState() f.sprites.update() colony.update() food.update() f.sprites.draw(screen) colony.draw(screen) food.draw(screen) # Alter pygame.display.flip() raw_input() main()