"""Plotting functions."""
import warnings
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
from ._layout import make_3d_axis
from ._artists import Arrow3D
from ..transformations import transform
from ..rotations import unitx, unitz, perpendicular_to_vectors, norm_vector
[docs]def plot_box(ax=None, size=np.ones(3), A2B=np.eye(4), ax_s=1, wireframe=True,
color="k", alpha=1.0):
"""Plot box.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
size : array-like, shape (3,), optional (default: [1, 1, 1])
Size of the box per dimension
A2B : array-like, shape (4, 4)
Center of the box
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of box and surface otherwise
color : str, optional (default: black)
Color in which the box should be plotted
alpha : float, optional (default: 1)
Alpha value of the mesh that will be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
corners = np.array([
[0, 0, 0],
[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1]
])
corners = (corners - 0.5) * size
corners = transform(
A2B, np.hstack((corners, np.ones((len(corners), 1)))))[:, :3]
if wireframe:
for i, j in [(0, 1), (0, 2), (1, 3), (2, 3),
(4, 5), (4, 6), (5, 7), (6, 7),
(0, 4), (1, 5), (2, 6), (3, 7)]:
ax.plot([corners[i, 0], corners[j, 0]],
[corners[i, 1], corners[j, 1]],
[corners[i, 2], corners[j, 2]],
c=color, alpha=alpha)
else:
p3c = Poly3DCollection(np.array([
[corners[0], corners[1], corners[2]],
[corners[1], corners[2], corners[3]],
[corners[4], corners[5], corners[6]],
[corners[5], corners[6], corners[7]],
[corners[0], corners[1], corners[4]],
[corners[1], corners[4], corners[5]],
[corners[2], corners[6], corners[7]],
[corners[2], corners[3], corners[7]],
[corners[0], corners[4], corners[6]],
[corners[0], corners[2], corners[6]],
[corners[1], corners[5], corners[7]],
[corners[1], corners[3], corners[7]],
]))
p3c.set_alpha(alpha)
p3c.set_facecolor(color)
ax.add_collection3d(p3c)
return ax
[docs]def plot_sphere(ax=None, radius=1.0, p=np.zeros(3), ax_s=1, wireframe=True,
n_steps=20, alpha=1.0, color="k"):
"""Plot sphere.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
radius : float, optional (default: 1)
Radius of the sphere
p : array-like, shape (3,), optional (default: [0, 0, 0])
Center of the sphere
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of sphere and surface otherwise
n_steps : int, optional (default: 20)
Number of discrete steps plotted in each dimension
alpha : float, optional (default: 1)
Alpha value of the sphere that will be plotted
color : str, optional (default: black)
Color in which the sphere should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
x = p[0] + radius * np.sin(phi) * np.cos(theta)
y = p[1] + radius * np.sin(phi) * np.sin(theta)
z = p[2] + radius * np.cos(phi)
if wireframe:
ax.plot_wireframe(
x, y, z, rstride=10, cstride=10, color=color, alpha=alpha)
else:
ax.plot_surface(x, y, z, color=color, alpha=alpha, linewidth=0)
return ax
[docs]def plot_cylinder(ax=None, length=1.0, radius=1.0, thickness=0.0,
A2B=np.eye(4), ax_s=1, wireframe=True, n_steps=100,
alpha=1.0, color="k"):
"""Plot cylinder.
A cylinder is the volume covered by a disk moving along a line segment.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
length : float, optional (default: 1)
Length of the cylinder
radius : float, optional (default: 1)
Radius of the cylinder
thickness : float, optional (default: 0)
Thickness of a cylindrical shell. It will be subtracted from the
outer radius to obtain the inner radius. The difference must be
greater than 0.
A2B : array-like, shape (4, 4)
Center of the cylinder
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of cylinder and surface otherwise
n_steps : int, optional (default: 100)
Number of discrete steps plotted in each dimension
alpha : float, optional (default: 1)
Alpha value of the cylinder that will be plotted
color : str, optional (default: black)
Color in which the cylinder should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
Raises
------
ValueError
If thickness is <= 0
"""
if ax is None:
ax = make_3d_axis(ax_s)
inner_radius = radius - thickness
if inner_radius <= 0.0:
raise ValueError("Thickness of cylindrical shell results in "
"invalid inner radius: %g" % inner_radius)
axis_start = A2B.dot(np.array([0, 0, -0.5 * length, 1]))[:3]
axis_end = A2B.dot(np.array([0, 0, 0.5 * length, 1]))[:3]
axis = axis_end - axis_start
axis /= length
not_axis = np.array([1, 0, 0])
if (axis == not_axis).all():
not_axis = np.array([0, 1, 0])
n1 = np.cross(axis, not_axis)
n1 /= np.linalg.norm(n1)
n2 = np.cross(axis, n1)
if wireframe:
t = np.linspace(0, length, n_steps)
else:
t = np.array([0, length])
theta = np.linspace(0, 2 * np.pi, n_steps)
t, theta = np.meshgrid(t, theta)
if thickness > 0.0:
X_outer, Y_outer, Z_outer = [
axis_start[i] + axis[i] * t
+ radius * np.sin(theta) * n1[i]
+ radius * np.cos(theta) * n2[i] for i in [0, 1, 2]]
X_inner, Y_inner, Z_inner = [
axis_end[i] - axis[i] * t
+ inner_radius * np.sin(theta) * n1[i]
+ inner_radius * np.cos(theta) * n2[i] for i in [0, 1, 2]]
X = np.hstack((X_outer, X_inner))
Y = np.hstack((Y_outer, Y_inner))
Z = np.hstack((Z_outer, Z_inner))
else:
X, Y, Z = [axis_start[i] + axis[i] * t
+ radius * np.sin(theta) * n1[i]
+ radius * np.cos(theta) * n2[i] for i in [0, 1, 2]]
if wireframe:
ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10, alpha=alpha,
color=color)
else:
ax.plot_surface(X, Y, Z, color=color, alpha=alpha, linewidth=0)
return ax
[docs]def plot_mesh(ax=None, filename=None, A2B=np.eye(4),
s=np.array([1.0, 1.0, 1.0]), ax_s=1, wireframe=False,
convex_hull=False, alpha=1.0, color="k"):
"""Plot mesh.
Note that this function requires the additional library 'trimesh'.
It will print a warning if trimesh is not available.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
filename : str, optional (default: None)
Path to mesh file.
A2B : array-like, shape (4, 4)
Pose of the mesh
s : array-like, shape (3,), optional (default: [1, 1, 1])
Scaling of the mesh that will be drawn
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of mesh and surface otherwise
convex_hull : bool, optional (default: False)
Show convex hull instead of the original mesh. This can be much
faster.
alpha : float, optional (default: 1)
Alpha value of the mesh that will be plotted
color : str, optional (default: black)
Color in which the mesh should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
if filename is None:
warnings.warn(
"No filename given for mesh. When you use the "
"UrdfTransformManager, make sure to set the mesh path or "
"package directory.")
return ax
try:
import trimesh
except ImportError:
warnings.warn(
"Cannot display mesh. Library 'trimesh' not installed.")
return ax
mesh = trimesh.load(filename)
if convex_hull:
mesh = mesh.convex_hull
vertices = mesh.vertices * s
vertices = np.hstack((vertices, np.ones((len(vertices), 1))))
vertices = transform(A2B, vertices)[:, :3]
vectors = np.array([vertices[[i, j, k]] for i, j, k in mesh.faces])
if wireframe:
surface = Line3DCollection(vectors)
surface.set_color(color)
else:
surface = Poly3DCollection(vectors)
surface.set_facecolor(color)
surface.set_alpha(alpha)
ax.add_collection3d(surface)
return ax
[docs]def plot_ellipsoid(ax=None, radii=np.ones(3), A2B=np.eye(4), ax_s=1,
wireframe=True, n_steps=20, alpha=1.0, color="k"):
"""Plot ellipsoid.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
radii : array-like, shape (3,)
Radii along the x-axis, y-axis, and z-axis of the ellipsoid.
A2B : array-like, shape (4, 4)
Transform from frame A to frame B
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of ellipsoid and surface otherwise
n_steps : int, optional (default: 20)
Number of discrete steps plotted in each dimension
alpha : float, optional (default: 1)
Alpha value of the ellipsoid that will be plotted
color : str, optional (default: black)
Color in which the ellipsoid should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
from ..transformations import transform, vectors_to_points
if ax is None:
ax = make_3d_axis(ax_s)
radius_x, radius_y, radius_z = radii
phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
x = radius_x * np.sin(phi) * np.cos(theta)
y = radius_y * np.sin(phi) * np.sin(theta)
z = radius_z * np.cos(phi)
shape = x.shape
P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
P = transform(A2B, vectors_to_points(P))[:, :3]
x = P[:, 0].reshape(*shape)
y = P[:, 1].reshape(*shape)
z = P[:, 2].reshape(*shape)
if wireframe:
ax.plot_wireframe(
x, y, z, rstride=10, cstride=10, color=color, alpha=alpha)
else:
ax.plot_surface(x, y, z, color=color, alpha=alpha, linewidth=0)
return ax
[docs]def plot_capsule(ax=None, A2B=np.eye(4), height=1.0, radius=1.0,
ax_s=1, wireframe=True, n_steps=20, alpha=1.0, color="k"):
"""Plot capsule.
A capsule is the volume covered by a sphere moving along a line segment.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
A2B : array-like, shape (4, 4)
Frame of the capsule, located at the center of the line segment.
height : float, optional (default: 1)
Height of the capsule along its z-axis.
radius : float, optional (default: 1)
Radius of the capsule.
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of capsule and surface otherwise
n_steps : int, optional (default: 20)
Number of discrete steps plotted in each dimension
alpha : float, optional (default: 1)
Alpha value of the mesh that will be plotted
color : str, optional (default: black)
Color in which the capsule should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
from ..transformations import transform, vectors_to_points
if ax is None:
ax = make_3d_axis(ax_s)
phi, theta = np.mgrid[0.0:np.pi:n_steps * 1j, 0.0:2.0 * np.pi:n_steps * 1j]
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
z[len(z) // 2:] -= 0.5 * height
z[:len(z) // 2] += 0.5 * height
shape = x.shape
P = np.column_stack((x.reshape(-1), y.reshape(-1), z.reshape(-1)))
P = transform(A2B, vectors_to_points(P))[:, :3]
x = P[:, 0].reshape(*shape)
y = P[:, 1].reshape(*shape)
z = P[:, 2].reshape(*shape)
if wireframe:
ax.plot_wireframe(
x, y, z, rstride=10, cstride=10, color=color, alpha=alpha)
else:
ax.plot_surface(x, y, z, color=color, alpha=alpha, linewidth=0)
return ax
[docs]def plot_cone(ax=None, height=1.0, radius=1.0, A2B=np.eye(4), ax_s=1,
wireframe=True, n_steps=20, alpha=1.0, color="k"):
"""Plot cone.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
height : float, optional (default: 1)
Height of the cone
radius : float, optional (default: 1)
Radius of the cone
A2B : array-like, shape (4, 4)
Center of the cone
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
wireframe : bool, optional (default: True)
Plot wireframe of cone and surface otherwise
n_steps : int, optional (default: 100)
Number of discrete steps plotted in each dimension
alpha : float, optional (default: 1)
Alpha value of the cone that will be plotted
color : str, optional (default: black)
Color in which the cone should be plotted
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
axis_start = A2B.dot(np.array([0, 0, 0, 1]))[:3]
axis_end = A2B.dot(np.array([0, 0, height, 1]))[:3]
axis = axis_end - axis_start
axis /= height
not_axis = np.array([1, 0, 0])
if (axis == not_axis).all():
not_axis = np.array([0, 1, 0])
n1 = np.cross(axis, not_axis)
n1 /= np.linalg.norm(n1)
n2 = np.cross(axis, n1)
if wireframe:
t = np.linspace(0, height, n_steps)
radii = np.linspace(radius, 0, n_steps)
else:
t = np.array([0, height])
radii = np.array([radius, 0])
theta = np.linspace(0, 2 * np.pi, n_steps)
t, theta = np.meshgrid(t, theta)
X, Y, Z = [axis_start[i] + axis[i] * t
+ radii * np.sin(theta) * n1[i]
+ radii * np.cos(theta) * n2[i] for i in [0, 1, 2]]
if wireframe:
ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10, alpha=alpha,
color=color)
else:
ax.plot_surface(X, Y, Z, color=color, alpha=alpha, linewidth=0)
return ax
[docs]def plot_vector(ax=None, start=np.zeros(3), direction=np.array([1, 0, 0]),
s=1.0, arrowstyle="simple", ax_s=1, **kwargs):
"""Plot Vector.
Draws an arrow from start to start + s * direction.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
start : array-like, shape (3,), optional (default: [0, 0, 0])
Start of the vector
direction : array-like, shape (3,), optional (default: [1, 0, 0])
Direction of the vector
s : float, optional (default: 1)
Scaling of the vector that will be drawn
arrowstyle : str, or ArrowStyle, optional (default: 'simple')
See matplotlib's documentation of arrowstyle in
matplotlib.patches.FancyArrowPatch for more options
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
kwargs : dict, optional (default: {})
Additional arguments for the plotting functions, e.g. alpha
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
axis_arrow = Arrow3D(
[start[0], start[0] + s * direction[0]],
[start[1], start[1] + s * direction[1]],
[start[2], start[2] + s * direction[2]],
mutation_scale=20, arrowstyle=arrowstyle, **kwargs)
ax.add_artist(axis_arrow)
return ax
[docs]def plot_length_variable(ax=None, start=np.zeros(3), end=np.ones(3), name="l",
above=False, ax_s=1, color="k", **kwargs):
"""Plot length with text at its center.
Parameters
----------
ax : Matplotlib 3d axis, optional (default: None)
If the axis is None, a new 3d axis will be created
start : array-like, shape (3,), optional (default: [0, 0, 0])
Start point
end : array-like, shape (3,), optional (default: [1, 1, 1])
End point
name : str, optional (default: 'l')
Text in the middle
above : bool, optional (default: False)
Plot name above line
ax_s : float, optional (default: 1)
Scaling of the new matplotlib 3d axis
color : str, optional (default: black)
Color in which the cylinder should be plotted
kwargs : dict, optional (default: {})
Additional arguments for the text, e.g. fontsize
Returns
-------
ax : Matplotlib 3d axis
New or old axis
"""
if ax is None:
ax = make_3d_axis(ax_s)
direction = end - start
length = np.linalg.norm(direction)
if above:
ax.plot([start[0], end[0]], [start[1], end[1]],
[start[2], end[2]], color=color)
else:
mid1 = start + 0.4 * direction
mid2 = start + 0.6 * direction
ax.plot([start[0], mid1[0]], [start[1], mid1[1]],
[start[2], mid1[2]], color=color)
ax.plot([end[0], mid2[0]], [end[1], mid2[1]],
[end[2], mid2[2]], color=color)
if np.linalg.norm(direction / length - unitz) < np.finfo(float).eps:
axis = unitx
else:
axis = unitz
mark = (norm_vector(perpendicular_to_vectors(direction, axis))
* 0.03 * length)
mark_start1 = start + mark
mark_start2 = start - mark
mark_end1 = end + mark
mark_end2 = end - mark
ax.plot([mark_start1[0], mark_start2[0]],
[mark_start1[1], mark_start2[1]],
[mark_start1[2], mark_start2[2]],
color=color)
ax.plot([mark_end1[0], mark_end2[0]],
[mark_end1[1], mark_end2[1]],
[mark_end1[2], mark_end2[2]],
color=color)
text_location = start + 0.45 * direction
if above:
text_location[2] += 0.3 * length
ax.text(text_location[0], text_location[1], text_location[2], name,
zdir="x", **kwargs)
return ax