Python 3 compatibility

A tiny change needed to make it work under Python 3.

I haven't modified the demo notebook in this commit, but it's as simple as changing `xrange` to `range`. I can commit that too if you like.