Improvements to plotting