Added new code for plotting