Tidied up the plotting classes