Animate scatter plot with colorbar

Question

I have spent a serious amount of time trying to animate scatter plot where the color 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.