updated plotting code