#410 2D random walk animation in python

import numpy
import random
import matplotlib.pyplot as plt
from matplotlib import animation

# define some plot variables
fig, ax = plt.subplots(figsize=(8,8))
bound = 25
ax.set_xlim(-bound,bound)
ax.set_ylim(-bound,bound)


# define a numpy array to hold the locations visited on the random walk
locations = numpy.zeros((1,2)) # 1 row, 2 columns

# define a function to make animation frames
def run(i):
  global locations
  global bound
  
  # first frame is handled separately
  if i == 0:
    line, = ax.plot([], [])
    ax.plot(0,0,'r.')
    return line
  
  # generate a step of the random walk
  a=random.uniform(-1,1)
  b=random.uniform(-1,1)
  dirs = numpy.array([[a,b]])

  move = dirs[0]
  nextloc = [locations[-1] + move]
  locations = numpy.append(locations, nextloc, axis=0)
  
  # set the plot data
  xdata = locations[:,0]
  ydata = locations[:,1]
  
  ax.cla() # clear the previous plot (necessary for removing old dots)
  
  # update the plot limits
  ax.set_xlim(min(-bound,min(xdata)-1), max(bound,max(xdata)+1))
  ax.set_ylim(min(-bound,min(ydata)-1), max(bound,max(ydata)+1))
  
  # redraw the plot
  line, = ax.plot([], [],c="black",lw=0.5)
  line.set_data(xdata, ydata)
  ax.plot(locations[-1,0], locations[-1,1], 'r.')
  
  return line

animation = animation.FuncAnimation(fig, run, frames=500, interval=5)