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()