random walk for different types of Drunks

#quote from 'introduction to computation and programming
#using Python, revised, MIT press'
import random
import pylab

def stdDev(X):
mean = sum(X)/len(X)
tot = 0.0
for x in X:
tot += (x-mean)**2
return (tot/len(X))**0.5

def CV(X):
mean = sum(X)/len(X)
return stdDev(X)/mean

class Location(object):

def __init__(self, x, y):
"""x and y are floats"""
self.x = x
self.y = y

def move(self, deltaX, deltaY):
"""deltaX and deltaY are floats"""
return Location(self.x + deltaX, self.y + deltaY)

def getX(self):
return self.x

def getY(self):
return self.y

def distFrom(self, other):
ox = other.x
oy = other.y
xDist = self.x - ox
yDist = self.y - oy
return (xDist**2 + yDist**2)**0.5

def __str__(self):
return '<' + str(self.x) + ', ' + str(self.y) + '>'

class Field(object):

def __init__(self):
self.drunks = {}

def addDrunk(self, drunk, loc):
if drunk in self.drunks:
raise ValueError('Duplicate drunk')
else:
self.drunks[drunk] = loc

def moveDrunk(self, drunk):
if drunk not in self.drunks:
raise ValueError('Drunk not in field')
xDist, yDist = drunk.takeStep()
currentLocation = self.drunks[drunk]
#use move method of Location to get new location
self.drunks[drunk] = currentLocation.move(xDist, yDist)

def getLoc(self, drunk):
if drunk not in self.drunks:
raise ValueError('Drunk not in field')
return self.drunks[drunk]

class Drunk(object):

def __init__(self, name = None):
"""Assumes name is a str"""
self.name = name

def __str__(self):
if self != None:
return self.name
return 'Anonymous'

class UsualDrunk(Drunk):
def takeStep(self):
stepChoices = [(0.0, 1.0), (0.0, -1.0), (1.0, 0.0), (-1.0, 0.0)]
return random.choice(stepChoices)

class ColdDrunk(Drunk):
def takeStep(self):
stepChoices = [(0.0, 1.0), (0.0, -2.0), (1.0, 0.0), (-1.0, 0.0)]
return random.choice(stepChoices)

class EWDrunk(Drunk):
def takeStep(self):
stepChoices = [(1.0, 0.0), (-1.0, 0.0)]
return random.choice(stepChoices)

class styleIterator(object):
def __init__(self, styles):
self.index = 0
self.styles = styles

def nextStyle(self):
result = self.styles[self.index]
if self.index == len(self.styles) - 1:
self.index = 0
else:
self.index += 1
return result

def walk(f, d, numSteps):
"""Assumes: f a field, d a Drunk in f, and numSteps an int >= 0.
Moves d numSteps times, and returns the difference between
the final location and the location at the start of the walk."""
start = f.getLoc(d)
for s in range(numSteps):
f.moveDrunk(d)
return start.distFrom(f.getLoc(d))

def simWalks(numSteps, numTrials, dClass):
"""Assumes numSteps an int >= 0, numTrials an int > 0,
dClass a subclass of Drunk
Simulates numTrials walks of numSteps steps each.
Returns a list of the final distances for each trial"""
Homer = dClass()
origin = Location(0.0, 0.0)
distances = []
for t in range(numTrials):
f = Field()
f.addDrunk(Homer, origin)
distances.append(walk(f, Homer, numSteps))
return distances

def simDrunk(numTrials, dClass, walkLengths):
meanDistances = []
cvDistances = []
for numSteps in walkLengths:
print 'Starting simulation of', numSteps, 'steps'
trials = simWalks(numSteps, numTrials, dClass)
mean = sum(trials)/float(len(trials))
meanDistances.append(mean)
cvDistances.append(stdDev(trials)/mean)
return (meanDistances, cvDistances)

def simAll(drunkKinds, walkLengths, numTrials):
styleChoice = styleIterator(('b-', 'r:', 'm-.'))
for dClass in drunkKinds:
curStyle = styleChoice.nextStyle()
print 'Starting simulation of', dClass.__name__
means, cvs = simDrunk(numTrials, dClass, walkLengths)
cvMean = sum(cvs)/float(len(cvs))
pylab.plot(walkLengths, means, curStyle,
label = dClass.__name__ + '(CV = ' + str(round(cvMean, 4)) + ')')
pylab.title('Mean Distances from origin ('
+ str(numTrials) + ' trials)')
pylab.xlabel('Number of Steps')
pylab.ylabel('Distances from Origin')
pylab.legend(loc = 'best')
pylab.semilogx()
pylab.semilogy()

simAll((UsualDrunk, ColdDrunk, EWDrunk), (10, 100, 1000, 10000), 100)
pylab.show()

%run "C:\Users\Administrator\test.py"

Starting simulation of UsualDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps

Starting simulation of 100000 steps

Starting simulation of ColdDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps

Starting simulation of 100000 steps

Starting simulation of EWDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps

Starting simulation of 100000 steps

%run "C:\Users\Administrator\test.py"

Starting simulation of UsualDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps

Starting simulation of ColdDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps

Starting simulation of EWDrunk

Starting simulation of 10 steps

Starting simulation of 100 steps

Starting simulation of 1000 steps

Starting simulation of 10000 steps