Animations in matplotlib

Sine curve

Let’s start with a simple plot of a sine curve, with the code that generates it. We’ll take a look at the code line by line afterwards:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
ax.set_xlim(0, 7)
ax.set_ylim(-1, 1)

xdata, ydata = [], []
line, = plt.plot([], [], 'ro')

def update(frame):
    xdata.append(frame)
    ydata.append(np.sin(frame))
    line.set_data(xdata, ydata)
    return line,

ani = FuncAnimation(
    fig,
    update,
    frames=np.linspace(0, 2*np.pi, 128),
    interval=25,
    repeat=False,
)

#with open('sine.html','w') as f:
#    f.write(ani.to_html5_video())

#ani.save('sine.mp4')

plt.show()

This example comes from the matplotlib documentation. Let’s go through it step by step:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

The one new import here is the FuncAnimation object from matplotlib’s animation library. It needs at minimum two arguments: FuncAnimation(fig, update).

fig is the current figure object to draw on, we get that information from plt.subplots():

fig, ax = plt.subplots()
ax.set_xlim(0, 7)
ax.set_ylim(-1, 1)

The other return value ax represents the x- and y-axes in the plot, we use it to set the axis limits here.

If you’re interested in more on subplots(), and how figures and axes work in matplotlib, see the documentation.

Next, we will use plt.plot() to set up an empty plot object that we will fill with data later:

xdata, ydata = [], []
line, = plt.plot([], [], 'ro')

We can use this empty plot to specify the plot formatting. Here we choose red dots (ro). The return value of plt.plot() is a tuple with one element, since we’re only plotting one line. We save that plot line in the variable line.

Next we need to define a function that will be called every time a new frame is needed. It can take one argument with information about the current frame. Here we use that frame as x-values for our function, and calculate the sinus as y-values:

def update(frame):
    xdata.append(frame)
    ydata.append(np.sin(frame))
    line.set_data(xdata, ydata)
    return line,

Every time when update() is called, we append frame to xdata and the sine of frame to ydata. Next we call line.set_data() to update the plot line. The return value of update() tells the animation which objects have changed in the image. For us, this is just the line

Now, we put it all together:

ani = FuncAnimation(
    fig,
    update,
    frames=np.linspace(0, 2*np.pi, 128),
    interval=25,
    repeat=False,
)

• fig: the matplotlib figure we are drawing on

• update: the name of the function to call with every frame

• frames: a list of numbers to pass to update(). The length of this list decides how many frames we have in the animation

• interval: minimum time between frames in milliseconds

• repeat: should the animation loop back to the beginning?

Try experimenting with these settings!

Like before with static images, we have several possibilities of saving the animation. We can save it as an HTML video element:

with open('sine.html','w') as f:
    f.write(ani.to_html5_video())

As an mp4 movie:

ani.save('sine.mp4')

Or, interactively on screen:

plt.show()

Things to try

• change colour or plot symbol

• make the animation longer, multiple oscillations up and down

• experiment with the settings in FuncAnimation

Moving dots

In our second example, we’ll build up an animation of individually moving dots. Let’s start with a minimal example, and build it up from there:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# format is xpos, ypos, xspeed, yspeed
dots = np.array([
    [0.3, 0.7, 0.02, -0.01],
    [0.7, 0.4, -0.02, 0.03],
    [0.1, 0.1, 0.02, 0.02],
])

#position = dots[:,:2]
#speed    = dots[:,2:]
xdata = dots[:,0]
ydata = dots[:,1]
xspeed = dots[:,2]
yspeed = dots[:,3]

graph, = plt.plot([], [], 'bo')

def update(frame):
    #position[:] += speed
    xdata[:] += xspeed
    ydata[:] += yspeed
    graph.set_data(xdata, ydata)
    return graph,

ani = FuncAnimation(
    fig,
    update,
    frames=40,
    interval=50,
    repeat=False,
)

# with open('dots.html','w') as f:
#     f.write(ani.to_html5_video())

# ani.save('dots.mp4')

plt.show()

Our main data structure is a numpy array with 3 lines of 4 values. Each line represents a dot we want to show, and the 4 values are the position in x and y, followed by the velocity in x and y:

# format is xpos, ypos, xspeed, yspeed
dots = np.array([
    [0.3, 0.7, 0.02, -0.01],   # start at 0.3, 0.7, moving right+down
    [0.7, 0.4, -0.02, 0.03],   # start at 0.7, 0.4, moving left+up
    [0.1, 0.1, 0.02, 0.02],    # start at 0.1, 0.1, moving right+up
])

To make this data easier to handle, we can define some useful shortcuts (using numpy slices). We can either take each coordinate separately into a 3x1 array:

xdata = dots[:,0]
ydata = dots[:,1]
xspeed = dots[:,2]
yspeed = dots[:,3]

We can also define a shortcurt for the positions and velocities as 3x2 arrays:

position = dots[:,:2]
speed    = dots[:,2:]

Like before, we set up an empty plot and save it as graph:

graph, = plt.plot([], [], 'bo')

The update function uses set_data like before to update the plot information.

The difference this time is that we update the current positions with the current velocity values. We can either do this on x and y separately:

xdata[:] += xspeed
ydata[:] += yspeed

or in one go:

position[:] += speed

The final bit is FuncAnimation, where we ask for 40 frames to be generated. The result looks like this:

More dots

Right now, we are plotting the movement of 3 points. Let’s try 30 points instead! We need 27 more lines in dots. Doing this by hand gets quite boring, so let’s use np.random instead:

dots = np.random.rand(30,4)

This gives us an array with 30 lines of 4 values. Each value will be between 0 and 1. When we try to run the animation, it looks blank and we see no dots. The velocity values are too large, and all dots are outside the image borders after just one or two steps. Let’s scale down the velocities a bit:

dots[:, 2:] *= 0.01

There’s still a problem with all dots going to the top right. np.random.rand produces only values between 0 and 1. If we want negative numbers, too, we need to shift the velocity values before scaling them:

dots[:, 2:] -= 0.5
dots[:, 2:] *= 0.01

Now the velocities will be chosen between -0.05 and 0.05.

To let the animation continue until the window is closed, remove the frames option in FuncAnimation:

# frames = 40,

Now you should see a cloud of dots slowly walking out of the picture in all directions.

Reflection

Let’s make the dots bounce off the floor. With numpy filters this can be done very conveniently. Inside the update() function, we look for all dots that are currently below the floor:

below_floor = ydata < 0

This is now an array of True/False values. (You can use print() to see what’s going on.)

All dots below the floor get their y-velocity swapped, and their position reset to 0:

def update(frame):

    xdata[:] += xspeed
    ydata[:] += yspeed

    below_floor = ydata < 0
    yspeed[below_floor] *= -1
    ydata[below_floor] = 0

    graph.set_data(xdata, ydata)
    return graph,

The complete code now looks like this:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# format is xpos, ypos, xspeed, yspeed
dots = np.random.rand(30,4)

# velocity should be positive and negative
dots[:, 2:] -= 0.5
# velocity much smaller than position
dots[:, 2:] *= 0.01

position = dots[:,:2]
speed    = dots[:,2:]

xdata = dots[:,0]
ydata = dots[:,1]
xspeed = dots[:,2]
yspeed = dots[:,3]

graph, = plt.plot([], [], 'bo')

def update(frame):

    xdata[:] += xspeed
    ydata[:] += yspeed

    below_floor = ydata < 0
    yspeed[below_floor] *= -1
    ydata[below_floor] = 0

    graph.set_data(xdata, ydata)
    return graph,

ani = FuncAnimation(
    fig,
    update,
    # frames=250,
    interval=30,
    repeat=False,
)

# with open('dots4.html','w') as f:
#    f.write(ani.to_html5_video())

# ani.save('dots.mp4')

plt.show()

Things to try

• change number of dots

• reflection on all walls

• gravitation (hint: change yspeed by a fixed amount in every frame)

• For those who like a challenge: collision between two dots (scipy.spatial.distance.pdist and scipy.spatial.distance.squareform are useful here)