handle animation of pose object with multi values

petercorke · petercorke · commit 41595391f167 · 2021-02-24T11:09:17.000+10:00
diff --git a/spatialmath/base/animate.py b/spatialmath/base/animate.py
@@ -13,7 +13,7 @@
 import matplotlib.pyplot as plt
 from matplotlib import animation
 from spatialmath import base
-from collections.abc import Iterable, Generator
+from collections.abc import Iterable, Generator, Iterator
 
 
 class Animate:
@@ -158,24 +158,32 @@ def run(self, movie=None, axes=None, repeat=False, interval=50, nframes=100, pau
             - invokes the draw() method of every object in the display list
         """
 
-        def update(frame, a):
-            # if contains trajectory:
+        def update(frame, animation):
             if self.trajectory is not None:
+                # passed a trajectory as an iterator or generator, get next
                 T = next(self.trajectory)
             else:
+                # passed a single transform, interpolate it
                 T = base.trinterp(start=self.start, end=self.end, s=frame / nframes)
-            a._draw(T)
+            # ensure result is SE(3)
+            if T.shape == (3,3):
+                T = base.r2t(T)
+
+            # update the scene
+            animation._draw(T)
+
+            # are we done yet
             if frame == nframes - 1:
-                a.done = True
-            return a.artists()
+                animation.done = True
+            return animation.artists()
 
         # blit leaves a trail and first frame
         if movie is not None:
             repeat = False
 
         self.done = False
         if self.trajectory is not None:
-            if not isinstance(self.trajectory, Iterable):
+            if not isinstance(self.trajectory, Iterator):
                 # make it iterable, eg. if a list or tuple
                 self.trajectory = iter(self.trajectory)
             frames = None
diff --git a/spatialmath/super_pose.py b/spatialmath/super_pose.py
@@ -851,10 +851,19 @@ def animate(self, *args, start=None, **kwargs):
         """
         if start is not None:
             start = start.A
-        if self.N == 2:
-            base.tranimate2(self.A, start=start, *args, **kwargs)
+        
+        if len(self) > 1:
+            # trajectory case
+            if self.N == 2:
+                base.tranimate2(self.data, *args, **kwargs)
+            else:
+                base.tranimate(self.data, *args, **kwargs)
         else:
-            base.tranimate(self.A, start=start, *args, **kwargs)
+            # singleton case
+            if self.N == 2:
+                base.tranimate2(self.A, start=start, *args, **kwargs)
+            else:
+                base.tranimate(self.A, start=start, *args, **kwargs)
 
 
 # ------------------------------------------------------------------------ #
