Saving scatterplot animations

I've been trying to save an animated scatterplot with matplotlib, and I would prefer that it didn't require totally different code for viewing as an animated figure and for saving a copy. The figure shows all the datapoints perfectly after the save completes.

This code is a modified version of Giggi's on Animating 3d scatterplot in matplotlib, with a fix for the colors from Yann's answer on Matplotlib 3D scatter color lost after redraw (because colors will be important on my video, so I want to make sure they work).

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

FLOOR = -10
CEILING = 10

class AnimatedScatter(object):
    def __init__(self, numpoints=5):
        self.numpoints = numpoints
        self.stream = self.data_stream()
        self.angle = 0

        self.fig = plt.figure()
        self.fig.canvas.mpl_connect('draw_event',self.forceUpdate)
        self.ax = self.fig.add_subplot(111,projection = '3d')
        self.ani = animation.FuncAnimation(self.fig, self.update, interval=100, 
                                       init_func=self.setup_plot, blit=True,frames=20)

    def change_angle(self):
        self.angle = (self.angle + 1)%360

    def forceUpdate(self, event):
        self.scat.changed()

    def setup_plot(self):
        X = next(self.stream)
        c = ['b', 'r', 'g', 'y', 'm']
        self.scat = self.ax.scatter(X[:,0], X[:,1], X[:,2] , c=c, s=200, animated=True)

        self.ax.set_xlim3d(FLOOR, CEILING)
        self.ax.set_ylim3d(FLOOR, CEILING)
        self.ax.set_zlim3d(FLOOR, CEILING)

        return self.scat,

    def data_stream(self):
        data = np.zeros(( self.numpoints , 3 ))
        xyz = data[:,:3]
        while True:
            xyz += 2 * (np.random.random(( self.numpoints,3)) - 0.5)
            yield data

    def update(self, i):
        data = next(self.stream)
        #data = np.transpose(data)

        self.scat._offsets3d = ( np.ma.ravel(data[:,0]) , np.ma.ravel(data[:,1]) , np.ma.ravel(data[:,2]) )

        plt.draw()
        return self.scat,

    def show(self):
        plt.show()

if __name__ == '__main__':
    a = AnimatedScatter()
    a.ani.save("movie.avi", codec='avi')
    a.show()

A perfectly valid .avi is generated by this, but it's blank for all of the four seconds except for the axes. The actual figure always shows exactly what I want to see. How can I populate the save function's plots the same way I populate a normally running animation, or is it possible in matplotlib?

EDIT: Using a scatter call in the update (without setting the bounds as in the initializer) causes the .avi to show the axes growing, showing that the data is being run each time, it's just not showing on the video itself. I am using matplotlib 1.1.1rc with Python 2.7.3.

import matplotlib.pyplot as plt import matplotlib.animation as animation import numpy as np from mpl_toolkits.mplot3d import Axes3D FLOOR = -10 CEILING = 10 class AnimatedScatter(object): def __init__(self, numpoints=5): self.numpoints = numpoints self.stream = self.data_stream() self.angle = 0 self.fig = plt.figure() self.fig.canvas.mpl_connect('draw_event',self.forceUpdate) self.ax = self.fig.add_subplot(111,projection = '3d') self.ani = animation.FuncAnimation(self.fig, self.update, interval=100, init_func=self.setup_plot, frames=20) def change_angle(self): self.angle = (self.angle + 1)%360 def forceUpdate(self, event): self.scat.changed() def setup_plot(self): X = next(self.stream) c = ['b', 'r', 'g', 'y', 'm'] self.scat = self.ax.scatter(X[:,0], X[:,1], X[:,2] , c=c, s=200) self.ax.set_xlim3d(FLOOR, CEILING) self.ax.set_ylim3d(FLOOR, CEILING) self.ax.set_zlim3d(FLOOR, CEILING) return self.scat, def data_stream(self): data = np.zeros(( self.numpoints , 3 )) xyz = data[:,:3] while True: xyz += 2 * (np.random.random(( self.numpoints,3)) - 0.5) yield data def update(self, i): data = next(self.stream) self.scat._offsets3d = ( np.ma.ravel(data[:,0]) , np.ma.ravel(data[:,1]) , np.ma.ravel(data[:,2]) ) return self.scat, def show(self): plt.show() if __name__ == '__main__': a = AnimatedScatter() a.ani.save("movie.avi", codec='avi') a.show()

Recommended topics

Hot tags