|
1 | 1 | import matplotlib.pyplot as plt
|
| 2 | +from numpy.core.defchararray import center |
2 | 3 | from spatialmath.base.vectors import getvector
|
3 | 4 | import numpy as np
|
| 5 | +import scipy as sp |
| 6 | + |
| 7 | +# TODO |
| 8 | +# axes_logic everywhere |
| 9 | +# dont do draw |
| 10 | +# return reference to the graphics object |
| 11 | +# don't have own color/style options, go for MPL ones |
| 12 | +# unit tests |
| 13 | +# seealso |
| 14 | +# example code |
| 15 | +# return a redrawer object, that can be used for animation |
| 16 | + |
4 | 17 |
|
5 | 18 | def plot_box(ax=None,
|
6 | 19 | bbox=None, bl=None, tl=None, br=None, tr=None, wh=None, centre=None,
|
7 |
| - color=None, fillcolor=None, alpha=None, thickness=None, **kwargs): |
| 20 | + color=None, filled=True, alpha=None, thickness=None, **kwargs): |
8 | 21 | """
|
9 | 22 | Plot a box using matplotlib
|
10 | 23 |
|
@@ -73,16 +86,20 @@ def plot_box(ax=None,
|
73 | 86 | w, h = wh
|
74 | 87 | xy = (tl[0], tl[1] - h)
|
75 | 88 |
|
76 |
| - if ax is None: |
77 |
| - ax = plt.gca() |
| 89 | + ax = _axes_logic(ax, 2) |
78 | 90 |
|
79 |
| - fill = fillcolor is not None |
80 |
| - rect = plt.Rectangle(xy, w, h, edgecolor=color, facecolor=fillcolor, fill=fill, |
81 |
| - alpha=alpha, linewidth=thickness, clip_on=True) |
82 |
| - ax.add_patch(rect) |
83 |
| - plt.draw() |
| 91 | + if filled: |
| 92 | + r = plt.Rectangle(xy, w, h, edgecolor=color, facecolor=fillcolor, fill=fill, |
| 93 | + alpha=alpha, linewidth=thickness, clip_on=True, **kwargs) |
| 94 | + ax.add_patch(rect) |
| 95 | + else: |
| 96 | + x1 = xy[0] |
| 97 | + x2 = x1 + w |
| 98 | + y1 = xy[1] |
| 99 | + y2 = y1 + h |
| 100 | + r = plt.plot([x1, x1, x2, x2, x1], [y1, y2, y2, y1, y1], **kwargs) |
84 | 101 |
|
85 |
| - return rect |
| 102 | + return r |
86 | 103 |
|
87 | 104 |
|
88 | 105 | def plot_text(pos, text=None, ax=None, color=None, **kwargs):
|
@@ -185,6 +202,159 @@ def plot_point(pos, marker='bs', text=None, ax=None, color=None, textargs=None,
|
185 | 202 | plt.text(x, y, ' ' + text, horizontalalignment='left', verticalalignment='center', color=color, **textopts)
|
186 | 203 |
|
187 | 204 |
|
| 205 | + |
| 206 | + |
| 207 | +def _axes_dimensions(ax): |
| 208 | + if hasattr(ax, 'get_zlim'): |
| 209 | + return 3 |
| 210 | + else: |
| 211 | + return 2 |
| 212 | + |
| 213 | +def circle(centre=(0, 0), radius=1, npoints=50): |
| 214 | + u = np.linspace(0.0, 2.0 * np.pi, npoints) |
| 215 | + x = radius * np.cos(u) + centre[0] |
| 216 | + y = radius * np.sin(u) + centre[1] |
| 217 | + |
| 218 | + return (x, y) |
| 219 | + |
| 220 | +def plot_circle(centre=(0, 0), radius=1, npoints=50, ax=None, filled=False): |
| 221 | + |
| 222 | + x, y = circle(centre, radius, npoints) |
| 223 | + ax = _axes_logic(ax, 2) |
| 224 | + if filled: |
| 225 | + patch = plt.Polygon(x, y, **kwargs) |
| 226 | + ax.add_patch(patch) |
| 227 | + else: |
| 228 | + plt.plot(x, y, **kwargs) |
| 229 | + |
| 230 | +def sphere(centre=(0,0,0), radius=1, npoints=50): |
| 231 | + u = np.linspace(0.0, 2.0 * np.pi, npoints) |
| 232 | + v = np.linspace(0.0, np.pi, npoints) |
| 233 | + |
| 234 | + x = radius * np.outer(np.cos(u), np.sin(v)) + centre[0] |
| 235 | + y = radius * np.outer(np.sin(u), np.sin(v)) + centre[1] |
| 236 | + z = radius * np.outer(np.ones_like(u), np.cos(v)) + centre[2] |
| 237 | + |
| 238 | + return (x, y, z) |
| 239 | + |
| 240 | +def plot_sphere(centre=(0,0,0), radius=1, npoints=50, ax=None, wireframe=False, **kwargs): |
| 241 | + (x, y, z) = _sphere(centre=centre, radius=radius, npoints=npoints) |
| 242 | + |
| 243 | + ax = _axes_logic(ax, 3) |
| 244 | + |
| 245 | + if wireframe: |
| 246 | + ax.plot_wireframe(x, y, z, **kwargs) |
| 247 | + else: |
| 248 | + ax.plot_surface(x, y, z, **kwargs) |
| 249 | + |
| 250 | + |
| 251 | +def ellipse(E, centre=(0,0,0), confidence=None, npoints=40, inverse=False): |
| 252 | + if E.shape != (2,2): |
| 253 | + raise ValueError('ellipse is defined by a 2x2 matrix') |
| 254 | + |
| 255 | + if inverse: |
| 256 | + E = np.linalg.inv(E) |
| 257 | + |
| 258 | + if confidence: |
| 259 | + # process the probability |
| 260 | + s = sqrt(chi2inv(confidence, 2)) |
| 261 | + else: |
| 262 | + s = 1 |
| 263 | + |
| 264 | + x, y = circle() # unit circle |
| 265 | + e = sp.linalg.sqrtm(E) @ np.array([x, y]) |
| 266 | + return e[0,:], e[1,:] |
| 267 | + |
| 268 | +def plot_ellipse(E, centre=(0,0), confidence=None, npoints=40, inverse=False, filled=None, **kwargs): |
| 269 | + |
| 270 | + # allow for centre[2] to plot ellipse in a plane in a 3D plot |
| 271 | + |
| 272 | + x, y = ellipse(E, centre, confidence, npoints, inverse) |
| 273 | + ax = _axes_logic(ax, 2) |
| 274 | + if filled: |
| 275 | + patch = plt.Polygon(x, y, **kwargs) |
| 276 | + ax.add_patch(patch) |
| 277 | + else: |
| 278 | + plt.plot(x, y, **kwargs) |
| 279 | + |
| 280 | +def ellipsoid(E, centre=(0,0,0), confidence=None, npoints=40, inverse=False): |
| 281 | + |
| 282 | + if E.shape != (3,3): |
| 283 | + raise ValueError('ellipsoid is defined by a 3x3 matrix') |
| 284 | + |
| 285 | + if inverse: |
| 286 | + E = np.linalg.inv(E) |
| 287 | + |
| 288 | + if confidence: |
| 289 | + # process the probability |
| 290 | + from scipy.stats.distributions import chi2 |
| 291 | + s = math.sqrt(chi2.ppf(s, df=2)) |
| 292 | + else: |
| 293 | + s = 1 |
| 294 | + |
| 295 | + x, y, z = sphere() # unit sphere |
| 296 | + e = sp.linalg.sqrtm(E) @ np.array([x.flatten(), y.flatten(), z.flatten()]) |
| 297 | + return e[0,:].reshape(x.shape), e[1,:].reshape(x.shape), e[2,:].reshape(x.shape) |
| 298 | + |
| 299 | +def plot_ellipsoid(E, centre=(0,0,0), confidence=None, npoints=40, inverse=False, ax=None, wireframe=False, stride=1, **kwargs): |
| 300 | + """ |
| 301 | + Draw an ellipsoid |
| 302 | +
|
| 303 | + :param E: ellipsoid |
| 304 | + :type E: ndarray(3,3) |
| 305 | + :param centre: [description], defaults to (0,0,0) |
| 306 | + :type centre: tuple, optional |
| 307 | + :param confidence: confidence interval, range 0 to 1 |
| 308 | + :type confidence: float |
| 309 | + :param npoints: [description], defaults to 40 |
| 310 | + :type npoints: int, optional |
| 311 | + :param inverse: [description], defaults to False |
| 312 | + :type inverse: bool, optional |
| 313 | + :param ax: [description], defaults to None |
| 314 | + :type ax: [type], optional |
| 315 | + :param wireframe: [description], defaults to False |
| 316 | + :type wireframe: bool, optional |
| 317 | + :param stride: [description], defaults to 1 |
| 318 | + :type stride: int, optional |
| 319 | +
|
| 320 | +
|
| 321 | + ``plot_ellipse(E)`` draws the ellipsoid defined by :math:`x^T \mat{E} x = 0` |
| 322 | + on the current plot. |
| 323 | +
|
| 324 | + Example: |
| 325 | +
|
| 326 | + H = plot_ellipse(diag([1 2]), [3 4]', 'r'); % draw red ellipse |
| 327 | + plot_ellipse(diag([1 2]), [5 6]', 'alter', H); % move the ellipse |
| 328 | + plot_ellipse(diag([1 2]), [5 6]', 'alter', H, 'LineColor', 'k'); % change color |
| 329 | +
|
| 330 | + plot_ellipse(COVAR, 'confidence', 0.95); % draw 95% confidence ellipse |
| 331 | +
|
| 332 | + .. note:: |
| 333 | +
|
| 334 | + - If a confidence interval is given then ``E`` is interpretted as a covariance |
| 335 | + matrix and the ellipse size is computed using an inverse chi-squared function. |
| 336 | + """ |
| 337 | + x, y, z = ellipsoid(E, centre, confidence, npoints, inverse) |
| 338 | + ax = _axes_logic(ax, 3) |
| 339 | + if wireframe: |
| 340 | + return ax.plot_wireframe(x, y, z, rstride=stride, cstride=stride, **kwargs) |
| 341 | + else: |
| 342 | + return ax.plot_surface(x, y, z, **kwargs) |
| 343 | + |
| 344 | +def _axes_logic(ax, dimensions, projection='ortho'): |
| 345 | + if ax is not None: |
| 346 | + # axis was given |
| 347 | + if _axes_dimensions == dimensions: |
| 348 | + return ax |
| 349 | + # mismatch, create new axes |
| 350 | + |
| 351 | + # no axis specified |
| 352 | + if dimensions == 2: |
| 353 | + ax = plt.axes() |
| 354 | + else: |
| 355 | + ax = plt.axes(projection='3d', proj_type=projection) |
| 356 | + return ax |
| 357 | + |
188 | 358 | def isnotebook():
|
189 | 359 | """
|
190 | 360 | Determine if code is being run from a Jupyter notebook
|
|
0 commit comments