8000 working version of circle, sphere, ellipse, ellipsoid primitives · krenshaw2018/spatialmath-python@d1dc9c4 · GitHub
[go: up one dir, main page]

Skip to content

Commit d1dc9c4

Browse files
committed
working version of circle, sphere, ellipse, ellipsoid primitives
needs work
1 parent 78ded48 commit d1dc9c4

File tree

2 files changed

+189
-9
lines changed

2 files changed

+189
-9
lines changed

spatialmath/base/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,8 @@
9090
'trinv',
9191
'tr2delta',
9292
'tr2jac',
93+
'rpy2jac',
94+
'eul2jac',
9395
'trprint',
9496
'trplot',
9597
'tranimate',
@@ -136,6 +138,14 @@
136138
'plot_point',
137139
'plot_text',
138140
'plot_box',
141+
'circle',
142+
'ellipse',
143+
'sphere',
144+
'ellipsoid',
145+
'plot_circle',
146+
'plot_ellipse',
147+
'plot_sphere',
148+
'plot_ellipsoid',
139149
'isnotebook',
140150

141151
]

spatialmath/base/graphics.py

Lines changed: 179 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,23 @@
11
import matplotlib.pyplot as plt
2+
from numpy.core.defchararray import center
23
from spatialmath.base.vectors import getvector
34
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+
417

518
def plot_box(ax=None,
619
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):
821
"""
922
Plot a box using matplotlib
1023
@@ -73,16 +86,20 @@ def plot_box(ax=None,
7386
w, h = wh
7487
xy = (tl[0], tl[1] - h)
7588

76-
if ax is None:
77-
ax = plt.gca()
89+
ax = _axes_logic(ax, 2)
7890

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)
84101

85-
return rect
102+
return r
86103

87104

88105
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,
185202
plt.text(x, y, ' ' + text, horizontalalignment='left', verticalalignment='center', color=color, **textopts)
186203

187204

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+
188358
def isnotebook():
189359
"""
190360
Determine if code is being run from a Jupyter notebook

0 commit comments

Comments
 (0)
0