plot_surface.py 3.7 KB

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