252 lines
8.5 KiB
Python
252 lines
8.5 KiB
Python
|
import pyglet.gl as pgl
|
||
|
|
from pyglet import font
|
||
|
|
|
||
|
|
from sympy.core import S
|
||
|
|
from sympy.plotting.pygletplot.plot_object import PlotObject
|
||
|
|
from sympy.plotting.pygletplot.util import billboard_matrix, dot_product, \
|
||
|
|
get_direction_vectors, strided_range, vec_mag, vec_sub
|
||
|
|
from sympy.utilities.iterables import is_sequence
|
||
|
|
|
||
|
|
|
||
|
|
class PlotAxes(PlotObject):
|
||
|
|
|
||
|
|
def __init__(self, *args,
|
||
|
|
style='', none=None, frame=None, box=None, ordinate=None,
|
||
|
|
stride=0.25,
|
||
|
|
visible='', overlay='', colored='', label_axes='', label_ticks='',
|
||
|
|
tick_length=0.1,
|
||
|
|
font_face='Arial', font_size=28,
|
||
|
|
**kwargs):
|
||
|
|
# initialize style parameter
|
||
|
|
style = style.lower()
|
||
|
|
|
||
|
|
# allow alias kwargs to override style kwarg
|
||
|
|
if none is not None:
|
||
|
|
style = 'none'
|
||
|
|
if frame is not None:
|
||
|
|
style = 'frame'
|
||
|
|
if box is not None:
|
||
|
|
style = 'box'
|
||
|
|
if ordinate is not None:
|
||
|
|
style = 'ordinate'
|
||
|
|
|
||
|
|
if style in ['', 'ordinate']:
|
||
|
|
self._render_object = PlotAxesOrdinate(self)
|
||
|
|
elif style in ['frame', 'box']:
|
||
|
|
self._render_object = PlotAxesFrame(self)
|
||
|
|
elif style in ['none']:
|
||
|
|
self._render_object = None
|
||
|
|
else:
|
||
|
|
raise ValueError(("Unrecognized axes style %s.") % (style))
|
||
|
|
|
||
|
|
# initialize stride parameter
|
||
|
|
try:
|
||
|
|
stride = eval(stride)
|
||
|
|
except TypeError:
|
||
|
|
pass
|
||
|
|
if is_sequence(stride):
|
||
|
|
if len(stride) != 3:
|
||
|
|
raise ValueError("length should be equal to 3")
|
||
|
|
self._stride = stride
|
||
|
|
else:
|
||
|
|
self._stride = [stride, stride, stride]
|
||
|
|
self._tick_length = float(tick_length)
|
||
|
|
|
||
|
|
# setup bounding box and ticks
|
||
|
|
self._origin = [0, 0, 0]
|
||
|
|
self.reset_bounding_box()
|
||
|
|
|
||
|
|
def flexible_boolean(input, default):
|
||
|
|
if input in [True, False]:
|
||
|
|
return input
|
||
|
|
if input in ('f', 'F', 'false', 'False'):
|
||
|
|
return False
|
||
|
|
if input in ('t', 'T', 'true', 'True'):
|
||
|
|
return True
|
||
|
|
return default
|
||
|
|
|
||
|
|
# initialize remaining parameters
|
||
|
|
self.visible = flexible_boolean(kwargs, True)
|
||
|
|
self._overlay = flexible_boolean(overlay, True)
|
||
|
|
self._colored = flexible_boolean(colored, False)
|
||
|
|
self._label_axes = flexible_boolean(label_axes, False)
|
||
|
|
self._label_ticks = flexible_boolean(label_ticks, True)
|
||
|
|
|
||
|
|
# setup label font
|
||
|
|
self.font_face = font_face
|
||
|
|
self.font_size = font_size
|
||
|
|
|
||
|
|
# this is also used to reinit the
|
||
|
|
# font on window close/reopen
|
||
|
|
self.reset_resources()
|
||
|
|
|
||
|
|
def reset_resources(self):
|
||
|
|
self.label_font = None
|
||
|
|
|
||
|
|
def reset_bounding_box(self):
|
||
|
|
self._bounding_box = [[None, None], [None, None], [None, None]]
|
||
|
|
self._axis_ticks = [[], [], []]
|
||
|
|
|
||
|
|
def draw(self):
|
||
|
|
if self._render_object:
|
||
|
|
pgl.glPushAttrib(pgl.GL_ENABLE_BIT | pgl.GL_POLYGON_BIT | pgl.GL_DEPTH_BUFFER_BIT)
|
||
|
|
if self._overlay:
|
||
|
|
pgl.glDisable(pgl.GL_DEPTH_TEST)
|
||
|
|
self._render_object.draw()
|
||
|
|
pgl.glPopAttrib()
|
||
|
|
|
||
|
|
def adjust_bounds(self, child_bounds):
|
||
|
|
b = self._bounding_box
|
||
|
|
c = child_bounds
|
||
|
|
for i in range(3):
|
||
|
|
if abs(c[i][0]) is S.Infinity or abs(c[i][1]) is S.Infinity:
|
||
|
|
continue
|
||
|
|
b[i][0] = c[i][0] if b[i][0] is None else min([b[i][0], c[i][0]])
|
||
|
|
b[i][1] = c[i][1] if b[i][1] is None else max([b[i][1], c[i][1]])
|
||
|
|
self._bounding_box = b
|
||
|
|
self._recalculate_axis_ticks(i)
|
||
|
|
|
||
|
|
def _recalculate_axis_ticks(self, axis):
|
||
|
|
b = self._bounding_box
|
||
|
|
if b[axis][0] is None or b[axis][1] is None:
|
||
|
|
self._axis_ticks[axis] = []
|
||
|
|
else:
|
||
|
|
self._axis_ticks[axis] = strided_range(b[axis][0], b[axis][1],
|
||
|
|
self._stride[axis])
|
||
|
|
|
||
|
|
def toggle_visible(self):
|
||
|
|
self.visible = not self.visible
|
||
|
|
|
||
|
|
def toggle_colors(self):
|
||
|
|
self._colored = not self._colored
|
||
|
|
|
||
|
|
|
||
|
|
class PlotAxesBase(PlotObject):
|
||
|
|
|
||
|
|
def __init__(self, parent_axes):
|
||
|
|
self._p = parent_axes
|
||
|
|
|
||
|
|
def draw(self):
|
||
|
|
color = [([0.2, 0.1, 0.3], [0.2, 0.1, 0.3], [0.2, 0.1, 0.3]),
|
||
|
|
([0.9, 0.3, 0.5], [0.5, 1.0, 0.5], [0.3, 0.3, 0.9])][self._p._colored]
|
||
|
|
self.draw_background(color)
|
||
|
|
self.draw_axis(2, color[2])
|
||
|
|
self.draw_axis(1, color[1])
|
||
|
|
self.draw_axis(0, color[0])
|
||
|
|
|
||
|
|
def draw_background(self, color):
|
||
|
|
pass # optional
|
||
|
|
|
||
|
|
def draw_axis(self, axis, color):
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def draw_text(self, text, position, color, scale=1.0):
|
||
|
|
if len(color) == 3:
|
||
|
|
color = (color[0], color[1], color[2], 1.0)
|
||
|
|
|
||
|
|
if self._p.label_font is None:
|
||
|
|
self._p.label_font = font.load(self._p.font_face,
|
||
|
|
self._p.font_size,
|
||
|
|
bold=True, italic=False)
|
||
|
|
|
||
|
|
label = font.Text(self._p.label_font, text,
|
||
|
|
color=color,
|
||
|
|
valign=font.Text.BASELINE,
|
||
|
|
halign=font.Text.CENTER)
|
||
|
|
|
||
|
|
pgl.glPushMatrix()
|
||
|
|
pgl.glTranslatef(*position)
|
||
|
|
billboard_matrix()
|
||
|
|
scale_factor = 0.005 * scale
|
||
|
|
pgl.glScalef(scale_factor, scale_factor, scale_factor)
|
||
|
|
pgl.glColor4f(0, 0, 0, 0)
|
||
|
|
label.draw()
|
||
|
|
pgl.glPopMatrix()
|
||
|
|
|
||
|
|
def draw_line(self, v, color):
|
||
|
|
o = self._p._origin
|
||
|
|
pgl.glBegin(pgl.GL_LINES)
|
||
|
|
pgl.glColor3f(*color)
|
||
|
|
pgl.glVertex3f(v[0][0] + o[0], v[0][1] + o[1], v[0][2] + o[2])
|
||
|
|
pgl.glVertex3f(v[1][0] + o[0], v[1][1] + o[1], v[1][2] + o[2])
|
||
|
|
pgl.glEnd()
|
||
|
|
|
||
|
|
|
||
|
|
class PlotAxesOrdinate(PlotAxesBase):
|
||
|
|
|
||
|
|
def __init__(self, parent_axes):
|
||
|
|
super().__init__(parent_axes)
|
||
|
|
|
||
|
|
def draw_axis(self, axis, color):
|
||
|
|
ticks = self._p._axis_ticks[axis]
|
||
|
|
radius = self._p._tick_length / 2.0
|
||
|
|
if len(ticks) < 2:
|
||
|
|
return
|
||
|
|
|
||
|
|
# calculate the vector for this axis
|
||
|
|
axis_lines = [[0, 0, 0], [0, 0, 0]]
|
||
|
|
axis_lines[0][axis], axis_lines[1][axis] = ticks[0], ticks[-1]
|
||
|
|
axis_vector = vec_sub(axis_lines[1], axis_lines[0])
|
||
|
|
|
||
|
|
# calculate angle to the z direction vector
|
||
|
|
pos_z = get_direction_vectors()[2]
|
||
|
|
d = abs(dot_product(axis_vector, pos_z))
|
||
|
|
d = d / vec_mag(axis_vector)
|
||
|
|
|
||
|
|
# don't draw labels if we're looking down the axis
|
||
|
|
labels_visible = abs(d - 1.0) > 0.02
|
||
|
|
|
||
|
|
# draw the ticks and labels
|
||
|
|
for tick in ticks:
|
||
|
|
self.draw_tick_line(axis, color, radius, tick, labels_visible)
|
||
|
|
|
||
|
|
# draw the axis line and labels
|
||
|
|
self.draw_axis_line(axis, color, ticks[0], ticks[-1], labels_visible)
|
||
|
|
|
||
|
|
def draw_axis_line(self, axis, color, a_min, a_max, labels_visible):
|
||
|
|
axis_line = [[0, 0, 0], [0, 0, 0]]
|
||
|
|
axis_line[0][axis], axis_line[1][axis] = a_min, a_max
|
||
|
|
self.draw_line(axis_line, color)
|
||
|
|
if labels_visible:
|
||
|
|
self.draw_axis_line_labels(axis, color, axis_line)
|
||
|
|
|
||
|
|
def draw_axis_line_labels(self, axis, color, axis_line):
|
||
|
|
if not self._p._label_axes:
|
||
|
|
return
|
||
|
|
axis_labels = [axis_line[0][::], axis_line[1][::]]
|
||
|
|
axis_labels[0][axis] -= 0.3
|
||
|
|
axis_labels[1][axis] += 0.3
|
||
|
|
a_str = ['X', 'Y', 'Z'][axis]
|
||
|
|
self.draw_text("-" + a_str, axis_labels[0], color)
|
||
|
|
self.draw_text("+" + a_str, axis_labels[1], color)
|
||
|
|
|
||
|
|
def draw_tick_line(self, axis, color, radius, tick, labels_visible):
|
||
|
|
tick_axis = {0: 1, 1: 0, 2: 1}[axis]
|
||
|
|
tick_line = [[0, 0, 0], [0, 0, 0]]
|
||
|
|
tick_line[0][axis] = tick_line[1][axis] = tick
|
||
|
|
tick_line[0][tick_axis], tick_line[1][tick_axis] = -radius, radius
|
||
|
|
self.draw_line(tick_line, color)
|
||
|
|
if labels_visible:
|
||
|
|
self.draw_tick_line_label(axis, color, radius, tick)
|
||
|
|
|
||
|
|
def draw_tick_line_label(self, axis, color, radius, tick):
|
||
|
|
if not self._p._label_axes:
|
||
|
|
return
|
||
|
|
tick_label_vector = [0, 0, 0]
|
||
|
|
tick_label_vector[axis] = tick
|
||
|
|
tick_label_vector[{0: 1, 1: 0, 2: 1}[axis]] = [-1, 1, 1][
|
||
|
|
axis] * radius * 3.5
|
||
|
|
self.draw_text(str(tick), tick_label_vector, color, scale=0.5)
|
||
|
|
|
||
|
|
|
||
|
|
class PlotAxesFrame(PlotAxesBase):
|
||
|
|
|
||
|
|
def __init__(self, parent_axes):
|
||
|
|
super().__init__(parent_axes)
|
||
|
|
|
||
|
|
def draw_background(self, color):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def draw_axis(self, axis, color):
|
||
|
|
raise NotImplementedError()
|