Skip to content

Commit e238570

Browse files
authored
Merge pull request rai-opensource#13 from mfkenson/master
tranimate/tranimate2/animate.trplot/animate.trplot2 support iterable
2 parents 9eda767 + 12cf3db commit e238570

1 file changed

Lines changed: 30 additions & 3 deletions

File tree

‎spatialmath/base/animate.py‎

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import matplotlib.pyplot as plt
1414
from matplotlib import animation
1515
from spatialmath import base
16+
from collections.abc import Iterable
1617

1718

1819
class Animate:
@@ -57,6 +58,7 @@ def __init__(self, axes=None, dims=None, projection='ortho', labels=('X', 'Y', '
5758
Will setup to plot into an existing or a new Axes3D instance.
5859
5960
"""
61+
self.trajectory = None
6062
self.displaylist = []
6163

6264
if axes is None:
@@ -104,6 +106,12 @@ def trplot(self, end, start=None, **kwargs):
104106
:seealso: :func:`run`
105107
106108
"""
109+
if not isinstance(end, (np.ndarray, np.generic) ) and isinstance(end, Iterable):
110+
if len(end) == 1:
111+
end = end[0]
112+
elif len(end) >= 2:
113+
self.trajectory = end
114+
107115
# stash the final value
108116
if base.isrot(end):
109117
self.end = base.r2t(end)
@@ -129,7 +137,7 @@ def run(self, movie=None, axes=None, repeat=False, interval=50, nframes=100, pau
129137
:type axes: Axes3D reference
130138
:param repeat: animate in endless loop [default False]
131139
:type repeat: bool
132-
:param nframes: number of steps in the animation [defaault 100]
140+
:param nframes: number of steps in the animation [default 100]
133141
:type nframes: int
134142
:param interval: number of milliseconds between frames [default 50]
135143
:type interval: int
@@ -147,7 +155,11 @@ def run(self, movie=None, axes=None, repeat=False, interval=50, nframes=100, pau
147155
"""
148156

149157
def update(frame, a):
150-
T = base.trinterp(start=self.start, end=self.end, s=frame / nframes)
158+
# if contains trajectory:
159+
if self.trajectory is not None:
160+
T = self.trajectory[frame]
161+
else:
162+
T = base.trinterp(start=self.start, end=self.end, s=frame / nframes)
151163
a._draw(T)
152164
if frame == nframes - 1:
153165
a.done = True
@@ -158,6 +170,8 @@ def update(frame, a):
158170
repeat = False
159171

160172
self.done = False
173+
if self.trajectory is not None:
174+
nframes = len(self.trajectory)
161175
ani = animation.FuncAnimation(fig=plt.gcf(), func=update, frames=range(0, nframes), fargs=(self,), blit=False, interval=interval, repeat=repeat)
162176
if movie is None:
163177
while repeat or not self.done:
@@ -398,6 +412,7 @@ def __init__(self, axes=None, dims=None, labels=('X', 'Y'), **kwargs):
398412
Will setup to plot into an existing or a new Axes3D instance.
399413
400414
"""
415+
self.trajectory = None
401416
self.displaylist = []
402417

403418
if axes is None:
@@ -439,6 +454,12 @@ def trplot2(self, end, start=None, **kwargs):
439454
:seealso: :func:`run`
440455
441456
"""
457+
if not isinstance(end, (np.ndarray, np.generic) ) and isinstance(end, Iterable):
458+
if len(end) == 1:
459+
end = end[0]
460+
elif len(end) >= 2:
461+
self.trajectory = end
462+
442463
# stash the final value
443464
if base.isrot2(end):
444465
self.end = base.r2t(end)
@@ -482,7 +503,11 @@ def run(self, movie=None, axes=None, repeat=False, interval=50, nframes=100, **k
482503
"""
483504

484505
def update(frame, a):
485-
T = base.trinterp2(start=self.start, end=self.end, s=frame / nframes)
506+
# if contains trajectory:
507+
if self.trajectory is not None:
508+
T = self.trajectory[frame]
509+
else:
510+
T = base.trinterp2(start=self.start, end=self.end, s=frame / nframes)
486511
a._draw(T)
487512
if frame == nframes - 1:
488513
a.done = True
@@ -493,6 +518,8 @@ def update(frame, a):
493518
repeat = False
494519

495520
self.done = False
521+
if self.trajectory is not None:
522+
nframes = len(self.trajectory)
496523
ani = animation.FuncAnimation(fig=plt.gcf(), func=update, frames=range(0, nframes), fargs=(self,), blit=False, interval=interval, repeat=repeat)
497524

498525
if movie is None:

0 commit comments

Comments
 (0)