sympy/plotting/plot_surface.py

   1 from pyglet.gl import *
   2 from plot_mode_base import PlotModeBase
   3 from sympy.core.basic import S
   4 from util import scale_value, scale_value_list
   5 #from time import sleep
   6
   7 class PlotSurface(PlotModeBase):
   8
   9     default_rot_preset = 'perspective'
  10
  11     def _on_calculate_verts(self):
  12         self.u_interval = self.intervals[0]
  13         self.u_set = list(self.u_interval.frange())
  14         self.v_interval = self.intervals[1]
  15         self.v_set = list(self.v_interval.frange())
  16         self.bounds = [ [S.Infinity,-S.Infinity,0],[S.Infinity,-S.Infinity,0],[S.Infinity,-S.Infinity,0] ]
  17         evaluate = self._get_evaluator()
  18
  19         self._calculating_verts_pos = 0.0
  20         self._calculating_verts_len = float(self.u_interval.v_len*self.v_interval.v_len)
  21
  22         verts = list()
  23         b = self.bounds
  24         for u in self.u_set:
  25             column = list()
  26             for v in self.v_set:
  27                 try: _e = evaluate(u, v) # calculate vertex
  28                 except: _e = None
  29                 if _e is not None:  # update bounding box
  30                     for axis in xrange(3):
  31                         b[axis][0] = min([b[axis][0], _e[axis]])
  32                         b[axis][1] = max([b[axis][1], _e[axis]])
  33                 column.append(_e)
  34                 self._calculating_verts_pos += 1.0
  35
  36             verts.append(column)
  37         for axis in xrange(3):
  38             b[axis][2] = b[axis][1] - b[axis][0]
  39             if b[axis][2] == 0.0: b[axis][2] = 1.0
  40
  41         self.verts = verts
  42         self.push_wireframe(self.draw_verts(False, False))
  43         self.push_solid(self.draw_verts(False, True))
  44
  45     def _on_calculate_cverts(self):
  46         if not self.verts or not self.color: return
  47         def set_work_len(n): self._calculating_cverts_len = float(n)
  48         def inc_work_pos(): self._calculating_cverts_pos += 1.0
  49         set_work_len(1); self._calculating_cverts_pos = 0
  50         self.cverts = self.color.apply_to_surface(self.verts, self.u_set, self.v_set,
  51                                                   set_len=set_work_len, inc_pos=inc_work_pos)
  52         self.push_solid(self.draw_verts(True, True))
  53
  54     def calculate_one_cvert(self, u, v):
  55         vert = self.verts[u][v]
  56         return self.color(vert[0], vert[1], vert[2],
  57                           self.u_set[u], self.v_set[v])
  58
  59     def draw_verts(self, use_cverts, use_solid_color):
  60         def f():
  61             for u in xrange( 1, len(self.u_set) ):
  62                 glBegin(GL_QUAD_STRIP)
  63                 for v in xrange( len(self.v_set) ):
  64                     pa = self.verts[u-1][v]
  65                     pb = self.verts[u][v]
  66                     if pa is None or pb is None:
  67                         glEnd()
  68                         glBegin(GL_QUAD_STRIP)
  69                         continue
  70                     if use_cverts:
  71                         ca = self.cverts[u-1][v]
  72                         cb = self.cverts[u][v]
  73                         if ca is None:
  74                             ca = (0,0,0)
  75                         if cb is None:
  76                             cb = (0,0,0)
  77                     else:
  78                         if use_solid_color: ca = cb = self.default_solid_color
  79                         else: ca = cb = self.default_wireframe_color
  80                     glColor3f(*ca)
  81                     glVertex3f(*pa)
  82                     glColor3f(*cb)
  83                     glVertex3f(*pb)
  84                 glEnd()
  85         return f
  86