plot.py 13 KB

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