Minor plotting improvments

[mlgill]

Minor improvements to plotting to:

• Allow figure filename to be changed
• Return figure handle so figure can be edited later
• Notebook updated to be compatible with figure handle