import numpy as np import random import matplotlib.pyplot as plt import matplotlib.animation as animation from matplotlib.colors import LogNorm ANTCOUNT = 200 INITFOOD = 100 EVAPORATE = 0.5 DIFFUSION = 1 INITDROP = 20 LOWERBOUND = 0.01 PROBRANDOM = 0.1 DROPRATE = 0.9 def von_neuman_neighbors(i, j, size): n = [] if i > 0: n.append((i - 1, j)) if j > 0: n.append((i, j - 1)) if i < size - 1: n.append((i + 1, j)) if j < size - 1: n.append((i, j + 1)) return n def moore_neighbors(i, j, size): n = [] if i > 0: n.append((i - 1, j)) if j > 0: n.append((i - 1, j - 1)) if j < size - 1: n.append((i - 1, j + 1)) if j > 0: n.append((i, j - 1)) if i < size - 1: n.append((i + 1, j)) if j > 0: n.append((i + 1, j - 1)) if j < size - 1: n.append((i + 1, j + 1)) if j < size - 1: n.append((i, j + 1)) return n # Keeps track of a food supply class Food(): def __init__(self, x, y, amount): self.x = x self.y = y self.amount = amount def eaten(self): if self.any_food(): self.amount -= 1 def add_food(self, amount): self.amount += amount def here(self, a): if self.x == a.x and self.y == a.y: return True else: return False def closer(self, a): xd = abs(self.x - a.x) yd = abs(self.y - a.y) if xd > yd: return (a.x -(a.x - self.x)/abs(a.x - self.x), a.y) else: return (a.x, a.y-(a.y - self.y)/abs(a.y - self.y)) def any_food(self): return self.amount > 0 # tracks the Ants that wander around the world looking for food class Ant(): states = {"FORAGING":0, "HOMING":1} def __init__(self, x, y, home, world, foods): self.x = x self.y = y self.home = home self.world = world self.foods = foods self.state = Ant.states["FORAGING"] self.drop = 0 def update(self): if self.state == Ant.states["FORAGING"]: # Look for food here first found = False for f in self.foods: if f.any_food() and f.here(self): f.eaten() self.state = Ant.states["HOMING"] self.drop = INITDROP found = True break # If no food here if not found: neighbors = moore_neighbors(self.x, self.y, self.world.size) if np.random.random() < PROBRANDOM: self.x, self.y = random.choice(neighbors) else: count = 0 i = 0 max = LOWERBOUND for a, b in neighbors: if self.world.cells[a, b] > max: max = self.world.cells[a, b] count = i i += 1 if max == LOWERBOUND: self.x, self.y = random.choice(neighbors) else: self.x, self.y = neighbors[count] elif self.state == Ant.states["HOMING"]: if self.home.here(self): self.state = Ant.states["FORAGING"] self.home.add_food(1) else: self.world.cells[self.x, self.y] += self.drop self.drop *= DROPRATE self.x, self.y = self.home.closer(self) class CA(): def __init__(self, size, food_l, num_ants): self.size = size self.x = np.arange(size) self.y = np.arange(size) self.cells = np.zeros((size, size)) self.time = 0 self.totalfood = 0 self.food = [] for i, j in food_l: self.food.append(Food(i, j, INITFOOD)) self.totalfood += INITFOOD self.home = Food(self.size / 2, self.size / 2, 0) self.ants = [] for i in range(num_ants): self.ants.append(Ant(size / 2, size / 2, self.home, self, self.food)) def pattern_setup(self, pattern): for p in pattern: self.cells[p[0] + self.size / 2, p[1] + self.size / 2] = INITDROP def random_setup(self, n): for i in range(n): a, b = np.random.random_integers(0, size - 1, 2) self.cells[a, b] = INITDROP def image_setup(self): self.plt = plt.imshow(self.cells, interpolation='nearest', origin='bottom', vmin=np.min(LOWERBOUND), vmax=np.max(INITDROP), norm=LogNorm(vmin=LOWERBOUND, vmax=INITDROP), cmap=plt.cm.jet) ax = [] ay = [] ac = [] for a in self.ants: ax.append(a.x) ay.append(a.y) ac.append("r") for f in self.food: ax.append(f.x) ay.append(f.y) ac.append("g") ax.append(self.home.x) ay.append(self.home.y) ac.append("y") self.ant_plot = plt.scatter(ax, ay, c=ac, s=30) def update(self): newcells = np.zeros((self.size, self.size)) for i in range(self.size): for j in range(self.size): neighbors = moore_neighbors(i, j, self.size) otherP = self.cells[i,j] for a, b in neighbors: otherP += self.cells[a,b] otherP /= len(neighbors) + 1 newcells[i,j] = (1 - EVAPORATE) * (self.cells[i,j] + (DIFFUSION * (otherP - self.cells[i,j]))) if newcells[i,j] < LOWERBOUND: newcells[i,j] = LOWERBOUND self.cells = newcells for a in self.ants: a.update() self.time += 1 def plot(self): self.plt.set_data(np.transpose(self.cells)) ap = [] for a in self.ants: ap.append((a.x, a.y)) for f in self.food: ap.append((f.x, f.y)) if self.totalfood == self.home.amount: print(self.time) raise Exception("Simulation", "Finished") ap.append((self.home.x, self.home.y)) self.ant_plot.set_offsets(ap) return self.plt, self.ant_plot size = 40 fig, ax = plt.subplots() ax.set_ylim(-1, size) ax.set_xlim(-1, size) food_locs = ((22, 11), (35, 8), (18, 33)) ca = CA(size, food_locs, ANTCOUNT) ca.image_setup() def update(data): ca.update() return ca.plot(), def data_gen(): while True: yield 1 ani = animation.FuncAnimation(fig, update, data_gen, blit=False, interval=40) plt.show()