Python, QT and matplotlib scatter plots with blitting

Question

I am trying to animate a scatter plot (it needs to be a scatter plot as I want to vary the circle sizes). I have gotten the matplotlib documentation tutorial matplotlib documentation tutorial to work in my PyQT application, but would like to introduce blitting into the equation as my application will likely run on slower machines where the animation may not be as smooth.

I have had a look at many examples of animations with blitting, but none ever use a scatter plot (they use plot or lines) and so I am really struggling to figure out how to initialise the animation (the bits that don't get re-rendered every time) and the ones that do. I have tried quite a few things, and seem to be getting nowhere (and I am sure they would cause more confusion than help!). I assume that I have missed something fairly fundamental. Has anyone done this before? Could anyone help me out splitting the figure into the parts that need to be initiated and the ones that get updates?

The code below works, but does not blit. Appending

blit=True

to the end of the animation call yields the following error:

RuntimeError: The animation function must return a sequence of Artist objects.

Any help would be great.

Regards

FP

import numpy as np
from PyQt4 import QtGui, uic
import sys
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

class MainWindow(QtGui.QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()

        self.setupAnim()

        self.show()

    def setupAnim(self):
        self.fig = plt.figure(figsize=(7, 7))
        self.ax = self.fig.add_axes([0, 0, 1, 1], frameon=False)
        self.ax.set_xlim(0, 1), self.ax.set_xticks([])
        self.ax.set_ylim(0, 1), self.ax.set_yticks([])

        # Create rain data
        self.n_drops = 50
        self.rain_drops = np.zeros(self.n_drops, dtype=[('position', float, 2),
                                              ('size',     float, 1),
                                              ('growth',   float, 1),
                                              ('color',    float, 4)])

        # Initialize the raindrops in random positions and with
        # random growth rates.
        self.rain_drops['position'] = np.random.uniform(0, 1, (self.n_drops, 2))
        self.rain_drops['growth'] = np.random.uniform(50, 200, self.n_drops)

        # Construct the scatter which we will update during animation
        # as the raindrops develop.
        self.scat = self.ax.scatter(self.rain_drops['position'][:, 0], self.rain_drops['position'][:, 1],
                          s=self.rain_drops['size'], lw=0.5, edgecolors=self.rain_drops['color'],
                          facecolors='none')

        self.animation = FuncAnimation(self.fig, self.update, interval=10)
        plt.show()

    def update(self, frame_number):
        # Get an index which we can use to re-spawn the oldest raindrop.
        self.current_index = frame_number % self.n_drops

        # Make all colors more transparent as time progresses.
        self.rain_drops['color'][:, 3] -= 1.0/len(self.rain_drops)
        self.rain_drops['color'][:, 3] = np.clip(self.rain_drops['color'][:, 3], 0, 1)

        # Make all circles bigger.
        self.rain_drops['size'] += self.rain_drops['growth']

        # Pick a new position for oldest rain drop, resetting its size,
        # color and growth factor.
        self.rain_drops['position'][self.current_index] = np.random.uniform(0, 1, 2)
        self.rain_drops['size'][self.current_index] = 5
        self.rain_drops['color'][self.current_index] = (0, 0, 0, 1)
        self.rain_drops['growth'][self.current_index] = np.random.uniform(50, 200)

        # Update the scatter collection, with the new colors, sizes and positions.
        self.scat.set_edgecolors(self.rain_drops['color'])
        self.scat.set_sizes(self.rain_drops['size'])
        self.scat.set_offsets(self.rain_drops['position'])

if __name__== '__main__':
    app = QtGui.QApplication(sys.argv)
    window = MainWindow()
    sys.exit(app.exec_())

Answer 1

You need to add return self.scat, at the end of the update method if you want to use FuncAnimation with blit=True. See also this nice StackOverflow post that presents an example of a scatter plot animation with matplotlib using blit.

As a side-note, if you wish to embed a mpl figure in a Qt application, it is better to avoid using the pyplot interface and to use instead the Object Oriented API of mpl as suggested in the matplotlib documentation.

This could be achieved, for example, as below, where mplWidget can be embedded as any other Qt widget in your main application. Note that I renamed the update method to update_plot to avoid conflict with the already existing method of the FigureCanvasQTAgg class.

import numpy as np
from PyQt4 import QtGui
import sys
import matplotlib as mpl
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt

class mplWidget(FigureCanvasQTAgg):
    def __init__(self):
        super(mplWidget, self).__init__(mpl.figure.Figure(figsize=(7, 7)))

        self.setupAnim()
        self.show()

    def setupAnim(self):
        ax = self.figure.add_axes([0, 0, 1, 1], frameon=False)
        ax.axis([0, 1, 0, 1])
        ax.axis('off')

        # Create rain data
        self.n_drops = 50
        self.rain_drops = np.zeros(self.n_drops, dtype=[('position', float, 2),
                                                        ('size',     float, 1),
                                                        ('growth',   float, 1),
                                                        ('color',    float, 4)
                                                        ])

        # Initialize the raindrops in random positions and with
        # random growth rates.
        self.rain_drops['position'] = np.random.uniform(0, 1, (self.n_drops, 2))
        self.rain_drops['growth'] = np.random.uniform(50, 200, self.n_drops)

        # Construct the scatter which we will update during animation
        # as the raindrops develop.
        self.scat = ax.scatter(self.rain_drops['position'][:, 0],
                               self.rain_drops['position'][:, 1],
                               s=self.rain_drops['size'],
                               lw=0.5, facecolors='none',
                               edgecolors=self.rain_drops['color'])

        self.animation = FuncAnimation(self.figure, self.update_plot,
                                       interval=10, blit=True)

    def update_plot(self, frame_number):
        # Get an index which we can use to re-spawn the oldest raindrop.
        indx = frame_number % self.n_drops

        # Make all colors more transparent as time progresses.
        self.rain_drops['color'][:, 3] -= 1./len(self.rain_drops)
        self.rain_drops['color'][:, 3] = np.clip(self.rain_drops['color'][:, 3], 0, 1)

        # Make all circles bigger.
        self.rain_drops['size'] += self.rain_drops['growth']

        # Pick a new position for oldest rain drop, resetting its size,
        # color and growth factor.
        self.rain_drops['position'][indx] = np.random.uniform(0, 1, 2)
        self.rain_drops['size'][indx] = 5
        self.rain_drops['color'][indx] = (0, 0, 0, 1)
        self.rain_drops['growth'][indx] = np.random.uniform(50, 200)

        # Update the scatter collection, with the new colors,
        # sizes and positions.
        self.scat.set_edgecolors(self.rain_drops['color'])
        self.scat.set_sizes(self.rain_drops['size'])
        self.scat.set_offsets(self.rain_drops['position'])

        return self.scat,


if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    window = mplWidget()
    sys.exit(app.exec_())