Steering agents in Pythonista

Steering.py

from scene import *
from collections import namedtuple
import copy
import random

class Vector2:
	def __init__(self, x = 0.0, y = 0.0):
		self.x = x
		self.y = y
	
	def clone(self):
		return copy.deepcopy(self)
		
	def set(self, x = 0.0, y = 0.0):
		self.x = x
		self.y = y
		
	def setAdd(self, v1, v2):
		self.x = v1.x + v2.x
		self.y = v1.y + v2.y
		
	def add(self, v):
		self.setAdd(self, v)
		
	def setSub(self, v1, v2):
		self.x = v1.x - v2.x
		self.y = v1.y - v2.y
		
	def sub(self, v):
		self.setSub(self, v)
	
	def mul(self, s):
		self.x *= s
		self.y *= s
	
	def setMulAdd(self, v1, s, v2):
		self.x = v1.x + s * v2.x
		self.y = v1.y + s * v2.y
		
	def mulAdd(self, s, v):
		self.setMulAdd(self, s, v)
		
	def sqrNorm(self):
		return self.x**2 + self.y**2
		
	def norm(self):
		return sqrt(self.sqrNorm())
		
	def normalize(self, len = 1.0):
		n = self.norm()
		if n > 0:
			self.mul(len/n)
		else:
			self.set(0, 0)
		return n
		
	def clamp(self, len):
		n = self.norm()
		if n > len:
			self.mul(len/n)

def dot(v1, v2):
	return v1.x * v2.x + v1.y * v2.y

Agent = namedtuple('Agent', ['p', 'v', 'dv', 'maxSpeed', 'acceleration', 'color', 'radius', 'perceptionDistance', 'neighbors'])
defaultAgent = Agent(p = Vector2(), 
		             v = Vector2(),
		             dv = Vector2(),
		             maxSpeed = 2,
		             acceleration = 10,
		             color = Color(0.90, 0.90, 0.90),
		             radius = .25,
		             perceptionDistance = 3,
		             neighbors=[])
		             
scaleFactor = 80
	
class Crowd (Scene):
	def setup(self):
		self.agents = []
		self.inCreation = False
		self.inCreationAgent = copy.deepcopy(defaultAgent)
	
	def updatePositions(self, dt, agents):
		newAgents = []
		for a in agents:
			newP = a.p.clone()
			newP.mulAdd(dt, a.v)
			
			# collision detection
			collision = False
			toOther = Vector2()
			futureOtherP = Vector2()
			for other in agents:
				if other != a:
					toOther.setSub(other.p, a.p)
					if toOther.norm() > other.radius + a.radius:
						futureOtherP.setMulAdd(other.p, dt, other.v)
						toOther.setSub(futureOtherP, newP)
						if toOther.norm() <= other.radius + a.radius:
							collision = True
							break
						
			if not collision:
				newAgents.append(a._replace(p=newP))
			else:
				newAgents.append(a._replace(v=Vector2()))
		return newAgents
	
	def updateNeighborhood(self, agents):
		newAgents = []
		toOther = Vector2()
		for a in agents:
			# neighbor detection
			neighbors = []
			for other in agents:
				if other != a:
					toOther.setSub(other.p, a.p)
					if toOther.norm() < a.perceptionDistance:
						neighbors.append(other)
			newAgents.append(a._replace(neighbors=neighbors))
		return newAgents
			
	def updateVelocities(self, dt, agents):
		newAgents = []
		futureP = Vector2()
		for a in agents:
			newDv = a.dv.clone()
					
			# compute separation/cohesion/alignment forces
			separation = Vector2()
			cohesion = Vector2()
			alignment = Vector2()
			toNeighbor = Vector2()
			toNeighborVelocity = Vector2()
			for n in a.neighbors:
				toNeighbor.setSub(n.p, a.p)
				cohesion.mulAdd(1.0/len(a.neighbors), toNeighbor)
				toNeighborDistance = toNeighbor.normalize() - a.radius - n.radius
				separation.mulAdd(-math.pow(toNeighborDistance, -1),toNeighbor)
				toNeighborVelocity.setSub(n.v, a.v)
				toNeighborVelocityLen = toNeighborVelocity.normalize()
				alignment.mulAdd(math.log1p(toNeighborVelocityLen), toNeighborVelocity)
				
			
			# compute a replusion force from the edges of the screen
			leftEdgeRepulsion=Vector2(1,0)
			leftEdgeRepulsion.mul(math.pow(a.p.x - self.bounds.left()/scaleFactor,-2))
			rightEdgeRepulsion=Vector2(-1,0)
			rightEdgeRepulsion.mul(math.pow(a.p.x - self.bounds.right()/scaleFactor,-2))
			topEdgeRepulsion=Vector2(0,-1)
			topEdgeRepulsion.mul(math.pow(a.p.y - self.bounds.top()/scaleFactor,-2))
			bottomEdgeRepulsion=Vector2(0,1)
			bottomEdgeRepulsion.mul(math.pow(a.p.y - self.bounds.bottom()/scaleFactor,-2))
			edgeRepulsion = Vector2()
			edgeRepulsion.add(leftEdgeRepulsion)
			edgeRepulsion.add(rightEdgeRepulsion)
			edgeRepulsion.add(topEdgeRepulsion)
			edgeRepulsion.add(bottomEdgeRepulsion)
			
			# compute an force to reach maxSpeed
			accelerationForce = Vector2()
			toMaxSpeed = a.dv.clone()
			toMaxSpeed.normalize(a.maxSpeed)
			toMaxSpeed.sub(a.v)
			toMaxSpeedLen = toMaxSpeed.normalize()
			accelerationForce.mulAdd(math.log1p(toMaxSpeedLen), toMaxSpeed)
			
			# Blend the forces
			steeringForce = Vector2()
			steeringForce.mulAdd(1, separation)
			steeringForce.mulAdd(1, cohesion)
			steeringForce.mulAdd(1, alignment)
			steeringForce.mulAdd(1, edgeRepulsion)
			steeringForce.mulAdd(1, accelerationForce)
			steeringForce.clamp(a.acceleration)
			
			# Compute desired velocity 
			newDv.mulAdd(dt, steeringForce)
			newDv.clamp(a.maxSpeed)
			
			# accelerate to desired velocity
			newV = a.v.clone()
			delta = Vector2()
			delta.setSub(newDv, newV)
			delta.clamp(a.acceleration * dt)
			newV.add(delta)
			newV.clamp(a.maxSpeed)
			
			newAgents.append(a._replace(v=newV, dv=newDv))
		return newAgents
	
	def drawAgent(self, agent):
		fill(agent.color.r, agent.color.g, agent.color.b)
		ellipse(agent.p.x - agent.radius, 
			    agent.p.y - agent.radius, 
			    agent.radius*2, 
			    agent.radius*2)
		stroke(0.90, 0.90, 0.90)
		stroke_weight(0.03)
		line(agent.p.x, agent.p.y, 
		     agent.p.x + agent.dv.x, agent.p.y + agent.dv.y)
		stroke(1.00, 0.50, 0.00)
		stroke_weight(0.02)
		line(agent.p.x, agent.p.y, 
		     agent.p.x + agent.v.x, agent.p.y + agent.v.y)
		no_stroke()
		
	def draw(self):
		push_matrix()
		scale(scaleFactor, scaleFactor)
		self.agents = self.updateNeighborhood(self.agents)
		self.agents = self.updateVelocities(self.dt, self.agents)
		self.agents = self.updatePositions(self.dt, self.agents)
		
		background(0, 0, 0)
		for agent in self.agents:
			self.drawAgent(agent)
		if self.inCreation:
			self.drawAgent(self.inCreationAgent)
		pop_matrix()
			
	
	def touch_began(self, touch):
		if not self.inCreation:
			self.inCreation = True
			self.inCreationAgent = self.inCreationAgent._replace(p=Vector2(touch.location.x/scaleFactor,touch.location.y/scaleFactor))
			self.inCreationAgent = self.inCreationAgent._replace(radius=random.uniform(.25,.35))
	
	def touch_moved(self, touch):
		if self.inCreation:
			self.inCreationAgent = self.inCreationAgent._replace(dv=Vector2(touch.location.x/scaleFactor - self.inCreationAgent.p.x, 
			                                        touch.location.y/scaleFactor - self.inCreationAgent.p.y))

	def touch_ended(self, touch):
		if self.inCreation:
			self.inCreationAgent = self.inCreationAgent._replace(
			    dv=Vector2(touch.location.x/scaleFactor - self.inCreationAgent.p.x, 
			               touch.location.y/scaleFactor - self.inCreationAgent.p.y),
			    v=Vector2(),
			    color=Color(r=touch.location.x/self.size.h, 
			                g=1.0-touch.location.y/self.size.w))
			self.inCreationAgent = self.inCreationAgent._replace(maxSpeed=self.inCreationAgent.dv.norm())
			self.inCreation = False
			self.agents.append(self.inCreationAgent)
			self.inCreationAgent = copy.deepcopy(defaultAgent)

run(Crowd())