Added an import matplotlib statement.
Commiting again, sorry I don't have a good description.