changed plotting function slightly, tidied code
new code just for plotting