Updated plotting code