Updated visualize code

Also fixed issue with saving files after a
number of iterations.