"""
pyVista plotting backend for MPT.
Implement the PlotterProtocol.
"""
# Plot a unit cube
from typing import Optional
import numpy as np
from scipy.linalg import block_diag
from mpt4py.base import Matrix, Vector
from mpt4py.tolerances import tolerances
from .plot import PolyhedralPlottingData, PlotterProtocol
# Note: pyvista takes a long time to load during debugging. Try to avoid importing this file if debugging
import pyvista as pv
# This ensures that the plotting is done in the notebook and not on a server, which is slower
pv.set_jupyter_backend('client')
[docs]
class PyvistaPlotter(PlotterProtocol):
def __init__(self, plotter: Optional[pv.Plotter] = None):
if isinstance(plotter, pv.Plotter):
self.plotter = plotter
else:
self.plotter = pv.Plotter(line_smoothing=True)
self.is_3d = False
[docs]
def plot_convexhull(self, points: Matrix, rays: Optional[Matrix] = None,
color: Optional[str] = 'lightblue',
opacity: Optional[float] = 1.,
show_edges: Optional[bool] = True,
edge_width: Optional[float] = 1.,
edge_color: Optional[str] = 'black',
show_vertices: Optional[bool] = False,
vertex_size: Optional[float] = 10.,
vertex_color: Optional[str] = 'black',
**kwargs):
if points.shape[1] == 3:
self.is_3d = True
data = PolyhedralPlottingData(points, rays)
polyhedron_connectivity = []
for facet in data.facets:
polyhedron_connectivity.extend([len(facet), *facet])
self.polyhedron_connectivity = [len(polyhedron_connectivity)+1, len(data.facets), *polyhedron_connectivity]
points = data.points
if points.shape[1] == 2:
points = np.hstack((data.points, np.zeros((data.points.shape[0], 1))))
mesh = pv.UnstructuredGrid([self.polyhedron_connectivity], [pv.CellType.POLYHEDRON], points)
# TODO: can we directly use the following line to add everything altogether? add_mesh already contain all the kwargs above, like:
# TODO: self.plotter.add_mesh(mesh, color=color, show_edges=show_edges, edge_color=edge_color, line_width=line_width, opacity=data.opacity * opacity, lighting=True, **kwargs)
# TODO: By doing this we can inherent the maximal flexibility from pyvista
self.plotter.add_mesh(
mesh,
color=color,
opacity=data.opacity * opacity,
lighting=True,
show_edges=False,
# vertices options
show_vertices=show_vertices,
vertex_color=vertex_color,
point_size=vertex_size,
render_points_as_spheres=True,
**kwargs)
if show_edges:
edge_mesh = pv.UnstructuredGrid([self.polyhedron_connectivity], [pv.CellType.POLYHEDRON], points)
self.plotter.add_mesh(edge_mesh, style='wireframe', line_width=edge_width, color=edge_color, opacity=data.opacity)
# edge_mesh.active_scalars[:,:-1] = 0 # Set the color to black
[docs]
def plot_ellipsoid(self,
P: Matrix,
c: Optional[Vector] = None,
r: float = 1.0,
color: Optional[str] = 'lightblue',
opacity: Optional[float] = 1.,
show_edges: Optional[bool] = False,
line_width: Optional[float] = 2,
edge_color: Optional[str] = 'black'):
assert P.ndim == 2 and P.shape[0] == P.shape[1], "P must be a square matrix."
assert P.shape[0] in [2, 3], "P must be a 2x2 or 3x3 matrix."
c = np.zeros(P.shape[0]) if c is None else np.array(c)
assert c.shape[0] == P.shape[0], "c must have the same dimension as P."
self.is_3d = True if P.shape[0] == 3 else False
if not self.is_3d:
P = block_diag(P, np.array([0.]))
c = np.concatenate((c, np.zeros(1)))
eigvals, eigvecs = np.linalg.eigh(P) # P = Q * Lambda * Q^T
radii = [r / np.sqrt(eigval) if abs(eigval) > tolerances['zero'] else 0.0 for eigval in eigvals] # deal with lower-dime ellipsoids
ellipsoid = pv.ParametricEllipsoid(*radii)
rotated_points = (ellipsoid.points @ eigvecs.T) # rotate the mesh using the eigenvectors
ellipsoid.points = rotated_points + c
self.plotter.add_mesh(ellipsoid, color=color, opacity=opacity, show_edges=show_edges, edge_color=edge_color, line_width=line_width)
[docs]
def show(self):
plt = self.plotter
plt.show_grid()
if not self.is_3d:
plt.camera_position = 'xy'
plt.enable_2d_style()
plt.enable_parallel_projection()
self.plotter.show()