sympy/plotting/plot.py

   1 from sympy import Integer
   2
   3 from threading import RLock
   4
   5 # it is sufficient to import "pyglet" here once
   6 from sympy.thirdparty import import_thirdparty
   7 pyglet = import_thirdparty("pyglet")
   8
   9 from pyglet.gl import *
  10
  11 from plot_object import PlotObject
  12 from plot_axes import PlotAxes
  13 from plot_window import PlotWindow
  14 from plot_mode import PlotMode
  15 import plot_modes
  16
  17 from time import sleep
  18 from os import getcwd, listdir
  19 from util import parse_option_string
  20
  21 from sympy.geometry.entity import GeometryEntity
  22
  23 class Plot(object):
  24     """
  25     Plot Examples
  26     =============
  27
  28     See examples/plotting.py for many more examples.
  29
  30
  31     >>> from sympy import symbols, Plot
  32     >>> x,y,z = symbols('xyz')
  33
  34     >>> Plot(x*y**3-y*x**3)
  35
  36     >>> p = Plot()
  37     >>> p[1] = x*y
  38     >>> p[1].color = z, (0.4,0.4,0.9), (0.9,0.4,0.4)
  39
  40     >>> p = Plot()
  41     >>> p[1] =  x**2+y**2
  42     >>> p[2] = -x**2-y**2
  43
  44
  45     Variable Intervals
  46     ==================
  47
  48     The basic format is [var, min, max, steps], but the
  49     syntax is flexible and arguments left out are taken
  50     from the defaults for the current coordinate mode:
  51
  52     >>> Plot(x**2) # implies [x,-5,5,100]
  53     >>> Plot(x**2, [], []) # [x,-1,1,40], [y,-1,1,40]
  54     >>> Plot(x**2-y**2, [100], [100]) # [x,-1,1,100], [y,-1,1,100]
  55     >>> Plot(x**2, [x,-13,13,100])
  56     >>> Plot(x**2, [-13,13]) # [x,-13,13,100]
  57     >>> Plot(x**2, [x,-13,13]) # [x,-13,13,100]
  58     >>> Plot(1*x, [], [x], mode='cylindrical')
  59     ... # [unbound_theta,0,2*Pi,40], [x,-1,1,20]
  60
  61
  62     Coordinate Modes
  63     ================
  64
  65     Plot supports several curvilinear coordinate modes, and
  66     they independent for each plotted function. You can specify
  67     a coordinate mode explicitly with the 'mode' named argument,
  68     but it can be automatically determined for cartesian or
  69     parametric plots, and therefore must only be specified for
  70     polar, cylindrical, and spherical modes.
  71
  72     Specifically, Plot(function arguments) and Plot[n] =
  73     (function arguments) will interpret your arguments as a
  74     cartesian plot if you provide one function and a parametric
  75     plot if you provide two or three functions. Similarly, the
  76     arguments will be interpreted as a curve is one variable is
  77     used, and a surface if two are used.
  78
  79     Supported mode names by number of variables:
  80
  81     1: parametric, cartesian, polar
  82     2: parametric, cartesian, cylindrical = polar, spherical
  83
  84     >>> Plot(1, mode='spherical')
  85
  86
  87     Calculator-like Interface
  88     =========================
  89
  90     >>> p = Plot(visible=False)
  91     >>> f = x**2
  92     >>> p[1] = f
  93     >>> p[2] = f.diff(x)
  94     >>> p[3] = f.diff(x).diff(x)
  95     >>> p
  96     [1]: x**2, 'mode=cartesian'
  97     [2]: 2*x, 'mode=cartesian'
  98     [3]: 2, 'mode=cartesian'
  99     >>> p.show()
 100     >>> p.clear()
 101     >>> p
 102     <blank plot>
 103     >>> p[1] =  x**2+y**2
 104     >>> p[1].style = 'solid'
 105     >>> p[2] = -x**2-y**2
 106     >>> p[2].style = 'wireframe'
 107     >>> p[1].color = z, (0.4,0.4,0.9), (0.9,0.4,0.4)
 108     >>> p[1].style = 'both'
 109     >>> p[2].style = 'both'
 110     >>> p.close()
 111
 112
 113     Plot Window Keyboard Controls
 114     =============================
 115
 116     Screen Rotation:
 117         X,Y axis      Arrow Keys, A,S,D,W, Numpad 4,6,8,2
 118         Z axis        Q,E, Numpad 7,9
 119
 120     Model Rotation:
 121         Z axis        Z,C, Numpad 1,3
 122
 123     Zoom:             R,F, PgUp,PgDn, Numpad +,-
 124
 125     Reset Camera:     X, Numpad 5
 126
 127     Camera Presets:
 128         XY            F1
 129         XZ            F2
 130         YZ            F3
 131         Perspective   F4
 132
 133     Sensitivity Modifier: SHIFT
 134
 135     Axes Toggle:
 136         Visible       F5
 137         Colors        F6
 138
 139     Close Window:     ESCAPE
 140
 141     =============================
 142     """
 143
 144     def __init__(self, *fargs, **win_args):
 145         """
 146         Positional Arguments
 147         ====================
 148
 149         Any given positional arguments are used to
 150         initialize a plot function at index 1. In
 151         other words...
 152
 153         >>> from sympy.core import Symbol
 154         >>> x = Symbol('x')
 155         >>> p = Plot(x**2, visible=False)
 156
 157         ...is equivalent to...
 158
 159         >>> p = Plot(visible=False)
 160         >>> p[1] = x**2
 161
 162         Note that in earlier versions of the plotting
 163         module, you were able to specify multiple
 164         functions in the initializer. This functionality
 165         has been dropped in favor of better automatic
 166         plot plot_mode detection.
 167
 168
 169         Named Arguments
 170         ===============
 171
 172         axes
 173             An option string of the form
 174             "key1=value1; key2 = value2" which
 175             can use the following options:
 176
 177             style = ordinate
 178                 none OR frame OR box OR ordinate
 179
 180             stride = 0.25
 181                 val OR (val_x, val_y, val_z)
 182
 183             overlay = True (draw on top of plot)
 184                 True OR False
 185
 186             colored = False (False uses Black,
 187                              True uses colors
 188                              R,G,B = X,Y,Z)
 189                 True OR False
 190
 191             label_axes = False (display axis names
 192                                 at endpoints)
 193                 True OR False
 194
 195         visible = True (show immediately
 196             True OR False
 197
 198
 199         The following named arguments are passed as
 200         arguments to window initialization:
 201
 202         antialiasing = True
 203             True OR False
 204
 205         ortho = False
 206             True OR False
 207
 208         invert_mouse_zoom = False
 209             True OR False
 210
 211         """
 212         self._win_args = win_args
 213         self._window = None
 214
 215         self._render_lock = RLock()
 216
 217         self._functions = {}
 218         self._pobjects = []
 219         self._screenshot = ScreenShot(self)
 220
 221         axe_options = parse_option_string(win_args.pop('axes', ''))
 222         self.axes = PlotAxes(**axe_options)
 223         self._pobjects.append(self.axes)
 224
 225         self[0] = fargs
 226         if win_args.get('visible', True):
 227             self.show()
 228
 229     ## Window Interfaces
 230
 231     def show(self):
 232         """
 233         Creates and displays a plot window, or activates it
 234         (gives it focus) if it has already been created.
 235         """
 236         if self._window and not self._window.has_exit:
 237             self._window.activate()
 238         else:
 239             self._win_args['visible'] = True
 240             self.axes.reset_resources()
 241             self._window = PlotWindow(self, **self._win_args)
 242
 243     def close(self):
 244         """
 245         Closes the plot window.
 246         """
 247         if self._window:
 248             self._window.close()
 249
 250     def saveimage(self, outfile=None, format='', size=(600, 500)):
 251         """
 252         Saves a screen capture of the plot window to an
 253         image file.
 254
 255         If outfile is given, it can either be a path
 256         or a file object. Otherwise a png image will
 257         be saved to the current working directory.
 258         If the format is omitted, it is determined from
 259         the filename extension.
 260         """
 261         self._screenshot.save(outfile, format, size)
 262
 263     ## Function List Interfaces
 264
 265     def clear(self):
 266         """
 267         Clears the function list of this plot.
 268         """
 269         self._render_lock.acquire()
 270         self._functions = {}
 271         self.adjust_all_bounds()
 272         self._render_lock.release()
 273
 274     def __getitem__(self, i):
 275         """
 276         Returns the function at position i in the
 277         function list.
 278         """
 279         return self._functions[i]
 280
 281     def __setitem__(self, i, args):
 282         """
 283         Parses and adds a PlotMode to the function
 284         list.
 285         """
 286         if not (isinstance(i, (int, Integer)) and i >= 0):
 287             raise ValueError("Function index must "
 288                              "be an integer >= 0.")
 289
 290         if isinstance(args, PlotObject):
 291             f = args
 292         else:
 293             if (not isinstance(args, (list, tuple))) or isinstance(args, GeometryEntity):
 294                 args = [args]
 295             if len(args) == 0:
 296                 return # no arguments given
 297             kwargs = dict(bounds_callback=self.adjust_all_bounds)
 298             f = PlotMode(*args, **kwargs)
 299
 300         if f:
 301             self._render_lock.acquire()
 302             self._functions[i] = f
 303             self._render_lock.release()
 304         else:
 305             raise ValueError("Failed to parse '%s'."
 306                     % ', '.join(str(a) for a in args))
 307
 308     def __delitem__(self, i):
 309         """
 310         Removes the function in the function list at
 311         position i.
 312         """
 313         self._render_lock.acquire()
 314         del self._functions[i]
 315         self.adjust_all_bounds()
 316         self._render_lock.release()
 317
 318     def firstavailableindex(self):
 319         """
 320         Returns the first unused index in the function list.
 321         """
 322         i = 0
 323         self._render_lock.acquire()
 324         while i in self._functions: i += 1
 325         self._render_lock.release()
 326         return i
 327
 328     def append(self, *args):
 329         """
 330         Parses and adds a PlotMode to the function
 331         list at the first available index.
 332         """
 333         self.__setitem__(self.firstavailableindex(), args)
 334
 335     def __len__(self):
 336         """
 337         Returns the number of functions in the function list.
 338         """
 339         return len(self._functions)
 340
 341     def __iter__(self):
 342         """
 343         Allows iteration of the function list.
 344         """
 345         return self._functions.itervalues()
 346
 347     def __repr__(self):
 348         return str(self)
 349
 350     def __str__(self):
 351         """
 352         Returns a string containing a new-line separated
 353         list of the functions in the function list.
 354         """
 355         s = ""
 356         if len(self._functions) == 0:
 357             s += "<blank plot>"
 358         else:
 359             self._render_lock.acquire()
 360             s += "\n".join(["%s[%i]: %s" % ("", i, str(self._functions[i]))
 361                               for i in self._functions])
 362             self._render_lock.release()
 363         return s
 364
 365     def adjust_all_bounds(self):
 366         self._render_lock.acquire()
 367         self.axes.reset_bounding_box()
 368         for f in self._functions:
 369             self.axes.adjust_bounds(self._functions[f].bounds)
 370         self._render_lock.release()
 371
 372     def wait_for_calculations(self):
 373         sleep(0)
 374         self._render_lock.acquire()
 375         for f in self._functions:
 376             a = self._functions[f]._get_calculating_verts
 377             b = self._functions[f]._get_calculating_cverts
 378             while a() or b(): sleep(0)
 379         self._render_lock.release()
 380
 381
 382 class ScreenShot:
 383     def __init__(self, plot):
 384         self._plot = plot
 385         self.screenshot_requested = False
 386         self.outfile = None
 387         self.format = ''
 388         self.invisibleMode = False
 389         self.flag = 0
 390
 391     def __nonzero__(self):
 392         if self.screenshot_requested:
 393             return 1
 394         return 0
 395
 396     def _execute_saving(self):
 397         if self.flag <3:
 398             self.flag += 1
 399             return
 400
 401         size_x, size_y = self._plot._window.get_size()
 402         size = size_x*size_y*4*sizeof(c_ubyte)
 403         image = create_string_buffer(size)
 404         glReadPixels(0,0,size_x,size_y, GL_RGBA, GL_UNSIGNED_BYTE, image)
 405         from PIL import Image
 406         im = Image.frombuffer('RGBA',(size_x,size_y),image.raw, 'raw', 'RGBA', 0, 1)
 407         if type(self.outfile) in (str, unicode):
 408             im.transpose(Image.FLIP_TOP_BOTTOM).save(self.outfile)
 409         elif type(self.outfile)==file:
 410             im.transpose(Image.FLIP_TOP_BOTTOM).save(self.outfile, self.format)
 411         self.flag = 0
 412         self.screenshot_requested = False
 413         if self.invisibleMode:
 414             self._plot._window.close()
 415
 416
 417     def save(self, outfile=None, format='', size=(600, 500)):
 418         self.outfile = outfile
 419         self.format = format
 420         self.size = size
 421         if not self._plot._window or self._plot._window.has_exit:
 422             self._plot._win_args['visible'] = False
 423
 424             self._plot._win_args['width'] = size[0]
 425             self._plot._win_args['height'] = size[1]
 426
 427             self._plot.axes.reset_resources()
 428             self._plot._window = PlotWindow(self._plot, **self._plot._win_args)
 429             self.invisibleMode = True
 430
 431         if type(self.outfile) in (str, unicode):
 432             self.screenshot_requested = True
 433         elif type(self.outfile)==file and self.format:
 434             self.screenshot_requested = True
 435         elif self.outfile==None:
 436             self.outfile=self._create_unique_path()
 437             self.screenshot_requested = True
 438             print self.outfile
 439
 440     def _create_unique_path(self):
 441         cwd = getcwd()
 442         l = listdir(cwd)
 443         path = ''
 444         i=0
 445         while True:
 446             if not 'plot_%s.png'%i in l:
 447                 path = cwd+'/plot_%s.png'%i
 448                 break
 449             i+=1
 450         return path