Cleaned up plotting code a bit