Fix plots

Plots were being added to multiple output images (figure() was not called every time)

Also simplifies plotting with new functions