Permalink
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
289 lines (238 sloc) 9.84 KB
# -*- coding: utf-8 -*-
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
"""
Arrows are a subclass of line visuals, which adds the ability to put several
heads on a line.
"""
from __future__ import division
import numpy as np
from ... import glsl, gloo
from ..transforms._util import as_vec4
from ..visual import Visual
from .line import LineVisual
ARROW_TYPES = (
'stealth',
'curved',
'angle_30',
'angle_60',
'angle_90',
'triangle_30',
'triangle_60',
'triangle_90',
'inhibitor_round'
)
class _ArrowHeadVisual(Visual):
"""Arrow head visual
Several shapes to put on the end of a line. This visual differs from
MarkersVisual in the sense that this visual calculates the orientation of
the visual on the GPU, by calculating the tangent of the line between two
given vertices.
This is not really a visual you would use on your own,
use :class:`ArrowVisual` instead.
Parameters
----------
parent : ArrowVisual
This actual ArrowVisual this arrow head is part of.
"""
ARROWHEAD_VERTEX_SHADER = glsl.get('arrowheads/arrowheads.vert')
ARROWHEAD_FRAGMENT_SHADER = glsl.get('arrowheads/arrowheads.frag')
_arrow_vtype = np.dtype([
('v1', np.float32, 4),
('v2', np.float32, 4),
('size', np.float32, 1),
('color', np.float32, 4),
('linewidth', np.float32, 1)
])
def __init__(self, parent):
Visual.__init__(self, self.ARROWHEAD_VERTEX_SHADER,
self.ARROWHEAD_FRAGMENT_SHADER)
self._parent = parent
self.set_gl_state(depth_test=False, blend=True,
blend_func=('src_alpha', 'one_minus_src_alpha'))
self._draw_mode = 'points'
self._arrow_vbo = gloo.VertexBuffer(
np.array([], dtype=self._arrow_vtype))
def _prepare_transforms(self, view):
xform = view.transforms.get_transform()
view.view_program.vert['transform'] = xform
def _prepare_draw(self, view=None):
if self._parent._arrows_changed:
self._prepare_vertex_data()
self.shared_program.bind(self._arrow_vbo)
self.shared_program['antialias'] = 1.0
self.shared_program.frag['arrow_type'] = self._parent.arrow_type
self.shared_program.frag['fill_type'] = "filled"
def _prepare_vertex_data(self):
arrows = self._parent.arrows
if arrows is None or arrows.size == 0:
self._arrow_vbo = gloo.VertexBuffer(
np.array([], dtype=self._arrow_vtype))
return
v = np.zeros(len(arrows), dtype=self._arrow_vtype)
# 2d // 3d v1 v2.
sh = int(arrows.shape[1] / 2)
v['v1'] = as_vec4(arrows[:, 0:sh])
v['v2'] = as_vec4(arrows[:, sh:int(2 * sh)])
v['size'][:] = self._parent.arrow_size
color, cmap = self._parent._interpret_color(self._parent.arrow_color)
v['color'][:] = color
v['linewidth'][:] = self._parent.width
self._arrow_vbo = gloo.VertexBuffer(v)
class ArrowVisual(LineVisual):
"""Arrow visual
A special line visual which can also draw optional arrow heads at the
specified vertices.
You add an arrow head by specifying two vertices `v1` and `v2` which
represent the arrow body. This visual will draw an arrow head using `v2`
as center point, and the orientation of the arrow head is automatically
determined by calculating the direction vector between `v1` and `v2`.
Parameters
----------
pos : array
Array of shape (..., 2) or (..., 3) specifying vertex coordinates.
color : Color, tuple, or array
The color to use when drawing the line. If an array is given, it
must be of shape (..., 4) and provide one rgba color per vertex.
Can also be a colormap name, or appropriate `Function`.
width:
The width of the line in px. Line widths > 1px are only
guaranteed to work when using 'agg' method.
connect : str or array
Determines which vertices are connected by lines.
* "strip" causes the line to be drawn with each vertex
connected to the next.
* "segments" causes each pair of vertices to draw an
independent line segment
* numpy arrays specify the exact set of segment pairs to
connect.
method : str
Mode to use for drawing.
* "agg" uses anti-grain geometry to draw nicely antialiased lines
with proper joins and endcaps.
* "gl" uses OpenGL's built-in line rendering. This is much faster,
but produces much lower-quality results and is not guaranteed to
obey the requested line width or join/endcap styles.
antialias : bool
Enables or disables antialiasing.
For method='gl', this specifies whether to use GL's line smoothing,
which may be unavailable or inconsistent on some platforms.
arrows : array
A (N, 4) or (N, 6) matrix where each row contains the (x, y) or the
(x, y, z) coordinate of the first and second vertex of the arrow
body. Remember that the second vertex is used as center point for
the arrow head, and the first vertex is only used for determining
the arrow head orientation.
arrow_type : string
Specify the arrow head type, the currently available arrow head types
are:
* stealth
* curved
* triangle_30
* triangle_60
* triangle_90
* angle_30
* angle_60
* angle_90
* inhibitor_round
arrow_size : float
Specify the arrow size
arrow_color : Color, tuple, or array
The arrow head color. If an array is given, it must be of shape
(..., 4) and provide one rgba color per arrow head. Can also be a
colormap name, or appropriate `Function`.
"""
def __init__(self, pos=None, color=(0.5, 0.5, 0.5, 1), width=1,
connect='strip', method='gl', antialias=False, arrows=None,
arrow_type='stealth', arrow_size=None,
arrow_color=(0.5, 0.5, 0.5, 1)):
# Do not use the self._changed dictionary as it gets overwritten by
# the LineVisual constructor.
self._arrows_changed = False
self._arrow_type = None
self._arrow_size = None
self._arrows = None
self.arrow_type = arrow_type
self.arrow_size = arrow_size
self.arrow_color = arrow_color
self.arrow_head = _ArrowHeadVisual(self)
# TODO: `LineVisual.__init__` also calls its own `set_data` method,
# which triggers an *update* event. This results in a redraw. After
# that we call our own `set_data` method, which triggers another
# redraw. This should be fixed.
LineVisual.__init__(self, pos, color, width, connect, method,
antialias)
ArrowVisual.set_data(self, arrows=arrows)
# Add marker visual for the arrow head
self.add_subvisual(self.arrow_head)
def set_data(self, pos=None, color=None, width=None, connect=None,
arrows=None):
"""Set the data used for this visual
Parameters
----------
pos : array
Array of shape (..., 2) or (..., 3) specifying vertex coordinates.
color : Color, tuple, or array
The color to use when drawing the line. If an array is given, it
must be of shape (..., 4) and provide one rgba color per vertex.
Can also be a colormap name, or appropriate `Function`.
width:
The width of the line in px. Line widths > 1px are only
guaranteed to work when using 'agg' method.
connect : str or array
Determines which vertices are connected by lines.
* "strip" causes the line to be drawn with each vertex
connected to the next.
* "segments" causes each pair of vertices to draw an
independent line segment
* numpy arrays specify the exact set of segment pairs to
connect.
arrows : array
A (N, 4) or (N, 6) matrix where each row contains the (x, y) or the
(x, y, z) coordinate of the first and second vertex of the arrow
body. Remember that the second vertex is used as center point for
the arrow head, and the first vertex is only used for determining
the arrow head orientation.
"""
if arrows is not None:
self._arrows = arrows
self._arrows_changed = True
LineVisual.set_data(self, pos, color, width, connect)
@property
def arrow_type(self):
return self._arrow_type
@arrow_type.setter
def arrow_type(self, value):
if value not in ARROW_TYPES:
raise ValueError(
"Invalid arrow type '{}'. Should be one of {}".format(
value, ", ".join(ARROW_TYPES)
)
)
if value == self._arrow_type:
return
self._arrow_type = value
self._arrows_changed = True
@property
def arrow_size(self):
return self._arrow_size
@arrow_size.setter
def arrow_size(self, value):
if value is None:
self._arrow_size = 5.0
else:
if value <= 0.0:
raise ValueError("Arrow size should be greater than zero.")
self._arrow_size = value
self._arrows_changed = True
@property
def arrow_color(self):
return self._arrow_color
@arrow_color.setter
def arrow_color(self, value):
if value is not None:
self._arrow_color = value
self._arrows_changed = True
@property
def arrows(self):
return self._arrows