Visualization code bug fixes