Fixing visualization bugs