Modified code and plots