Improved plot code assignment 2