Cleaning up code for the new visualization.