moved pylab import to plot function