legend animation example

spread.py

#!/usr/bin/env python
from matplotlib import animation
from time import time
from random import randint
import matplotlib.pyplot as plt
import numpy as np

class LiveCounter(object):
	def init(self):
		plt.ylim((0, self.maxy))
		plt.xlim((0, self.bars))
		return list()

	def frames(self, *l, **m):
		while True:
			yield str(randint(0, 5))

	def animate(self, result):
		rolled = False
		if not result in self.counts_per_result:
			if self.legend:
				self.legend.remove()
			label = 'Result {}'.format(result)
			self.counts_per_result[result] = np.zeros(self.bars, int)
			self.bars_per_result[result] = [r for r in plt.bar(
				self.position,
				self.counts_per_result[result],
				align='center',
				color='#ff0000',
				label=label,
				)]
			self.legend = plt.legend(loc='upper center')
		self.counts_per_result[result][-1] += 1
		current = time()
		if current - self.last > self.delta:
			self.last = current
			for k in self.counts_per_result:
				self.counts_per_result[k] = np.roll(self.counts_per_result[k], -1)
				self.counts_per_result[k][-1] = 0
			rolled = True
		bottom = np.zeros(self.bars, int)
		for k in self.counts_per_result:
			if rolled:
				for h, r, b in zip(self.counts_per_result[k], self.bars_per_result[k], bottom):
					r.set_height(h)
					r.set_y(b)
			else:
				h, r, b = self.counts_per_result[k][-1], self.bars_per_result[k][-1], bottom[-1]
				r.set_height(h)
				r.set_y(b)
			bottom += self.counts_per_result[k]
		return [m for n in self.bars_per_result.values() for m in n] + [self.legend, ]

	def main(self,
		bars=15,
		delta=12,
		maxy=500,
	):
		plt.ioff()
		plt.rcParams['toolbar'] = 'None'
		#self.writer = animation.writers['ffmpeg'](fps=15, metadata=dict(artist='Me'), bitrate=1800)
		self.last = time()
		self.legend = None
		self.bars = int(bars)
		self.maxy = int(maxy)
		self.delta = int(delta)
		self.position = np.arange(self.bars) + .5
		self.counts_per_result = dict()
		self.bars_per_result = dict()
		self.figure = plt.figure()
		self.anim = animation.FuncAnimation(
			self.figure,
			self.animate,
			init_func=self.init,
			frames=self.frames(),
			interval=0,
			blit=True,
			fargs=dict(),
			)
		#self.anim.save('/tmp/les_error_distribution.mp4', writer=self.writer)
		self.figure.canvas.draw()
		plt.show()

if __name__ == '__main__':
	from sys import argv
	LiveCounter().main(*argv[1:])