class MotionPlotter(scene.SceneCanvas):
def __init__(self,
keys='interactive',
size=(640, 480),
show=True,
**kwargs):
super().__init__(keys=keys, size=size, show=show, **kwargs)
self.unfreeze()
self._viewbox = self.central_widget.add_view(camera='turntable')
self._baseAxis = visuals.XYZAxis(parent=self._viewbox.scene, width=5)
self._gridLines = visuals.GridLines()
self._viewbox.add(self._gridLines)
self._cubeAxis = visuals.XYZAxis(parent=self._viewbox.scene, width=5)
Plot3D = scene.visuals.create_visual_node(LinePlotVisual)
self._plot = Plot3D(([0], [0], [0]),
width=3.0,
color='y',
edge_color='w',
symbol='x',
face_color=(0.2, 0.2, 1, 0.8),
parent=self._viewbox.scene)
self._xPos = np.array([0], dtype=np.float32)
self._yPos = np.array([0], dtype=np.float32)
self._zPos = np.array([0], dtype=np.float32)
self._cube = visuals.Cube(parent=self._viewbox.scene,
color=(0.5, 0.5, 1, 0.5),
edge_color=(0.6, 0.2, 0.8, 1))
self._transform = MatrixTransform()
self._cube.transform = self._transform
self._cubeAxis.transform = self._transform
self.freeze()
def transformCube(self, t, q):
quaternion = Quaternion(*q)
self._transform.reset()
self._transform.matrix = quaternion.get_matrix()
self._transform.translate(t)
x, y, z = t
self._xPos = np.append(self._xPos, x)
self._yPos = np.append(self._yPos, y)
self._zPos = np.append(self._zPos, z)
self._plot.set_data((self._xPos.transpose(), self._yPos.transpose(),
self._zPos.transpose()),
symbol=None)
def resetTransform(self):
self._xPos = np.array([0], dtype=np.float32)
self._yPos = np.array([0], dtype=np.float32)
self._zPos = np.array([0], dtype=np.float32)
class Canvas(app.Canvas):
def __init__(self):
app.Canvas.__init__(self, keys='interactive', size=(800, 550))
vertices, faces, outline = create_box(width=1,
height=1,
depth=1,
width_segments=4,
height_segments=8,
depth_segments=16)
self.box = visuals.BoxVisual(width=1,
height=1,
depth=1,
width_segments=4,
height_segments=8,
depth_segments=16,
vertex_colors=vertices['color'],
edge_color='b')
self.theta = 0
self.phi = 0
self.transform = MatrixTransform()
self.box.transform = self.transform
self.show()
self.timer = app.Timer(connect=self.rotate)
self.timer.start(0.016)
def rotate(self, event):
self.theta += .5
self.phi += .5
self.transform.reset()
self.transform.rotate(self.theta, (0, 0, 1))
self.transform.rotate(self.phi, (0, 1, 0))
self.transform.scale((100, 100, 0.001))
self.transform.translate((200, 200))
self.update()
def on_resize(self, event):
# Set canvas viewport and reconfigure visual transforms to match.
vp = (0, 0, self.physical_size[0], self.physical_size[1])
self.context.set_viewport(*vp)
self.box.transforms.configure(canvas=self, viewport=vp)
def on_draw(self, ev):
gloo.clear(color='white', depth=True)
self.box.draw()
class Canvas(app.Canvas):
def __init__(self):
app.Canvas.__init__(self, keys='interactive', size=(800, 550))
vertices, faces, outline = create_box(width=1, height=1, depth=1,
width_segments=4,
height_segments=8,
depth_segments=16)
self.box = visuals.BoxVisual(width=1, height=1, depth=1,
width_segments=4,
height_segments=8,
depth_segments=16,
vertex_colors=vertices['color'],
edge_color='b')
self.theta = 0
self.phi = 0
self.transform = MatrixTransform()
self.box.transform = self.transform
self.show()
self.timer = app.Timer(connect=self.rotate)
self.timer.start(0.016)
def rotate(self, event):
self.theta += .5
self.phi += .5
self.transform.reset()
self.transform.rotate(self.theta, (0, 0, 1))
self.transform.rotate(self.phi, (0, 1, 0))
self.transform.scale((100, 100, 0.001))
self.transform.translate((200, 200))
self.update()
def on_resize(self, event):
# Set canvas viewport and reconfigure visual transforms to match.
vp = (0, 0, self.physical_size[0], self.physical_size[1])
self.context.set_viewport(*vp)
self.box.transforms.configure(canvas=self, viewport=vp)
def on_draw(self, ev):
gloo.clear(color='white', depth=True)
self.box.draw()
class Colorbar(scene.visuals.Image):
""" A colorbar visual fixed to the right side of the canvas. This is
based on the rendering from Matplotlib, then display this rendered
image as a scene.visuals.Image visual node on the canvas.
Parameters:
"""
def __init__(self,
size=(500, 10),
cmap='grays',
clim=None,
label_str="Colorbar",
label_color='black',
label_size=12,
tick_size=10,
border_width=1.0,
border_color='black',
visible=True,
parent=None):
assert clim is not None, 'clim must be specified explicitly.'
# Create a scene.visuals.Image (without parent by default).
scene.visuals.Image.__init__(self,
parent=None,
interpolation='nearest',
method='auto')
self.unfreeze()
self.visible = visible
self.canvas_size = None # will be set when parent is linked
# Record the important drawing parameters.
self.pos = (0, 0)
self.bar_size = size # tuple
self.cmap = get_colormap(cmap) # vispy Colormap
self.clim = clim # tuple
# Record the styling parameters.
self.label_str = label_str
self.label_color = label_color
self.label_size = label_size
self.tick_size = tick_size
self.border_width = border_width
self.border_color = border_color
# Draw colorbar using Matplotlib.
self.set_data(self._draw_colorbar())
# Give a Matrix transform to self in order to move around canvas.
self.transform = MatrixTransform()
self.freeze()
def on_resize(self, event):
""" When window is resized, only need to move the position in vertical
direction, because the coordinate is relative to the secondary ViewBox
that stays on the right side of the canvas.
"""
pos = np.array(self.pos).astype(np.single)
pos[1] *= event.size[1] / self.canvas_size[1]
self.pos = tuple(pos)
# Move the colorbar to specified position (with half-size padding, because
# Image visual uses a different anchor (top-left corner) rather than the
# center-left corner used by ColorBar visual.).
self.transform.reset()
self.transform.translate((
self.pos[0] / 2.618, # make the gap smaller :)
self.pos[1] - self.size[1] / 2.))
# Update the canvas size.
self.canvas_size = event.size
def _draw_colorbar(self):
""" Draw a Matplotlib colorbar, save this figure without any boundary to a
rendering buffer, and return this buffer as a numpy array.
"""
dpi = get_dpi()
# Compute the Matplotlib figsize: note the order of width, height.
# The width is doubled because we will only keep the colorbar portion
# before we export this image to buffer!
figsize = (self.bar_size[1] / dpi * 2, self.bar_size[0] / dpi)
# Convert cmap and clim to Matplotlib format.
rgba = self.cmap.colors.rgba
# Blend to white to avoid this Matplotlib rendering issue:
# https://github.com/matplotlib/matplotlib/issues/1188
for i in range(3):
rgba[:, i] = (1 - rgba[:, -1]) + rgba[:, -1] * rgba[:, i]
rgba[:, -1] = 1.
if len(rgba) < 2: # in special case of 'grays' cmap!
rgba = np.array([[0, 0, 0, 1.], [1, 1, 1, 1.]])
cmap = LinearSegmentedColormap.from_list('vispy_cmap', rgba)
norm = mpl.colors.Normalize(vmin=self.clim[0], vmax=self.clim[1])
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
# Put the colorbar at proper location on the Matplotlib fig.
fig = plt.figure(figsize=figsize, dpi=dpi)
ax = plt.Axes(fig, [0, 0, 1, 1])
ax.set_axis_off()
fig.add_axes(ax)
divider = make_axes_locatable(ax)
cax = divider.append_axes('right', size='100%', pad=0.)
cb = fig.colorbar(sm, cax=cax)
ax.remove()
# Apply styling to the colorbar.
cb.set_label(self.label_str,
color=self.label_color,
fontsize=self.label_size)
plt.setp(plt.getp(cb.ax.axes, 'yticklabels'),
color=self.label_color,
fontsize=self.tick_size)
cb.ax.yaxis.set_tick_params(color=self.label_color)
cb.outline.set_linewidth(self.border_width)
cb.outline.set_edgecolor(self.border_color)
# Export the rendering to a numpy array in the buffer.
buf = io.BytesIO()
fig.savefig(buf,
format='png',
bbox_inches='tight',
pad_inches=0,
dpi=dpi,
transparent=True)
buf.seek(0)
return plt.imread(buf)
class mbVector(scene.visuals.Vector):
"""
a simple wrapper class which includes the transformation done on the vector
"""
def __init__(self, view, face_color, state_vec, orient_vec):
super(mbVector, self).__init__(10,
10,
0.05,
1.,
0.1,
0.25,
color=face_color,
shading="smooth",
parent=view)
self.unfreeze()
self.n = 0
self.n_max = len(state_vec)
self.trafo = MatrixTransform()
self.state_vec = state_vec
self.orient_vec = orient_vec
def _update(self):
"""
update
called for every time-step
"""
self.n += 1
if self.n > self.n_max - 1:
self.n = 0
return
# first flip the vector in the x direction, such that further
# rotations are handled correctly
self.trafo.set_rotation((0, 0, 1, 0, 1, 0, 1, 0, 0))
# calculates the scaling of the object on bases of the orientation
# vector
scale_x = self.orient_vec[self.n][0]
scale_y = self.orient_vec[self.n][1]
scale_z = self.orient_vec[self.n][2]
scale = np.sqrt(scale_x * scale_x + scale_y * scale_y +
scale_z * scale_z)
scale = (scale, 1., 1.)
# calculates the translation of the object on bases of the state vector
dx = self.state_vec[self.n][0]
dy = self.state_vec[self.n][1]
dz = self.state_vec[self.n][2]
# calculates the rotation (orthogonal O3 Trafo) on bases of the
# orientation vector)
ex, ey, ez = self._get_ortho_base((scale_x, scale_y, scale_z))
new_base = (ex[0], ex[1], ex[2], ey[0], ey[1], ey[2], ez[0], ez[1],
ez[2])
self._doTrafo(dx, dy, dz, base=new_base, scale=scale, reset=False)
def _doTrafo(self, x=0., y=0., z=0., base=None, scale=None, reset=True):
"""
doing first the scale then rotation and afterwards translate
"""
if reset:
self.trafo.reset()
if scale is not None:
self.trafo.scale(scale)
if base is not None:
self.trafo.mult_rotation(base)
self.trafo.translate((x, y, z))
self.transform = self.trafo
def _norm(self, n):
""" normalizing a vector n = (x,y,z) """
norm = np.linalg.norm(n)
if norm > 0:
return n / norm
def _cross(self, n0, n1):
""" doing a crossproduct i.e. n1 x n2 """
return np.cross(n0, n1)
def _ortho(self, n):
""" finding an arbitrary orthogonal vector to another in 3d """
x, y, z = n
if z != 0. and y != 0.:
return np.array((0., z, -y))
elif x != 0. and y != 0.:
return np.array((y, -x, 0.))
elif x != 0. and z != 0.:
return np.array((z, 0., -x))
elif x == 0 and y == 0:
return np.array((1., 0., 0.))
elif x == 0 and z == 0:
return np.array((1., 0., 0.))
elif y == 0 and z == 0:
return np.array((0., 1., 0.))
else:
return np.array((0., 0., 0.))
def _get_ortho_base(self, n):
"""
calc an ortho base for one direction, such that ex is pointing in the
end to that direction
"""
ex = self._norm(n)
ey = self._norm(self._ortho(ex))
ez = self._cross(ex, ey)
return ex, ey, ez
class AxisAlignedImage(scene.visuals.Image):
""" Visual subclass displaying an image that aligns to an axis.
This image should be able to move along the perpendicular direction when
user gives corresponding inputs.
Parameters:
"""
def __init__(self, image_funcs, axis='z', pos=0, limit=None,
seismic_coord_system=True,
cmaps=['grays'], clims=None,
interpolation='nearest', method='auto'):
assert clims is not None, 'clim must be specified explicitly.'
# Create an Image obj and unfreeze it so we can add more
# attributes inside.
# First image (from image_funcs[0])
scene.visuals.Image.__init__(self, parent=None, # no image func yet
cmap=cmaps[0], clim=clims[0],
interpolation=interpolation, method=method)
self.unfreeze()
self.interactive = True
# Other images ...
self.overlaid_images = [self]
for i_img in range(1, len(image_funcs)):
overlaid_image = scene.visuals.Image(parent=self,
cmap=cmaps[i_img], clim=clims[i_img],
interpolation=interpolation, method=method)
self.overlaid_images.append(overlaid_image)
# Set GL state. Must check depth test, otherwise weird in 3D.
self.set_gl_state(depth_test=True, depth_func='lequal',
blend_func=('src_alpha', 'one_minus_src_alpha'))
# Determine the axis and position of this plane.
self.axis = axis
# Check if pos is within the range.
if limit is not None:
assert (pos>=limit[0]) and (pos<=limit[1]), \
'pos={} is outside limit={} range.'.format(pos, limit)
self.pos = pos
self.limit = limit
self.seismic_coord_system = seismic_coord_system
# Get the image_func that returns either image or image shape.
self.image_funcs = image_funcs # a list of functions!
shape = self.image_funcs[0](self.pos, get_shape=True)
# The selection highlight (a Plane visual with transparent color).
# The plane is initialized before any rotation, on '+z' direction.
self.highlight = scene.visuals.Plane(parent=self,
width=shape[0], height=shape[1], direction='+z',
color=(1, 1, 0, 0.1)) # transparent yellow color
# Move the plane to align with the image.
self.highlight.transform = STTransform(
translate=(shape[0]/2, shape[1]/2, 0))
# This is to make sure we can see highlight plane through the images.
self.highlight.set_gl_state('additive', depth_test=True)
self.highlight.visible = False # only show when selected
# Set the anchor point (2D local world coordinates). The mouse will
# drag this image by anchor point moving in the normal direction.
self.anchor = None # None by default
self.offset = 0
# Apply SRT transform according to the axis attribute.
self.transform = MatrixTransform()
# Move the image plane to the corresponding location.
self._update_location()
self.freeze()
@property
def axis(self):
"""The dimension that this image is perpendicular aligned to."""
return self._axis
@axis.setter
def axis(self, value):
value = value.lower()
if value not in ('z', 'y', 'x'):
raise ValueError('Invalid value for axis.')
self._axis = value
def set_anchor(self, mouse_press_event):
""" Set an anchor point (2D coordinate on the image plane) when left click
in the selection mode (<Ctrl> pressed). After that, the dragging called
in func 'drag_visual_node' will try to move along the normal direction
and let the anchor follows user's mouse position as close as possible.
"""
# Get the screen-to-local transform to get camera coordinates.
tr = self.canvas.scene.node_transform(self)
# Get click (camera) coordinate in the local world.
click_pos = tr.map([*mouse_press_event.pos, 0, 1])
click_pos /= click_pos[3] # rescale to cancel out the pos.w factor
# Get the view direction (camera-to-target) vector in the local world.
view_vector = tr.map([*mouse_press_event.pos, 1, 1])[:3]
view_vector /= np.linalg.norm(view_vector) # normalize to unit vector
# Get distance from camera to the drag anchor point on the image plane.
# Eq 1: click_pos + distance * view_vector = anchor
# Eq 2: anchor[2] = 0 <- intersects with the plane
# The following equation can be derived by Eq 1 and Eq 2.
distance = (0. - click_pos[2]) / view_vector[2]
self.anchor = click_pos[:2] + distance * view_vector[:2] # only need vec2
def drag_visual_node(self, mouse_move_event):
""" Drag this visual node while holding left click in the selection mode
(<Ctrl> pressed). The plane will move in the normal direction
perpendicular to this image, and the anchor point (set with func
'set_anchor') will move along the normal direction to stay as close to
the mouse as possible, so that user feels like 'dragging' the plane.
"""
# Get the screen-to-local transform to get camera coordinates.
tr = self.canvas.scene.node_transform(self)
# Unlike in 'set_anchor', we now convert most coordinates to the screen
# coordinate system, because it's more intuitive for user to do operations
# in 2D and get 2D feedbacks, e.g. mouse leading the anchor point.
anchor = [*self.anchor, self.pos, 1] # 2D -> 3D
anchor_screen = tr.imap(anchor) # screen coordinates of the anchor point
anchor_screen /= anchor_screen[3] # rescale to cancel out 'w' term
anchor_screen = anchor_screen[:2] # only need vec2
# Compute the normal vector, starting from the anchor point and
# perpendicular to the image plane.
normal = [*self.anchor, self.pos+1, 1] # +[0,0,1,0] from anchor
normal_screen = tr.imap(normal) # screen coordinates of anchor + [0,0,1,0]
normal_screen /= normal_screen[3] # rescale to cancel out 'w' term
normal_screen = normal_screen[:2] # only need vec2
normal_screen -= anchor_screen # end - start = vector
normal_screen /= np.linalg.norm(normal_screen) # normalize to unit vector
# Use the vector {anchor_screen -> mouse.pos} and project to the
# normal_screen direction using dot product, we can get how far the plane
# should be moved (on the screen!).
drag_vector = mouse_move_event.pos[:2] - anchor_screen
drag = np.dot(drag_vector, normal_screen) # normal_screen must be length 1
# We now need to convert the move distance from screen coordinates to
# local world coordinates. First, find where the anchor is on the screen
# after dragging; then, convert that screen point to a local line shooting
# across the normal vector; finally, find where the line comes directly
# above/below the anchor point (before dragging) and get that distance as
# the true dragging distance in local coordinates.
new_anchor_screen = anchor_screen + normal_screen * drag
new_anchor = tr.map([*new_anchor_screen, 0, 1])
new_anchor /= new_anchor[3] # rescale to cancel out the pos.w factor
view_vector = tr.map([*new_anchor_screen, 1, 1])[:3]
view_vector /= np.linalg.norm(view_vector) # normalize to unit vector
# Solve this equation:
# new_anchor := new_anchor + view_vector * ?,
# ^^^ describe a 3D line of possible new anchor positions
# arg min (?) |new_anchor[:2] - anchor[:2]|
# ^^^ find a point on that 3D line that minimize the 2D distance between
# new_anchor and anchor.
numerator = anchor[:2] - new_anchor[:2]
numerator *= view_vector[:2] # element-wise multiplication
numerator = np.sum(numerator)
denominator = view_vector[0]**2 + view_vector[1]**2
shoot_distance = numerator / denominator
# Shoot from new_anchor to get the new intersect point. The z- coordinate
# of this point will be our dragging offset.
offset = new_anchor[2] + view_vector[2] * shoot_distance
# Note: must reverse normal direction from -y direction to +y!
if self.axis == 'y': offset = -offset
# Limit the dragging within range.
if self.limit is not None:
if self.pos + offset < self.limit[0]: offset = self.limit[0] - self.pos
if self.pos + offset > self.limit[1]: offset = self.limit[1] - self.pos
self.offset = offset
# Note: must reverse normal direction from +y direction to -y!
if self.axis == 'y': offset = -offset
self._update_location()
def _update_location(self):
""" Update the image plane to the dragged location and redraw this image.
"""
self.pos += self.offset
self.pos = int(np.round(self.pos)) # must round to nearest integer location
# Update the transformation in order to move to new location.
self.transform.reset()
if self.axis == 'z':
# 1. No rotation to do for z axis (y-x) slice. Only translate.
self.transform.translate((0, 0, self.pos))
elif self.axis == 'y':
# 2. Rotation(s) for the y axis (z-x) slice, then translate:
self.transform.rotate(90, (1, 0, 0))
self.transform.translate((0, self.pos, 0))
elif self.axis == 'x':
# 3. Rotation(s) for the x axis (z-y) slice, then translate:
self.transform.rotate(90, (1, 0, 0))
self.transform.rotate(90, (0, 0, 1))
self.transform.translate((self.pos, 0, 0))
# Update image on the slice based on current position. The numpy array
# is transposed due to a conversion from i-j to x-y axis system.
# First image, the primary one:
self.set_data(self.image_funcs[0](self.pos).T)
# Other images, overlaid on the primary image:
for i_img in range(1, len(self.image_funcs)):
self.overlaid_images[i_img].set_data(
self.image_funcs[i_img](self.pos).T)
# Reset attributes after dragging completes.
self.offset = 0
self._bounds_changed() # update the bounds with new self.pos
def _compute_bounds(self, axis_3d, view):
""" Overwrite the original 2D bounds of the Image class. This will correct
the automatic range setting for the camera in the scene canvas. In the
original Image class, the code assumes that the image always lies in x-y
plane; here we generalize that to x-z and y-z plane.
Parameters:
axis_3d: int in {0, 1, 2}, represents the axis in 3D view box.
view: the ViewBox object that connects to the parent.
The function returns a tuple (low_bounds, high_bounds) that represents
the spatial limits of self obj in the 3D scene.
"""
# Note: self.size[0] is slow dim size, self.size[1] is fast dim size.
if self.axis == 'z':
if axis_3d==0: return (0, self.size[0])
elif axis_3d==1: return (0, self.size[1])
elif axis_3d==2: return (self.pos, self.pos)
elif self.axis == 'y':
if axis_3d==0: return (0, self.size[0])
elif axis_3d==1: return (self.pos, self.pos)
elif axis_3d==2: return (0, self.size[1])
elif self.axis == 'x':
if axis_3d==0: return (self.pos, self.pos)
elif axis_3d==1: return (0, self.size[0])
elif axis_3d==2: return (0, self.size[1])
class Body():
"""
binds tranfos to the bodies and the order of doing it
"""
def __init__(self, state_vec, p):
self.state_vec = state_vec
self.n = 0
self.n_max = len(state_vec)
self.p = p
self.radius = p/6.0
self.rot = MatrixTransform()
self.v_orient = None
self.x = 0.
self.y = 0.
self.z = 0.
#print("-------", self.n_max, self.p)
def set_orient(self,v_orient):
"""
set orientation to v_orient.
:param v_orient: new orientation of point object
"""
self.v_orient = v_orient
def _update(self):
"""
update
rotation and then translation
"""
global OO
self.n += 1
if self.n > self.n_max-1:
self.n = 0
return
dx = self.state_vec[self.n][0]
dy = self.state_vec[self.n][1]
dz = self.state_vec[self.n][2]
if self.v_orient:
new_base = self.v_orient[self.n]
else:
new_base = None
self.trafo(dx,dy,dz,base=new_base)
def _norm(self, n):
""" normalizing a vector n = (x,y,z) """
x, y, z = n
norm = np.sqrt(x*x + y*y + z*z)
return x/norm, y/norm, z/norm
def _cross(self, n0, n1):
""" doing a crossproduct i.e. n1 x n2 """
x0, y0, z0 = n0
x1, y1, z1 = n1
x = y0*z1 - z0*y1
y = z0*x1 - x0*z1
z = x0*y1 - y0*x1
return x,y,z
def _ortho(self, n):
""" finding an arbitrary orthogonal vector to another in 3d """
x,y,z = n
if z!=0. and y!=0.:
return 0., z, -y
elif x!=0. and y!=0.:
return y, -x, 0.
elif x!=0. and z!=0.:
return z, 0., -x
else:
return x,y,z+1
def _get_ortho_base(self, n):
"""
calc an ottho base for one direction, such that ex is pointing in the end to that direction
"""
ex = self._norm(n)
ey = self._norm(self._ortho(ex))
ez = self._cross(ex, ey)
return ex, ey, ez
def trafo(self, x=0.,y=0.,z=0., base=None, scale=None, reset=True):
"""
doing first the scale then rotation and afterwards translate
"""
if reset:
self.rot.reset()
if scale is not None:
self.rot.scale(scale)
if base is not None:
self.rot.mult_rotation(base)
self.rot.translate((x,y,z))
self.transform = self.rot
