examples/plotting.py

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