cleaned up the matplotlib animation code