Animate scatter plot with colorbar using matplotlib

问题

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)

回答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:

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)
    # 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)

plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.grid(b=None)
ani = animation.FuncAnimation(fig, animate, interval=100, frames=range(time_steps))

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

Producing the following GIF:

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.

来源：https://stackoverflow.com/questions/63453435/animate-scatter-plot-with-colorbar-using-matplotlib

标签

python

matplotlib

animation