Kizirian

A random walk animation with matplotlib

I have met with but one or two persons in the course of my life who understood the art of Walking
- Henry David Thoreau, Walking

If you start at a light pole, spin to face a random direction, and walk 1 meter. Then spin to face a new direction and walk 1 meter. Then you repeat that forever. What does that look like? Where do you end up? Back at the pole? If not, how far away are you?

This is the question of what a random walk looks like. We won't answer all of these questions, but let's look at how we can generate a random walk animation with matplotlib.

In [1]:
from typing import Tuple

import numpy as np
from numpy import random

from matplotlib import animation, rc
import matplotlib.pyplot as plt

First, a little magic for embedding matplotlib animations in a notebook. Thanks to this post for helping me figure this out.

In [2]:
rc('animation', html='html5')

We can start with a function that returns a random vector with unit length:

In [3]:
def random_unit_vector() -> Tuple[np.float, np.float]:
    """Gets a random point on the unit circle"""
    angle = random.random() * 2 * np.pi
    return np.cos(angle), np.sin(angle)

To create a random walk with $N$ steps, we simply start at $(0, 0)$, add a random_unit_vector, and continue to add a random_unit_vector to the current position for $N$ iterations.

The get_random_walk function does exactly that and returns all steps from the random walk in a $2xN$ matrix. The values in row $0$ contain the $x$ coordinates and the values in row $1$ contain the $y$ coordinates.

In [4]:
def get_random_walk(num_steps: int) -> np.ndarray:
    locations = np.zeros((2, num_steps))
    for i in range(1, num_steps):
        next_x, next_y = random_unit_vector()
        locations[0, i] = locations[0, i-1] + next_x
        locations[1, i] = locations[1, i-1] + next_y
    return locations

Our plot will contain four elements:

  1. A bold red dot at the current step
  2. Faded red dots at all previous steps
  3. Text displaying the $(x, y)$ coordinates of the current step
  4. Text displaying the current distance from the origin
In [5]:
%%capture
fig = plt.figure()
ax = fig.add_subplot(
    aspect="equal",
    xlim=(-10, 10),
    ylim=(-10, 10)
)

previous_steps, = plt.plot([], [], 'ro', alpha=0.2)
current_step, = plt.plot([], [], 'ro', alpha=1.0)

loc_text = ax.text(0.02, 0.95, "", transform=ax.transAxes)
distance_text = ax.text(0.02, 0.90, "", transform=ax.transAxes)

The update_plot function is called at each step of the animation and will return the four elements above. We start by generating all of the points in our random walk. We'll only generate 50 steps to keep the animation short, but we easily could have generated a lot more.

The FuncAnimation object repeatedly calls update_plot with the current frame number. The frame number starts at $0$ and counts up to num_steps. At each step, we'll use the frame number to index into the points generated by get_random_walk.

In [6]:
num_steps = 50
all_steps = get_random_walk(num_steps)

def update_plot(step_num: int):
    loc = all_steps[..., step_num]
    
    previous_steps.set_data(all_steps[..., :step_num])
    current_step.set_data(loc)
    
    loc_text.set_text(
        f"Location = ({loc[0]:0.2f}, {loc[1]:0.2f})"
    )
    distance_text.set_text(
        f"Distance from origin = {np.linalg.norm(loc):0.2f}"
    )
    
    return previous_steps, current_step, loc_text, distance_text

The last step is to create the FuncAnimation object which takes:

  1. The plt.figure() object
  2. The function that updates the plot at each frame
  3. The number of frames to run the animation1
  4. The number of milliseconds to pause between frames
In [7]:
walk = animation.FuncAnimation(
    fig,
    update_plot,
    num_steps,
    interval=500,
)

walk
Out[7]:

Notes

  1. The frames argument to FuncAnimation confused me at first because I had seen it used in a few different ways and I hadn't read the documentation closely. That argument is meant to be a generator which yields values that get passed to the update_plot function. When an integer is passed (like we have done), it's converted into range(num_steps). Our FuncAnimation object could have been created with:
walk = animation.FuncAnimation(
    fig,
    update_plot,
    range(num_steps),
    interval=500,
)

We also could have turned get_random_walk into a generator that yields a new step each time it's called. Rather than indexing into previously generated data, the update_plot function would take a new step location and plot it along with the steps generated in previous frames.

python jupyter matplotlib