Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
import numpy as np
from .. import gloo
from .visual import Visual
_VERTEX_SHADER = """
attribute vec2 a_pos;
varying vec4 v_color;
void main() {
vec4 pos = vec4(a_pos, 0, 1);
if($is_vertical==1)
{
pos.y = $render_to_visual(pos).y;
}
else
{
pos.x = $render_to_visual(pos).x;
}
gl_Position = $transform(pos);
gl_PointSize = 10.;
v_color = $color;
}
"""
_FRAGMENT_SHADER = """
varying vec4 v_color;
void main() {
gl_FragColor = v_color;
}
"""
class LinearRegionVisual(Visual):
"""Infinite horizontal or vertical region for 2D plots.
Parameters
----------
pos : list, tuple or numpy array
Bounds of the region along the axis. len(pos) must be >=2.
color : list, tuple, or array
The color to use when drawing the line. It must have a shape of
(1, 4) for a single color region or (len(pos), 4) for a multicolor
region.
vertical:
True for drawing a vertical region, False for an horizontal region
"""
_shaders = {
'vertex': _VERTEX_SHADER,
'fragment': _FRAGMENT_SHADER,
}
def __init__(self, pos=None, color=[1.0, 1.0, 1.0, 1.0],
vertical=True, **kwargs):
"""
"""
Visual.__init__(self, vcode=self._shaders['vertex'], fcode=self._shaders['fragment'])
self._changed = {'pos': False, 'color': False}
self.pos_buf = gloo.VertexBuffer()
self.color_buf = gloo.VertexBuffer()
# The Visual superclass contains a MultiProgram, which is an object
# that behaves like a normal shader program (you can assign shader
# code, upload values, set template variables, etc.) but internally
# manages multiple ModularProgram instances, one per view.
# The MultiProgram is accessed via the `shared_program` property, so
# the following modifications to the program will be applied to all
# views:
self.shared_program['a_pos'] = self.pos_buf
self._program.vert['is_vertical'] = 1 if vertical else 0
self._need_upload = False
self._is_vertical = bool(vertical)
self._pos = np.zeros((4, 2), dtype=np.float32)
self._color = np.ones((1, 4), dtype=np.float32)
# Visual keeps track of draw mode, index buffer, and GL state. These
# are shared between all views.
self._draw_mode = 'triangle_strip'
self.set_gl_state('translucent', depth_test=False)
self.set_data(pos=pos, color=color)
def set_data(self, pos=None, color=None):
"""Set the data
Parameters
----------
pos : list, tuple or numpy array
Bounds of the region along the axis. len(pos) must be >=2.
color : list, tuple, or array
The color to use when drawing the line. It must have a shape of
(1, 4) for a single color region or (len(pos), 4) for a multicolor
region.
"""
new_pos = self._pos
new_color = self._color
if pos is not None:
num_elements = len(pos)
pos = np.array(pos, dtype=np.float32)
if pos.ndim != 1:
raise ValueError('Expected 1D array')
vertex = np.empty((num_elements * 2, 2), dtype=np.float32)
if self._is_vertical:
vertex[:, 0] = np.repeat(pos, 2)
vertex[:, 1] = np.tile([-1, 1], num_elements)
else:
vertex[:, 1] = np.repeat(pos, 2)
vertex[:, 0] = np.tile([1, -1], num_elements)
new_pos = vertex
self._changed['pos'] = True
if color is not None:
color = np.array(color, dtype=np.float32)
num_elements = new_pos.shape[0] / 2
if color.ndim == 2:
if color.shape[0] != num_elements:
raise ValueError('Expected a color for each pos')
if color.shape[1] != 4:
raise ValueError('Each color must be a RGBA array')
color = np.repeat(color, 2, axis=0).astype(np.float32)
elif color.ndim == 1:
if color.shape[0] != 4:
raise ValueError('Each color must be a RGBA array')
color = np.repeat([color], new_pos.shape[0], axis=0)
color = color.astype(np.float32)
else:
raise ValueError('Expected a numpy array of shape '
'(%d, 4) or (1, 4)' % num_elements)
new_color = color
self._changed['color'] = True
# Ensure pos and color have the same size
if new_pos.shape[0] != new_color.shape[0]:
raise ValueError('pos and color does must have the same size')
self._color = new_color
self._pos = new_pos
@property
def color(self):
return self._color[::2]
@property
def pos(self):
if self._is_vertical:
return self._pos[:, 0].ravel()[::2]
else:
return self._pos[:, 1].ravel()[::2]
def _compute_bounds(self, axis, view):
"""Return the (min, max) bounding values of this visual along *axis*
in the local coordinate system.
"""
is_vertical = self._is_vertical
pos = self._pos
if axis == 0 and is_vertical:
return (pos[0, 0], pos[-1, 0])
elif axis == 1 and not is_vertical:
return (pos[0, 1], pos[-1, 1])
return None
@property
def is_vertical(self):
return self._is_vertical
def _prepare_transforms(self, view=None):
program = view.view_program
transforms = view.transforms
program.vert['render_to_visual'] = transforms.get_transform('render',
'visual')
program.vert['transform'] = transforms.get_transform('visual',
'render')
def _prepare_draw(self, view=None):
"""This method is called immediately before each draw.
The *view* argument indicates which view is about to be drawn.
"""
if self._changed['pos']:
self.pos_buf.set_data(self._pos)
self._changed['pos'] = False
if self._changed['color']:
self.color_buf.set_data(self._color)
self._program.vert['color'] = self.color_buf
self._changed['color'] = False
return True