Animations with matplotlib - 2: Random walks

January 31, 2021

This post is the second in a series on how to use the matplotlib.animation module to create animated plots. The examples here use this module to create animations of 2D random walks.

Introduction #

matplotlib¹ is a python library for creating high quality scientific plots. It contains a module matplotlib.animation to create animations. This post shows examples of using this module to create some interesting visualizations of 2D random walks.

Random walks #

Random walks² are models for many physical phenomena, such as the movement of molecules in a liquid or gas, and the diffusion of proteins on the cell membrane. To simulate a random walk, a particle is placed at some initial position $\bold{r}_0 = (x_0, y_0)$ and advanced using the recursion relation:

$$ \bold{r}_{j+1} = (x_j + \delta x, y_j + \delta y) $$

The displacement along each axis is a normally distributed random variable: $$ \delta x, \delta y \sim \cal{N}(0, \sigma) $$

This algorithm produces a Gaussian random walk. The standard deviation $\sigma$ of the normal distribution is related to the diffusion coefficient of the particle³.

In the examples below, particles are initially placed within a small central region of the simulation domain. Each particle moves independently, following the random walk algorithm above. Over time, the aggregated random motion of the particles disperses them like a drop of ink mixing in water.

Examples #

These examples follow a common template:

Set up an empty plot
Initialize with plot elements, eg: points in their initial location
Update plot element attributes, eg: change point coordinates

Step 3 is repeated at each frame to advance the animation. The function FuncAnimation() implements these steps.

First, Load the needed libraries

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.animation as anim
from IPython.display import HTML
from pathlib import Path

Example 1: Diffusing particles #

The first simulation shows the position of particles over time. In the next code block, init() plots the particles in their initial position, move() updates particle positions using the random walk algorithm above, and FuncAnimation() uses these functions as inputs to generate the animation:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


# Initialize particle positions within a central square 
nparticles = 50
r0 = 0.45 + 0.1 * np.random.rand(nparticles, 2)
    
# Create figure
fig, ax = plt.subplots(figsize = (5, 5))
pos = [ax.plot([], [], "o", color = "C3", ms = 6, alpha = 0.8)[0] 
       for j in range(nparticles)]

# Plot initial particle positions
def init():
    for jj, p in enumerate(pos):
        p.set_data(r0[jj, 0], r0[jj, 1])
    # Fix plot appearance
    ax.set_xlim([0, 1])
    ax.set_xticks([])
    ax.set_ylim([0, 1])
    ax.set_yticks([])
    ax.set_title("Example 1: Brownian motion")
    return(pos)

# Update particle positions
def move(k, sigma = 0.01):
    # Apply random normal displacement along each axis
    for p in pos:
        x, y = p.get_data()
        dx, dy = tuple(sigma * np.random.randn(2))
        p.set_data(x + dx, y + dy)
    return(pos)

ani = anim.FuncAnimation(fig, move, init_func = init,
                         frames = 300, interval = 50, blit = True)

The resulting output shows how particles diffuse and disperse away from their initial locations.

Example 2: Particle tracks #

The next example plots the trajectory of each particle. To do this efficiently, first simulate the entire trajectory for each particle, and store all trajectories in a numpy array:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


def simulate_rw(nparticles = 15, nframes = 200, sigma = 0.01):
    """Simulate trajectories of diffusing particles"""
    # Initialize particle positions
    displacements = np.zeros((nframes, 2, nparticles))
    displacements[0, :, :] = 0.45 + 0.1 * np.random.rand(2, nparticles)
    # Simulate Brownian random walk
    displacements[1:, :, :] = np.random.normal(scale = sigma, 
                                               size = (nframes-1, 2, nparticles))
    # Compute positions
    r = np.cumsum(displacements, axis = 0)
    return(r)

Then plot the trajectories by showing the progress up to each frame. Note how the function set_markevery() is used to display both the marker and the line for each track.

12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42


# Generate tracks
positions = simulate_rw()
nframes, _, nparticles = positions.shape

# Create figure
fig, ax = plt.subplots(figsize = (5, 5))
pts = [ax.plot([], [], "-o", alpha = 0.8)[0] 
       for j in range(nparticles)]

# Plot initial particle positions
def init():
    r0 = positions[0,:,:]
    for jj, p in enumerate(pts):
        p.set_data(r0[0, jj], r0[1, jj])
    # Set plot appearance
    ax.set_xlim([0, 1])
    ax.set_xticks([])
    ax.set_ylim([0, 1])
    ax.set_yticks([])
    return(pts)

# Update particle positions
def move(k):
    for jj, p in enumerate(pts):
        r = positions[:k+1, :, jj]
        p.set_data(r[:, 0], r[:, 1])
        p.set_markevery((k, k+1))
    return(pts)

ani = anim.FuncAnimation(fig, move, init_func = init,
                         frames = nframes, interval = 50, blit = True)

The output shows how diffusing particles explore space:

Example 3: Particle trails #

In this example, the animaiton shows a “trail” behind each particle rather than the complete trajectory. The maximum trail length can be changed in the move() function below.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


# Create figure
fig, ax = plt.subplots(figsize = (5, 5))
pts = [ax.plot([], [], "-o", alpha = 0.85)[0] 
       for j in range(nparticles)]

# Plot initial particle positions
def init():
    r0 = positions[0,:,:]
    for jj, p in enumerate(pts):
        p.set_data(r0[0, jj], r0[1, jj])
    # Set plot appearance
    ax.set_xlim([0, 1])
    ax.set_xticks([])
    ax.set_ylim([0, 1])
    ax.set_yticks([])
    return(pts)

# Update particle positions
def move(k, trail = 30):
    for jj, p in enumerate(pts):
        if (k <= trail):
            r = positions[:k+1, :, jj]
            n = k
        else:
            r = positions[(k+1 - trail):k+1, :, jj]
            n = trail-1
        p.set_data(r[:, 0], r[:, 1])
        p.set_markevery((n, n+1))
    return(pts)

ani = anim.FuncAnimation(fig, move, init_func = init,
                         frames = nframes, interval = 50, blit = True)

which produces the following output

Summary and references #

These examples demonstrate how the matplotlib.animation module can be used to create dynamic visualizations. The complete code for this post is available in this jupyter notebook

References:

https://matplotlib.org/ ↩︎
https://en.wikipedia.org/wiki/Random_walk ↩︎
$\sigma = \sqrt{2Dt}$ where $D$ is the diffusion coefficient, and $t$ is the time interval between observation ↩︎