Updated visualization code.