examples/plotting.py

   1 """
   2 Plotting Examples
   3
   4 Note: In Python < 2.5, you will need the ctypes library
   5 to use plotting. It is included with Python 2.5 and later.
   6
   7 Suggested Usage:    python -i plotting.py
   8 """
   9 import sys
  10 sys.path.append("..")
  11
  12 from sympy import symbols
  13 from sympy import Plot
  14 from sympy import sin, cos, pi, sqrt, exp
  15
  16 from time import sleep, clock
  17
  18 if __name__ == "__main__":
  19
  20     x,y,z = symbols('xyz')
  21
  22     # toggle axes visibility with F5, colors with F6
  23     axes_options = 'visible=false; colored=true; label_ticks=true; label_axes=true; overlay=true; stride=0.5'
  24     #axes_options = 'colored=false; overlay=false; stride=(1.0, 0.5, 0.5)'
  25
  26     p = Plot(width=600, height=500, ortho=False, invert_mouse_zoom=False, axes=axes_options, antialiasing=True)
  27
  28     examples = []
  29     def example_wrapper(f):
  30         examples.append(f)
  31         return f
  32
  33     @example_wrapper
  34     def mirrored_saddles():
  35         p[5] = x**2-y**2, [20], [20]
  36         p[6] = y**2-x**2, [20], [20]
  37
  38     @example_wrapper
  39     def mirrored_saddles_saveimage():
  40         p[5] = x**2-y**2, [20], [20]
  41         p[6] = y**2-x**2, [20], [20]
  42         p.wait_for_calculations()
  43         # although the calculation is complete,
  44         # we still need to wait for it to be
  45         # rendered, so we'll sleep to be sure.
  46         sleep(1)
  47         p.saveimage("plot_example.png")
  48
  49     @example_wrapper
  50     def mirrored_ellipsoids():
  51         p[2] =  x**2+y**2, [40], [40], 'color=zfade'
  52         p[3] = -x**2-y**2, [40], [40], 'color=zfade'
  53
  54     @example_wrapper
  55     def saddle_colored_by_derivative():
  56         f = x**2-y**2
  57         p[1] = f, 'style=solid'
  58         p[1].color = abs(f.diff(x)), abs(f.diff(x) + f.diff(y)), abs(f.diff(y))
  59
  60     @example_wrapper
  61     def ding_dong_surface():
  62         f = sqrt(1.0-y)*y
  63         p[1] = f, [x,0,2*Pi,40], [y,-1,4,100], 'mode=cylindrical; style=solid; color=zfade4'
  64
  65     @example_wrapper
  66     def polar_circle():
  67         p[7] = 1, 'mode=polar'
  68
  69     @example_wrapper
  70     def polar_flower():
  71         p[8] = 1.5*sin(4*x), [160], 'mode=polar'
  72         p[8].color = z, x, y, (0.5,0.5,0.5), (0.8,0.8,0.8), (x,y,None,z) # z is used for t
  73
  74     @example_wrapper
  75     def simple_cylinder():
  76         p[9] = 1, 'mode=cylindrical'
  77
  78     @example_wrapper
  79     def cylindrical_hyperbola():
  80         ## (note that polar is an alias for cylindrical)
  81         p[10] = 1/y, 'mode=polar', [x], [y,-2,2,20]
  82
  83     @example_wrapper
  84     def extruded_hyperbolas():
  85         p[11] = 1/x, [x,-10,10,100], [1], 'style=solid'
  86         p[12] = -1/x, [x,-10,10,100], [1], 'style=solid'
  87
  88     @example_wrapper
  89     def torus():
  90         a,b = 1, 0.5 # radius, thickness
  91         p[13] = (a+b*cos(x))*cos(y), (a+b*cos(x))*sin(y), b*sin(x), [x,0,Pi*2,40], [y,0,Pi*2,40]
  92
  93     @example_wrapper
  94     def warped_torus():
  95         a,b = 2, 1 # radius, thickness
  96         p[13] = (a+b*cos(x))*cos(y), (a+b*cos(x))*sin(y), b*sin(x)+0.5*sin(4*y), [x,0,Pi*2,40], [y,0,Pi*2,40]
  97
  98     @example_wrapper
  99     def parametric_spiral():
 100         p[14] = cos(y), sin(y), y/10.0, [y,-4*Pi,4*Pi,100]
 101         p[14].color = x,(0.1,0.9),y,(0.1,0.9),z,(0.1,0.9)
 102
 103     @example_wrapper
 104     def multistep_gradient():
 105         p[1] = 1, 'mode=spherical', 'style=both'
 106         #p[1] = exp(-x**2-y**2+(x*y)/4), [-1.7,1.7,100], [-1.7,1.7,100], 'style=solid'
 107         #p[1] = 5*x*y*exp(-x**2-y**2), [-2,2,100], [-2,2,100]
 108         gradient = [ 0.0,  (0.3, 0.3, 1.0),
 109                      0.30, (0.3, 1.0, 0.3),
 110                      0.55, (0.95,1.0, 0.2),
 111                      0.65, (1.0,0.95, 0.2),
 112                      0.85, (1.0, 0.7, 0.2),
 113                      1.0,  (1.0, 0.3, 0.2) ]
 114         p[1].color = z, [None, None, z], gradient
 115         #p[1].color = 'zfade'
 116         #p[1].color = 'zfade3'
 117
 118     @example_wrapper
 119     def lambda_vs_sympy_evaluation():
 120         start = clock()
 121         p[4] = x**2+y**2, [100], [100], 'style=solid'
 122         p.wait_for_calculations()
 123         print "lambda-based calculation took %s seconds." % (clock()-start)
 124
 125         start = clock()
 126         p[4] = x**2+y**2, [100], [100], 'style=solid; use_sympy_eval'
 127         p.wait_for_calculations()
 128         print "sympy substitution-based calculation took %s seconds." % (clock()-start)
 129
 130     @example_wrapper
 131     def gradient_vectors():
 132         def gradient_vectors_inner(f, i):
 133             from sympy import lambdify
 134             from sympy.plotting.plot_interval import PlotInterval
 135             from pyglet.gl import glBegin, glColor3f
 136             from pyglet.gl import glVertex3f, glEnd, GL_LINES
 137
 138             def draw_gradient_vectors(f, iu, iv):
 139                 """
 140                 Create a function which draws vectors
 141                 representing the gradient of f.
 142                 """
 143                 dx, dy, dz = f.diff(x), f.diff(y), 0
 144                 FF = lambdify( [x,y,f],    [x,y] )
 145                 FG = lambdify( [dx,dy,dz], [x,y] )
 146                 iu.v_steps /= 5
 147                 iv.v_steps /= 5
 148                 Gvl = list(list([FF(u, v), FG(u, v)]
 149                         for v in iv.frange())
 150                         for u in iu.frange())
 151
 152                 def draw_arrow(p1, p2):
 153                     """
 154                     Draw a single vector.
 155                     """
 156                     glColor3f(0.4, 0.4, 0.9)
 157                     glVertex3f(*p1)
 158
 159                     glColor3f(0.9, 0.4, 0.4)
 160                     glVertex3f(*p2)
 161
 162                 def draw():
 163                     """
 164                     Iterate through the calculated
 165                     vectors and draw them.
 166                     """
 167                     glBegin(GL_LINES)
 168                     for u in Gvl:
 169                         for v in u:
 170                             point = [ [v[0][0], v[0][1], v[0][2]],
 171                                       [v[0][0] + v[1][0], v[0][1] + v[1][1], v[0][2] + v[1][2]] ]
 172                             draw_arrow(point[0], point[1])
 173                     glEnd()
 174
 175                 return draw
 176             p[i] = f, [-0.5,0.5,25], [-0.5,0.5,25], 'style=solid'
 177             iu = PlotInterval(p[i].intervals[0])
 178             iv = PlotInterval(p[i].intervals[1])
 179             p[i].postdraw.append(draw_gradient_vectors(f, iu, iv))
 180
 181         gradient_vectors_inner( x**2 + y**2, 1)
 182         gradient_vectors_inner(-x**2 - y**2, 2)
 183
 184     def help_str():
 185         s =  ("\nPlot p has been created. Useful commands: \n"
 186               "    help(p), p[1] = x**2, print p, p.clear() \n\n"
 187               "Available examples (see source in plotting.py):\n\n")
 188         for i in xrange(len(examples)):
 189             s += "(%i) %s\n" % (i, examples[i].__name__)
 190         s += "\n"
 191         s += "e.g. >>> example(2)\n"
 192         s += "     >>> ding_dong_surface()\n"
 193         return s
 194
 195     def example(i):
 196         if callable(i):
 197             p.clear()
 198             i()
 199         elif i >= 0 and i < len(examples):
 200             p.clear()
 201             examples[i]()
 202         else: print "Not a valid example.\n"
 203         print p
 204
 205     #ding_dong_surface()
 206     mirrored_saddles()
 207     #parametric_spiral()
 208     #multistep_gradient()
 209     #gradient_vectors()
 210     #example(0)
 211     print help_str()
 212
 213     #def profile_plotting():
 214         #import cProfile
 215         #from pstats import Stats
 216         #cProfile.run("p.append(1, 'mode=polar')", 'plot.profile2')
 217         #cProfile.run("p.append(x**2+y**2)", 'plot.profile2')
 218         #s = Stats('plot.profile2')
 219         #s.sort_stats('cumulative').print_stats(20)