Source code for mpt4py.geometry.visualization.plot

"""
Plotting protocol for MPT library.
"""

import time

import numpy as np
from typing import Tuple, List, Optional
from dataclasses import dataclass
from collections import defaultdict
from abc import ABC, abstractmethod
from scipy.spatial.distance import pdist

from scipy.spatial import ConvexHull
from scipy.linalg import svd
from sklearn.cluster import AgglomerativeClustering

import cvxpy as cp

from mpt4py.base import Matrix, Vector, IntVector
from mpt4py.tolerances import tolerances
from mpt4py.exceptions import GeometryError

[docs] class PlotterProtocol(ABC):
[docs] def plot(self, object, **kwargs): object.plot(self, **kwargs)
[docs] @abstractmethod def plot_convexhull(self, points: Matrix, rays: Optional[Matrix] = None, color: Optional[str] = None, opacity: Optional[float] = 1.): pass
[docs] @abstractmethod 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'): pass
[docs] @abstractmethod def show(self): pass
# @abstractmethod # def next_color(self): # pass
[docs] @dataclass class PolyhedralPlottingData: """ Represents the data needed to plot a polyhedron. The representation is always in 3d. A 2d polyhedron will be represented as a single facet with an affine hull of z=0. """ hull: ConvexHull # The convex hull of the points and rays points: Matrix # The original data specified by the user rays: Matrix # The original rays specified by the user combined_points: Matrix # [points; rays] # A list of indices into the vertices matrix that represent the facets # in counter-clockwise order facets: List[IntVector] # Outward facing unit normals to the facets facet_eqns: Matrix # [normal offset] -> <normal, x> = offset ray_length: float = 1.0 # Length of the rays def __init__(self, points: Matrix, rays: Optional[Matrix] = None): """ Computes all data needed to plot the Polyhedron on initialization. """ if points.size == 0: raise GeometryError("We can only plot pointed polyhedra (those with at least one vertex).") self.points = points self.num_points = self.points.shape[0] self.rays = rays.astype(np.float64) if rays is not None else np.empty((0, points.shape[1]), dtype=np.float64) self.rays /= np.linalg.norm(self.rays, axis=1).reshape((-1, 1)) # Normalize the rays self.num_rays = self.rays.shape[0] # We use a ball around this point to determine how big to make the rays self.center_point = np.mean(self.points, axis=0) if self.points.shape[1] < 2 or self.points.shape[1] > 3: raise ValueError('Can only plot 2D and 3D polyhedra. Project or slice your polyhedron first.') # Compute the diameter of the bounded part of the set dist = pdist(self.points) if dist.size == 0: self.ray_length = 1. else: self.diameter = float(np.max(dist)) self.ray_length = self.diameter * 3. # Homogenize the points and rays self.homo = self._homogenize(self.points, self.rays) # Points used in the actual plot self.combined_points = np.vstack((self.points, self.rays)) if self.combined_points.shape[1] == 2: npoints = self.combined_points.shape[0] self.combined_points = np.hstack((self.combined_points, np.zeros((npoints, 1)))) # Project the homogenized points into the affine hull p0, aff_normal, aff_basis = self._compute_affine_hull(self.homo) projected_points = self.homo @ aff_basis # Compute the convex hull (This also does redundancy elimination) self.hull = ConvexHull(projected_points) dim = self.hull.simplices.shape[1] self.facets = [] # List of facets = list of vertices in counterclockwise order if dim == 0: pass elif dim == 1: pass elif dim == 2: # 2D object vertices2d = projected_points[self.hull.vertices] order = self._order_vertices_counterclockwise(vertices2d, self.hull.vertices) self.facets.append(order) elif dim == 3: # Compute the facets and the vertices on each facet normalized_equations = self.hull.equations group_labels = AgglomerativeClustering(n_clusters=None, distance_threshold=1e-6, linkage='single').fit_predict(normalized_equations) # Facets is a map from facet to the list of vertices on that facet num_facets = max(group_labels) + 1 facets = [set() for i in range(num_facets)] facet_eqns = [None for i in range(num_facets)] for i, facet_num in enumerate(group_labels): facets[facet_num].update(self.hull.simplices[i].tolist()) normal = normalized_equations[i] if facet_eqns[facet_num] is None: facet_eqns[facet_num] = normal else: if np.linalg.norm(normal - facet_eqns[facet_num]) > 1e-6: raise GeometryError("Facet normals are not consistent - could not plot polyhedron") self.facet_eqns = np.array(facet_eqns) facets = [np.array(list(facet)) for facet in facets] polyhedron_connectivity = [] self.facets = [] for facet, eqn in zip(facets, self.facet_eqns): if np.all(facet > self.num_points - 1): # This is a facet at infinity as it only involves rays -> we don't need to plot it continue points2d = self._project_into_facet(projected_points[facet], eqn[:-1]) ordered_facet = self._order_vertices_counterclockwise(points2d, facet) self.facets.append(ordered_facet) else: raise GeometryError("We got an affine dimension larger than three. This shouldn't be possible") # Fade out any rays self.opacity = np.ones(self.points.shape[0]) self.opacity[self.num_points:] = 0 @property def extreme_points(self) -> Tuple[Matrix, Matrix]: """ Returns (Vertices, Rays) """ return self.points[self.hull.vertices], np.empty((0, self.points.shape[1])) @property def vertices(self) -> Matrix: """ Returns vertices """ return self.extreme_points[0] @property def extreme_rays(self) -> Matrix: """ Returns the extreme rays (irredundant rays) """ return self.extreme_points[1] @property def num_facets(self): return len(self.facets) @property def num_vertices(self): return self.extreme_points[0].shape[0] @property def num_extreme_rays(self): return self.extreme_points[1].shape[0] @property def affine_hull(self) -> Tuple[Vector, Matrix, Matrix]: r""" Returns :math:`p_0, A_e`, Basis such that :math:`A_e * (x - p_0) = 0` and :math:`x = B * \lambda` for all :math:`x\in P`. """ return self._affine_hull def _compute_affine_hull(self, points: Matrix) -> Tuple[Vector, Matrix, Matrix]: # Choose the first point as the reference p0 = points[0] # Should be [0,0,0,0] # Translate points by subtracting p0 translated_points = points - p0 # We want an Ae s.t. Ae * (points - p0) = 0 # U * S * V' = (points - p0) # S * V' = U' * (points - p0) # So we take the last columns of U as the basis of the null space U, s, _ = svd(translated_points.T, full_matrices=True) rank = np.sum(np.abs(s) > tolerances['affine_hull'], axis=0) self._affine_hull = p0, U[:, rank:].T, U[:, :rank] return self.affine_hull def _homogenize(self, points: Matrix, rays: Matrix) -> Matrix: """ Homogenize the points and rays by appending 1 and 0 respectively. Given the polyhedron {x | x = V' * lam + R' * gam, lam, gam >= 0, 1' * lam = 1} The homogenized cone is: {y | y = [V 1; R 0]' lam, lam >= 0} We cut the cone by the hyperplane {x | alpha' * x = 1} to get the polyhedron {y | y = [V 1; R 0]' lam, lam >= 0, [V 1; R 0] * alpha = 1} where alpha = mean([V 1; R 0]) so that we maintain the connectivity diagram of the homogenized cone, but we have a bounded polytope. """ def make_homo(points: Matrix, lifter) -> Matrix: return np.hstack((points, lifter(points.shape[0])[:, np.newaxis])) homo = np.vstack((make_homo(points, np.ones), make_homo(rays, np.zeros))) # Cut every ray through alpha'x = 1 alpha = np.mean(homo, axis=0) alpha /= np.linalg.norm(alpha) scl = homo @ alpha if np.any(scl < 1e-3): # The mean isn't a good choice, we've got to do this the hard way alpha = cp.Variable(alpha.size) epsilon = cp.Variable() prob = cp.Problem(cp.Maximize(epsilon), [homo @ alpha.T >= epsilon, cp.norm(alpha, "inf") <= 1]) result = prob.solve() if prob.status != cp.OPTIMAL: raise GeometryError("Polyhedron appears to be unpointed (no vertices)") alpha = alpha.value alpha /= np.linalg.norm(alpha) scl = homo @ alpha if np.any(scl < tolerances['zero']): raise GeometryError("Polyhedron appears to be unpointed (no vertices)") return homo / scl[:, np.newaxis] def _project_into_facet(self, points: Matrix, normal: Vector) -> Matrix: """ Project the points[facet] into the 2D plane defined by the normal. 3D -> 2D projection """ # Find two orthogonal vectors on the plane u = np.cross(normal, [1, 0, 0]) if np.linalg.norm(u) < 1e-6: u = np.cross(normal, [0, 1, 0]) v = np.cross(normal, u) u = u / np.linalg.norm(u) v = v / np.linalg.norm(v) # Project the facet points onto the 2D plane return np.dot(points, np.vstack([u, v]).T) def _scale_ray(self, point: Vector, ray: Vector) -> Vector: """ Compute a scaling of the ray such that: ||self.center_points - (point + ray * scl)|| == self.ray_length """ c = self.center_point p = point r = ray[:len(point)] A = c - p b = np.dot(A, r) A_squared = np.dot(A, A) r_squared = np.dot(r, r) # Compute coefficients for the quadratic equation C = A_squared - self.ray_length**2 # Compute the discriminant discriminant = b**2 - r_squared * C if discriminant < 0: raise ValueError("The discriminant is negative. No real solution exists.") # Compute the two possible values for scl sqrt_discriminant = np.sqrt(discriminant) scl1 = (b + sqrt_discriminant) / r_squared scl2 = (b - sqrt_discriminant) / r_squared return max(scl1, scl2) * ray def _order_vertices_counterclockwise(self, points2d: Matrix, indices: IntVector) -> IntVector: # Compute the centroid of the projected points centroid = np.mean(points2d, axis=0) # Compute the angles of the points relative to the centroid angles = np.arctan2(points2d[:, 1] - centroid[1], points2d[:, 0] - centroid[0]) # Sort the vertices by angle in counter-clockwise order vertices = indices[np.argsort(angles)] # We're in 2D here, so there can be a maximum of one ray in the facet ray_indices = np.sort(np.where(vertices > self.num_points - 1)[0]) if len(ray_indices) == 1: vertices = np.insert(vertices, ray_indices[0], [vertices[ray_indices[0]]]) ray_indices = np.sort(np.where(vertices > self.num_points - 1)[0]) if len(ray_indices) > 2: raise GeometryError("Facet has more than two rays. This should not be possible.") if len(ray_indices) == 2: # Add new vertices to the facet for the rays # Rotate the rays so that they appear at the end if ray_indices[0] == 0 and ray_indices[1] == len(vertices) - 1: # We span the end of the list vertices = np.roll(vertices, -1) else: vertices = np.roll(vertices, len(vertices)-2-ray_indices[0]) ray_indices = np.sort(np.where(vertices > self.num_points - 1)[0]) assert all(ray_indices == [len(vertices) - 2, len(vertices) - 1]), "Rays are not at the end of the facet... something screwed up" # Compute the adjacent vertices + ray, and add to the facet rays = self.combined_points[vertices[-2:]] before = vertices[-3] ray = self._scale_ray(self.points[before], rays[0])[:self.points.shape[1]] self.points = np.vstack((self.points, self.points[before] + ray)) vertices[-2] = len(self.points) - 1 after = vertices[0] ray = self._scale_ray(self.points[after], rays[1])[:self.points.shape[1]] self.points = np.vstack((self.points, self.points[after] + ray)) vertices[-1] = len(self.points) - 1 return vertices