from Numeric import *
from plplot import *

NS = 20
NX = 35
NY = 46
PERIMETERPTS = 100

XSPA = 2./(NX-1)
YSPA = 2./(NY-1)

tr = array((XSPA, 0.0, -1.0, 0.0, YSPA, -1.0))

def mypltr(x, y, data):
    result0 = data[0] * x + data[1] * y + data[2]
    result1 = data[3] * x + data[4] * y + data[5]
    return array((result0, result1))
		
def main():

    fill_width = 2
    cont_color = 2
    cont_width = 3

    # Set up data array

    x = (arrayrange(NX) - (NX/2)) / float(NX/2)
    x.shape = (-1,1)
    y = (arrayrange(NY) - (NY/2)) / float(NY/2) - 1.
    zz = -sin(7.*x) * cos(7.*y) + x*x - y*y
    ww = -cos(7.*x) * sin(7.*y) + 2.*x*y
    
    zmin = min(zz.flat)
    zmax = max(zz.flat)

    clevel = zmin + (zmax - zmin) * (arrayrange(NS)+0.5)/NS
    shedge = zmin + (zmax - zmin) * (arrayrange(NS+1))/NS

    # Build the identity transformation between grid and world coordinates
    # using mypltr.
    # Note that *for the given* tr, mypltr(i,j,tr)[0] is only a function of i
    # and mypltr(i,j,tr)[1] is only function of j.
    xg0 = mypltr(arange(NX),0,tr)[0]
    yg0 = mypltr(0,arange(NY),tr)[1]
	    
    # Build the 1-d coord arrays.

    distort = .4
    
    xx =  xg0
    xg1 = xx + distort * cos( .5 * pi * xx )
    yy =  yg0
    yg1 = yy - distort * cos( .5 * pi * yy )

    # Build the 2-d coord arrays.
    xx.shape = (-1,1)
    xg2 = xx + distort*cos((0.5*pi)*xx)*cos((0.5*pi)*yy)
    yg2 = yy - distort*cos((0.5*pi)*xx)*cos((0.5*pi)*yy)

    # Plot using identity transform

    pladv(0)
    plvpor(0.1, 0.9, 0.1, 0.9)
    plwind(-1.0, 1.0, -1.0, 1.0)
    plpsty(0)

#   Note another alternative to produce the identical result in a different way
#   is to use the command:
#   plshades(zz, -1.0, 1.0, -1.0, 1.0, shedge, fill_width, 1)
#   The above command works because xmin, xmax, ymin, and ymax do an effective
#   linear transformation on the index ranges.  (We could have dropped
#   xmin, xmax, ymin, ymax since the defaults are -1.0, 1.0, -1.0, 1.0, but
#   for pedagogical reasons we left them in.)  The alternative below using 
#   mypltr does the exact same linear transformation because of the way that
#   the tr array has been defined.  Note that when pltr and pltr_data are 
#   defined, xmin, xmax, ymin, ymax are completely ignored so we can drop 
#   them from the argument list.
    plshades(zz, shedge, fill_width, 1, mypltr, tr)

    plcol0(1)
    plbox( "bcnst", 0., 0, "bcnstv", 0., 0 )
    plcol0(2)

    pllab( "distance", "altitude", "Bogon density" )
    
    # Plot using 1d coordinate transform

    pladv(0)
    plvpor(0.1, 0.9, 0.1, 0.9)
    plwind(-1.0, 1.0, -1.0, 1.0)
    plpsty(0)

    plshades(zz, shedge, fill_width, 1, pltr1, xg1, yg1)

    plcol0(1)
    plbox( "bcnst", 0.0, 0, "bcnstv", 0.0, 0 )
    plcol0(2)
	       
    pllab( "distance", "altitude", "Bogon density" )

    # Plot using 2d coordinate transform

    pladv(0)
    plvpor(0.1, 0.9, 0.1, 0.9)
    plwind(-1.0, 1.0, -1.0, 1.0)
    plpsty(0)

    plshades(zz, shedge, fill_width, 0, pltr2, xg2, yg2)

    plcol0(1)
    plbox( "bcnst", 0.0, 0, "bcnstv", 0.0, 0 )
    plcol0(2)

    plcont(ww, clevel, pltr2, xg2, yg2)
    pllab( "distance", "altitude", "Bogon density, with streamlines" )

    # Plot using 2d coordinate transform (with contours generated by plshades)

    pladv(0)
    plvpor(0.1, 0.9, 0.1, 0.9)
    plwind(-1.0, 1.0, -1.0, 1.0)
    plpsty(0)

    # Note default cont_color and cont_width are zero so that no contours
    # are done with other calls to plshades.  But for this call we specify
    # non-zero values so that contours are drawn.
    plshades(zz, shedge, fill_width, cont_color, cont_width,\
    0, pltr2, xg2, yg2)

    plcol0(1)
    plbox( "bcnst", 0.0, 0, "bcnstv", 0.0, 0 )
    plcol0(2)

#    plcont(ww, clevel, "pltr2", xg2, yg2, 0)
    pllab( "distance", "altitude", "Bogon density" )

    # Example with polar coordinates that demonstrates python wrap support.

    pladv(0)
    plvpor(0.1, 0.9, 0.1, 0.9)
    plwind(-1.0, 1.0, -1.0, 1.0)

    # Build new coordinate matrices.
    
    r = arrayrange(NX)/(NX-1.)
    r.shape = (-1,1)
    t = (2.*pi/(NY-1.))*arrayrange(NY-1)
    xg = r*cos(t)
    yg = r*sin(t)
    z = exp(-r*r)*cos(5.*pi*r)*cos(5.*t)

    # Need a new shedge to go along with the new data set.

    zmin = min(z.flat)
    zmax = max(z.flat)
    shedge = zmin + ((zmax - zmin)/NS) * (arrayrange(NS+1))

    plpsty(0)

    # Now we can shade the interior region.  Use wrap=2 to simulate additional
    # point at t = 2 pi.
    plshades(z, shedge, fill_width, 0, pltr2, xg, yg, 2)

    # Now we can draw the perimeter.  (If do before, plshades may overlap.)
    t = 2.*pi*arrayrange(PERIMETERPTS)/(PERIMETERPTS-1.)
    px = cos(t)
    py = sin(t)


    plcol0(1)
    plline( px, py )

    # And label the plot.

    plcol0(2)
    pllab( "", "",  "Tokamak Bogon Instability" )

    # Restore defaults
    plcol0(1)

main()