%matplotlib inline
import numpy as np
import pylab as plt
from matplotlib import animation
from JSAnimation import IPython_display
swarm_params = { 'Ncarbon': 6, 'Nsulfur': 8, 'Niron': 5,'Nseed': 1 }
color_dict = { 'carbon': 'k', 'sulfur' : 'y', 'iron' : 'g', 'seed' :'r' }
xlims = [-5,5]
ylims = [-5,5]
class agent():
def __init__(self,id_number,pars):
self.id = id_number
self.xpos = pars['x_init']
self.ypos = pars['y_init']
self.cat = pars['cat']
self.fixed = pars['fixed']
return None
def update_position(self, newx, newy):
self.xpos = newx
self.ypos = newy
return None
def checkBCs(self):
if self.xpos > xlims[1]:
self.update_position(-5,self.ypos)
if self.xpos < xlims[0]:
self.update_position(5,self.ypos)
if self.ypos > ylims[1]:
self.update_position(self.xpos,-5)
if self.ypos < ylims[0]:
self.update_position(self.xpos, 5)
def moveRight(self):
newx = self.xpos + 1
newy = self.ypos
self.update_position(newx,newy)
def moveLeft(self):
newx = self.xpos - 1
newy = self.ypos
self.update_position(newx,newy)
def moveUp(self):
newx = self.xpos
newy = self.ypos + 1
self.update_position(newx,newy)
def moveDown(self):
newx = self.xpos
newy = self.ypos - 1
self.update_position(newx,newy)
def agentSep(agent1,agent2):
return np.sqrt((agent1.xpos-agent2.xpos)**2+(agent1.ypos-agent2.ypos)**2)
class swarm():
def __init__(self, pars):
self.members = []
self.fig = plt.figure()
self.ax = self.fig.gca()
for i in np.arange(pars['Ncarbon']):
temp = agent(i,{'x_init':i, 'y_init':0, 'cat': 'carbon', 'fixed': False})
self.members.append(temp)
"""
for i in np.arange(pars['Nsulfur']):
temp = agent(i,{'x_init':i, 'y_init':1, 'cat': 'sulfur', 'fixed': False})
self.members.append(temp)
for i in np.arange(pars['Niron']):
temp = agent(i,{'x_init':i, 'y_init':2, 'cat': 'iron', 'fixed': False})
self.members.append(temp)
"""
for i in np.arange(pars['Nseed']):
temp = agent(i,{'x_init':i, 'y_init':3, 'cat': 'seed', 'fixed': True})
self.members.append(temp)
def checkMove(self,current,newx,newy):
for other in self.members:
if newx == other.xpos and newy == other.ypos:
return 1
return 0
def timeEvolve(self, N):
counter = 0
while counter < N:
for member in self.members:
if not member.fixed:
test = np.random.rand(1)
if test >= 0 and test < 0.25:
member.moveLeft()
if test >= 0.25 and test < 0.5:
member.moveRight()
if test >= 0.5 and test < 0.75:
member.moveUp()
if test >= 0.75 and test <= 1:
member.moveDown()
member.checkBCs()
member.fixed = self.shouldFixQ(member)
counter += 1
def shouldFixQ(self,member):
for other in self.members:
if agentSep(member,other) == 1 and other.fixed:
return True
return False
def showCurrentState(self):
self.ax.cla()
self.ax.set_xlim(xlims[0],xlims[1])
self.ax.set_ylim(ylims[0],ylims[1])
xvec = []
yvec = []
colvec = []
for member in self.members:
xvec.append(member.xpos)
yvec.append(member.ypos)
colvec.append(color_dict[member.cat])
scat = self.ax.scatter(xvec,yvec,c=colvec,)
return scat
def animInit(self):
self.ax.grid(b=True, which='major', color='0.65',linestyle='-')
def timeEvolveAndShowCurrentState(self):
print '1'
self.timeEvolve(1)
print '2'
scat = self.showCurrentState()
return scat
def animateSwarm(self):
anim = animation.FuncAnimation(self.fig, self.timeEvolveAndShowCurrentState, init_func=self.animInit,
frames=xrange(200))