Python Graphing Libraries Benchmark

benchmark.py

# -*- coding: utf-8 -*-

from random import random, uniform
from sys import argv, exit, path
from time import clock

from PySide import QtCore, QtGui

from numpy import arange, mean

HEIGHT, WIDTH = 480, 768
COLOR_PALETTE = (
    QtGui.QColor('#2B2B2B'),
    QtGui.QColor('#162674'),
    QtGui.QColor('#001166'),
    QtGui.QColor('#000022'),
    QtGui.QColor('#99AAFF'),
    QtGui.QColor('#8888AA'),
    QtGui.QColor('#3D3F4C'),
    QtGui.QColor('#222266'),
    QtGui.QColor('#101011'),
)


class DataGen(object):
    def __init__(self, init=50):
        self.data = self.init = init

    def next(self):
        delta = uniform(-0.5, 0.5)
        r = random()

        if r > 0.9:
            self.data += delta * 15
        elif r > 0.8:
            delta += (0.5 if self.init > self.data else -0.5)
            self.data += delta
        else:
            self.data += delta

        return self.data


class Benchmark(object):
    def __init__(self):
        super(Benchmark, self).__init__()

        self.datagens = [DataGen(50 - (10 * i)) for i in xrange(9)]
        self.data     = [[0.0] * 128 for _ in xrange(9)]

        self.curves = []
        
        self.clipped_frames = False
        self.frame_rates    = []
        self.last_draw      = clock()

        self.init()
        
        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.update)
        self.timer.start(0)

    def init(self):
        pass

    def calculate_frame_rate(self):
        frame_rate, self.last_draw = (1.0 / (clock() - self.last_draw)), clock()
        self.frame_rates.append(frame_rate)
        
        if not self.clipped_frames and len(self.frame_rates) > 4:
            self.clipped_frames = True
            self.frame_rates    = self.frame_rates[4:]

        print 'Current: %0.2f fps, Average: %0.2f fps' % (frame_rate, mean(self.frame_rates))

    def generate_data(self, i):
        self.data[i].insert(0, self.datagens[i].next())
        self.data[i].pop()

    def update(self):
        pass

class PyQtGraphBenchmark(Benchmark):
    def init(self):
        from pyqtgraph import mkPen, plot, setConfigOption, ViewBox

        setConfigOption('background', (255,255,255))
        setConfigOption('antialias', True)
        setConfigOption('downsample', True)

        self.plt = plot()
        self.plt.showGrid(True, True)
        self.plt.hideButtons()
        self.plt.setMenuEnabled(False)
        self.plt.setMouseEnabled(False, False)
        self.plt.setRange(xRange=(0,128), padding=0.0)
        self.plt.resize(WIDTH, HEIGHT)
        self.plt.getViewBox().enableAutoRange(ViewBox.YAxis, True)

        self.center_and_raise()
        
        self.curves = [
            self.plt.plot(self.data[i], pen={'color':COLOR_PALETTE[i], 'width': 2})
            for i in xrange(9)
        ]

    def update(self):
        self.calculate_frame_rate()
        
        for i in xrange(9):
            self.generate_data(i)
            self.curves[i].setData(self.data[i])
    
    def center_and_raise(self):
        qr = self.plt.win.frameGeometry()
        cp = QtGui.QDesktopWidget().availableGeometry().center()
        qr.moveCenter(cp)
        self.plt.win.move(qr.topLeft())
        self.plt.win.raise_()


class ChacoBenchmark(Benchmark, QtGui.QWidget):
    def init(self):
        from enable.api import Window
        from traits.api import HasTraits
        from chaco.api  import OverlayPlotContainer, add_default_axes, add_default_grids

        class AnimatedPlot(HasTraits):
            def __init__(self, x, y, color):
                from chaco.api import create_line_plot

                self.x_values, self.y_values = x, y
                self.plot = create_line_plot((self.x_values,self.y_values),
                                                color=color, bgcolor='white', width=2.0)
                self.plot.downsample = True
                self.plot.resizable  = 'hv'
                self.plot.fit_window = True

        self.container = OverlayPlotContainer(
            padding_top=4, padding_right=4, padding_bottom=24, padding_left=32,
            bgcolor='white', use_backbuffer=False, border_visible=False, fill_padding=True
        )

        x = arange(128)
        value_mapper = None
        for i in xrange(9):
            curve = AnimatedPlot(x, self.data[i], COLOR_PALETTE[i]).plot
            
            if value_mapper is None:
                value_mapper = curve.value_mapper
                self.grids = add_default_grids(curve)
                self.axes  = add_default_axes(curve)
            else:
                curve.value_mapper = value_mapper
                value_mapper.range.add(curve.value)
            
            self.container.add(curve)
            self.curves.append(curve)

        self.enable_win = Window(self, -1, component=self.container)

        layout = QtGui.QGridLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(self.enable_win.control)

        self.setLayout(layout)
        self.setMinimumSize(WIDTH, HEIGHT)
        self.show()
        self.raise_()

    def update(self):
        self.calculate_frame_rate()
    
        for i in xrange(9):
            self.generate_data(i)
            self.curves[i].value.set_data(self.data[i])
            self.curves[i].request_redraw()


if __name__ == '__main__':
    #QtGui.QApplication.setGraphicsSystem('native')
    app = QtGui.QApplication(argv)

    #benchmark = PyQtGraphBenchmark()
    benchmark = ChacoBenchmark()

    exit(QtGui.QApplication.exec_())