Bug fixes and new plot function