import numpy
import matplotlib.pyplot as pyplot

def GenData(count) :
    x = numpy.linspace(-10, 10, count)
    X_grid, Y_grid = numpy.meshgrid(x, x)
    return numpy.sqrt(X_grid**2 + Y_grid**2)

def PlotFig(Z) :
    fig = pyplot.figure(figsize=(11, 9))
    print "pcolor rast"
    ax = fig.add_subplot(2, 2, 1)
    ax.pcolor(Z, rasterized=True)
    ax.set_title("pcolor() Rast")

    print "pcolormesh rast"
    ax = fig.add_subplot(2, 2, 2)
    ax.pcolormesh(Z, rasterized=True)
    ax.set_title("pcolormesh() Rast")

    print "pcolor no rast"
    ax = fig.add_subplot(2, 2, 3)
    ax.pcolor(Z, rasterized=False)
    ax.set_title("pcolor() No Rast")

    print "pcolormesh no rast"
    ax = fig.add_subplot(2, 2, 4)
    ax.pcolormesh(Z, rasterized=False)
    ax.set_title("pcolormesh() No Rast")

    return fig

    
if __name__ == "__main__" :

    dpi = 200

    fileprefix = "pcolortest_lores"
    Z = GenData(10)
    fig = PlotFig(Z)
    fig.suptitle("Low Res")

    print "Save PNG"
    fig.savefig(fileprefix + ".png", dpi=dpi)
    print "Save PDF"
    fig.savefig(fileprefix + ".pdf", dpi=dpi)
    print "Save EPS"
    fig.savefig(fileprefix + ".eps", dpi=dpi)


    fileprefix = "pcolortest_hires"
    Z = GenData(100)
    fig = PlotFig(Z)
    fig.suptitle("High Res")

    print "Save PNG"
    fig.savefig(fileprefix + ".png", dpi=dpi)
    print "Save PDF"
    fig.savefig(fileprefix + ".pdf", dpi=dpi)
    print "Save EPS"
    fig.savefig(fileprefix + ".eps", dpi=dpi)