python - How to animate a scatter plot

Question

I'm trying to do an animation of a scatter plot where colors and size of the points changes at different stage of the animation. For data I have two numpy ndarray with an x value and y value:

data.shape = (ntime, npoint)
x.shape = (npoint)
y.shape = (npoint)

Now I want to plot a scatter plot of the type

pylab.scatter(x,y,c=data[i,:])

and create an animation over the index i. How do I do this?

Answer 1

Suppose you have a scatter plot, scat = ax.scatter(...), then you can

• change the positions

        scat.set_offsets(array)
  

where array is a N x 2 shaped array of x and y coordinates.

• change the sizes

        scat.set_sizes(array)
  

where array is a 1D array of sizes in points.

• change the color

        scat.set_array(array)
  

where array is a 1D array of values which will be colormapped.

Here's a quick example using the animation module.
It's slightly more complex than it has to be, but this should give you a framework to do fancier things.

(Code edited April 2019 to be compatible with current versions. For the older code see revision history)

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

class AnimatedScatter(object):
    """An animated scatter plot using matplotlib.animations.FuncAnimation."""
    def __init__(self, numpoints=50):
        self.numpoints = numpoints
        self.stream = self.data_stream()

        # Setup the figure and axes...
        self.fig, self.ax = plt.subplots()
        # Then setup FuncAnimation.
        self.ani = animation.FuncAnimation(self.fig, self.update, interval=5, 
                                          init_func=self.setup_plot, blit=True)

    def setup_plot(self):
        """Initial drawing of the scatter plot."""
        x, y, s, c = next(self.stream).T
        self.scat = self.ax.scatter(x, y, c=c, s=s, vmin=0, vmax=1,
                                    cmap="jet", edgecolor="k")
        self.ax.axis([-10, 10, -10, 10])
        # For FuncAnimation's sake, we need to return the artist we'll be using
        # Note that it expects a sequence of artists, thus the trailing comma.
        return self.scat,

    def data_stream(self):
        """Generate a random walk (brownian motion). Data is scaled to produce
        a soft "flickering" effect."""
        xy = (np.random.random((self.numpoints, 2))-0.5)*10
        s, c = np.random.random((self.numpoints, 2)).T
        while True:
            xy += 0.03 * (np.random.random((self.numpoints, 2)) - 0.5)
            s += 0.05 * (np.random.random(self.numpoints) - 0.5)
            c += 0.02 * (np.random.random(self.numpoints) - 0.5)
            yield np.c_[xy[:,0], xy[:,1], s, c]

    def update(self, i):
        """Update the scatter plot."""
        data = next(self.stream)

        # Set x and y data...
        self.scat.set_offsets(data[:, :2])
        # Set sizes...
        self.scat.set_sizes(300 * abs(data[:, 2])**1.5 + 100)
        # Set colors..
        self.scat.set_array(data[:, 3])

        # We need to return the updated artist for FuncAnimation to draw..
        # Note that it expects a sequence of artists, thus the trailing comma.
        return self.scat,


if __name__ == '__main__':
    a = AnimatedScatter()
    plt.show()

If you're on OSX and using the OSX backend, you'll need to change blit=True to blit=False in the FuncAnimation initialization below. The OSX backend doesn't fully support blitting. The performance will suffer, but the example should run correctly on OSX with blitting disabled.

---

For a simpler example, which just updates the colors, have a look at the following:

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

def main():
    numframes = 100
    numpoints = 10
    color_data = np.random.random((numframes, numpoints))
    x, y, c = np.random.random((3, numpoints))

    fig = plt.figure()
    scat = plt.scatter(x, y, c=c, s=100)

    ani = animation.FuncAnimation(fig, update_plot, frames=range(numframes),
                                  fargs=(color_data, scat))
    plt.show()

def update_plot(i, data, scat):
    scat.set_array(data[i])
    return scat,

main()