Animate scatter plot with colorbar using matplotlib

Question

I have spent a serious amount of time trying to animate scatter plot where the colour of the marker is defined by a number.

Below is my attempt, which sort of works but not really as planned:

• After each animation step, the old points should be removed. Instead, the new points are simply added to the points already on plot.
• The colorbar should also be updated at each step according to the values (just like time text is).

However, whatever I seem to do, produces a blank chart. Is this really the best I can do with python when it comes to animating scatter?

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

fig = plt.figure()
marker_size = 1
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def init():
    """ Initialize animation. """
    scat = ax.scatter(positions[0][0], positions[0][1], s = marker_size, c = solutions[0], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(scat)
    time_text.set_text('Time step = %d' % 0)

    return scat, time_text

def animate(i):
    """ Perform animation step. """
    scat = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    time_text.set_text('Time step = %d' % i)

    return scat, time_text

plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.grid(b=None)
plt.show()
ani = animation.FuncAnimation(fig, animate, interval=100, blit=True, repeat=True, init_func=init)

ani.save('animation.gif', writer='imagemagick', fps = 8)

Answer 1

I'm not sure if you were indicating this from your post, but I couldn't get your code to run as is. However, I believe the main issue relates to the first point you mention: " After each animation step, the old points should be removed. " You do need to be explicit about this when drawing the animation. Currently, your code is repeatedly creating a scatter for the same Axes . Just as if you were to do this outside of an animation, this will result in multiple sets of data being drawn over each other.

I have seen 2 major ways people do this: either using some set_... methods of the plot to update the data (seen here for scatter plots or here in general ) or clearing the Axes or Figure each iteration in order to plot new data. I find the latter easier/more universal (if lazier). Here is an approach for your example doing so (I've edited this code to remove calls to plt.grid and plt.label , as those were not functional):

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

fig, ax = plt.subplots()
marker_size = 5 #upped this to make points more visible

def animate(i):
    """ Perform animation step. """
    #important - the figure is cleared and new axes are added
    fig.clear()
    ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
    #the new axes must be re-formatted
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)
    ax.grid(b=None)
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    # and the elements for this frame are added
    ax.text(0.02, 0.95, 'Time step = %d' % i, transform=ax.transAxes)
    s = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(s)

ani = animation.FuncAnimation(fig, animate, interval=100, frames=range(time_steps))

ani.save('animation.gif', writer='pillow')

Producing the following GIF:

Here, I use fig.clear() to clear the colorbar each frame; otherwise, many of them will be drawn. This means you have to re-add the Axes and the formatting each time. In other cases, using ax.clear() can be fine and save the step of add_subplot .

There is another way to do this however, following here . If you have the handle for the colorbar Axes , you can just clear them (rather than clearing the entire Figure ), similar to the scatter plot axes:

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

# init the figure, so the colorbar can be initially placed somewhere
marker_size = 5
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
s = ax.scatter(positions[0][0], positions[0][1], s = marker_size, c = solutions[0], cmap = "RdBu_r", marker = ".", edgecolor = None)
cb = fig.colorbar(s)

# get the axis for the colobar
cax = cb.ax

def animate(i):
    """ Perform animation step. """
    # clear both plotting axis and colorbar axis
    ax.clear()
    cax.cla()
    #the new axes must be re-formatted
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)
    ax.grid(b=None)
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    # and the elements for this frame are added
    ax.text(0.02, 0.95, 'Time step = %d' % i, transform=ax.transAxes)
    s = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(s, cax=cax)

ani = animation.FuncAnimation(fig, animate, interval=100, frames=range(time_steps))

ani.save('animation2.gif', writer='pillow')

Producing the same figure.