Source code for pytransform3d.plot_utils._artists

"""Matplotlib artists."""
import numpy as np
from matplotlib import artist
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
from mpl_toolkits.mplot3d.art3d import Line3D, Text3D


[docs]class Frame(artist.Artist): """A Matplotlib artist that displays a frame represented by its basis. Parameters ---------- A2B : array-like, shape (4, 4) Transform from frame A to frame B label : str, optional (default: None) Name of the frame s : float, optional (default: 1) Length of basis vectors Other arguments except 'c' and 'color' are passed on to Line3D. """
[docs] def __init__(self, A2B, label=None, s=1.0, **kwargs): super(Frame, self).__init__() if "c" in kwargs: kwargs.pop("c") if "color" in kwargs: kwargs.pop("color") self.s = s self.x_axis = Line3D([], [], [], color="r", **kwargs) self.y_axis = Line3D([], [], [], color="g", **kwargs) self.z_axis = Line3D([], [], [], color="b", **kwargs) self.draw_label = label is not None self.label = label if self.draw_label: self.label_indicator = Line3D([], [], [], color="k", **kwargs) self.label_text = Text3D(0, 0, 0, text="", zdir="x") self.set_data(A2B, label)
[docs] def set_data(self, A2B, label=None): """Set the transformation data. Parameters ---------- A2B : array-like, shape (4, 4) Transform from frame A to frame B label : str, optional (default: None) Name of the frame """ R = A2B[:3, :3] p = A2B[:3, 3] for d, b in enumerate([self.x_axis, self.y_axis, self.z_axis]): b.set_data(np.array([p[0], p[0] + self.s * R[0, d]]), np.array([p[1], p[1] + self.s * R[1, d]])) b.set_3d_properties(np.array([p[2], p[2] + self.s * R[2, d]])) if self.draw_label: if label is None: label = self.label label_pos = p + 0.5 * self.s * (R[:, 0] + R[:, 1] + R[:, 2]) self.label_indicator.set_data( np.array([p[0], label_pos[0]]), np.array([p[1], label_pos[1]])) self.label_indicator.set_3d_properties( np.array([p[2], label_pos[2]])) self.label_text.set_text(label) self.label_text.set_position([label_pos[0], label_pos[1]]) self.label_text.set_3d_properties(label_pos[2], zdir="x")
[docs] @artist.allow_rasterization def draw(self, renderer, *args, **kwargs): """Draw the artist.""" for b in [self.x_axis, self.y_axis, self.z_axis]: b.draw(renderer, *args, **kwargs) if self.draw_label: self.label_indicator.draw(renderer, *args, **kwargs) self.label_text.draw(renderer, *args, **kwargs) super(Frame, self).draw(renderer, *args, **kwargs)
[docs] def add_frame(self, axis): """Add the frame to a 3D axis.""" for b in [self.x_axis, self.y_axis, self.z_axis]: axis.add_line(b) if self.draw_label: axis.add_line(self.label_indicator) axis._add_text(self.label_text)
[docs]class LabeledFrame(Frame): """Displays a frame represented by its basis with axis labels. Parameters ---------- A2B : array-like, shape (4, 4) Transform from frame A to frame B label : str, optional (default: None) Name of the frame s : float, optional (default: 1) Length of basis vectors Other arguments except 'c' and 'color' are passed on to Line3D. """
[docs] def __init__(self, A2B, label=None, s=1.0, **kwargs): self.x_label = Text3D(0, 0, 0, text="", zdir="x") self.y_label = Text3D(0, 0, 0, text="", zdir="x") self.z_label = Text3D(0, 0, 0, text="", zdir="x") super(LabeledFrame, self).__init__( A2B, label=label, s=s, **kwargs)
[docs] def set_data(self, A2B, label=None): """Set the transformation data. Parameters ---------- A2B : array-like, shape (4, 4) Transform from frame A to frame B label : str, optional (default: None) Name of the frame """ super(LabeledFrame, self).set_data(A2B, label) R = A2B[:3, :3] p = A2B[:3, 3] x_label_location = p + 1.1 * self.s * R[:, 0] y_label_location = p + 1.1 * self.s * R[:, 1] z_label_location = p + 1.1 * self.s * R[:, 2] self.x_label.set_text("x") self.x_label.set_position(x_label_location[:2]) self.x_label.set_3d_properties(x_label_location[2], zdir="x") self.y_label.set_text("y") self.y_label.set_position(y_label_location[:2]) self.y_label.set_3d_properties(y_label_location[2], zdir="x") self.z_label.set_text("z") self.z_label.set_position(z_label_location[:2]) self.z_label.set_3d_properties(z_label_location[2], zdir="x")
[docs] @artist.allow_rasterization def draw(self, renderer, *args, **kwargs): """Draw the artist.""" self.x_label.draw(renderer, *args, **kwargs) self.y_label.draw(renderer, *args, **kwargs) self.z_label.draw(renderer, *args, **kwargs) super(LabeledFrame, self).draw(renderer, *args, **kwargs)
[docs] def add_frame(self, axis): """Add the frame to a 3D axis.""" super(LabeledFrame, self).add_frame(axis) axis._add_text(self.x_label) axis._add_text(self.y_label) axis._add_text(self.z_label)
[docs]class Trajectory(artist.Artist): """A Matplotlib artist that displays a trajectory. Parameters ---------- H : array-like, shape (n_steps, 4, 4) Sequence of poses represented by homogeneous matrices show_direction : bool, optional (default: True) Plot an arrow to indicate the direction of the trajectory n_frames : int, optional (default: 10) Number of frames that should be plotted to indicate the rotation s : float, optional (default: 1) Scaling of the frames that will be drawn Other arguments are passed onto Line3D. """
[docs] def __init__(self, H, show_direction=True, n_frames=10, s=1.0, **kwargs): super(Trajectory, self).__init__() self.show_direction = show_direction self.trajectory = Line3D([], [], [], **kwargs) self.key_frames = [Frame(np.eye(4), s=s, **kwargs) for _ in range(n_frames)] if self.show_direction: self.direction_arrow = Arrow3D( [0, 0], [0, 0], [0, 0], mutation_scale=20, lw=1, arrowstyle="-|>", color="k") self.set_data(H)
[docs] def set_data(self, H): """Set the trajectory data. Parameters ---------- H : array-like, shape (n_steps, 4, 4) Sequence of poses represented by homogeneous matrices """ positions = H[:, :3, 3] self.trajectory.set_data(positions[:, 0], positions[:, 1]) self.trajectory.set_3d_properties(positions[:, 2]) key_frames_indices = np.linspace( 0, len(H) - 1, len(self.key_frames), dtype=np.int) for i, key_frame_idx in enumerate(key_frames_indices): self.key_frames[i].set_data(H[key_frame_idx]) if self.show_direction: start = 0.8 * positions[0] + 0.2 * positions[-1] end = 0.2 * positions[0] + 0.8 * positions[-1] self.direction_arrow.set_data( [start[0], end[0]], [start[1], end[1]], [start[2], end[2]])
[docs] @artist.allow_rasterization def draw(self, renderer, *args, **kwargs): """Draw the artist.""" self.trajectory.draw(renderer, *args, **kwargs) for key_frame in self.key_frames: key_frame.draw(renderer, *args, **kwargs) if self.show_direction: self.direction_arrow.draw(renderer) super(Trajectory, self).draw(renderer, *args, **kwargs)
[docs] def add_trajectory(self, axis): """Add the trajectory to a 3D axis.""" axis.add_line(self.trajectory) for key_frame in self.key_frames: key_frame.add_frame(axis) if self.show_direction: axis.add_artist(self.direction_arrow)
[docs]class Arrow3D(FancyArrowPatch): """A Matplotlib patch that represents an arrow in 3D. Source: http://stackoverflow.com/a/11156353/915743 """
[docs] def __init__(self, xs, ys, zs, *args, **kwargs): super(Arrow3D, self).__init__((0, 0), (0, 0), *args, **kwargs) self._verts3d = xs, ys, zs
[docs] def set_data(self, xs, ys, zs): """Set the arrow data. Parameters ---------- xs : iterable List of x positions ys : iterable List of y positions zs : iterable List of z positions """ self._verts3d = xs, ys, zs
[docs] def draw(self, renderer): """Draw the patch.""" xs3d, ys3d, zs3d = self._verts3d try: M = self.axes.M except AttributeError: # deprecated since matplotlib 3.4, will be removed in 3.6 M = renderer.M xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, M) self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) super(Arrow3D, self).draw(renderer)