Source code for spatialmath.base.animate

# Part of Spatial Math Toolbox for Python
# Copyright (c) 2000 Peter Corke
# MIT Licence, see details in top-level file: LICENCE

# matplotlib inline

# line.set_data()
# text.set_position()
# quiver.set_offsets(), quiver.set_UVC()
# FancyArrow.set_xy()
from __future__ import annotations
import os.path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
import spatialmath.base as smb
from collections.abc import Iterable, Iterator
from spatialmath.base.types import *

# global variable holds reference to FuncAnimation object, this is essential
# for animatiion to work
_ani = None


[docs]class Animate: """ Animate objects for matplotlib 3d An instance of this class behaves like an Axes3D and supports proxies for - ``plot`` - ``quiver`` - ``text`` - ``scatter`` which renders them and also places corresponding objects into a display list. These objects are ``Line``, ``Quiver`` and ``Text``. Only these primitives will be animated. The objects are all drawn relative to the origin, and will be transformed according to the transform that is being animated. Example:: anim = animate.Animate(dims=[0,2]) # set up the 3D axes anim.trplot(T, frame='A', color='green') # draw the frame anim.run(repeat=True) # animate it """
[docs] def __init__( self, ax: Optional[plt.Axes] = None, dim: Optional[ArrayLike] = None, projection: Optional[str] = "ortho", labels: Optional[Tuple[str, str, str]] = ("X", "Y", "Z"), **kwargs, ): """ Construct an Animate object :param ax: the axes to plot into, defaults to current axes :type ax: Axes3D reference :param dim: dimension of plot volume as [xmin, xmax, ymin, ymax, zmin, zmax]. If dims is [min, max] those limits are applied to the x-, y- and z-axes. :type dim: array_like(6) or array_like(2) :param projection: 3D projection: ortho [default] or persp :type projection: str :param labels: labels for the axes, defaults to X, Y and Z :type labels: 3-tuple of strings Will setup to plot into an existing or a new Axes3D instance. """ self.trajectory = None self.displaylist = [] if ax is None: # # no axes specified # fig = plt.gcf() # # check any current axes # for a in fig.axes: # if a.name != "3d": # # if they are not 3D axes, remove them, otherwise will # # get plot errors # a.remove() # if len(fig.axes) == 0: # # no axes in the figure, create a 3D axes # axes = fig.add_subplot(111, projection="3d", proj_type=projection) # ax.set_xlabel(labels[0]) # ax.set_ylabel(labels[1]) # ax.set_zlabel(labels[2]) # ax.autoscale(enable=True, axis="both") # else: # # reuse an existing axis # axes = plt.gca() # if dims is not None: # if len(dims) == 2: # dims = dims * 3 # ax.set_xlim(dims[0:2]) # ax.set_ylim(dims[2:4]) # ax.set_zlim(dims[4:6]) # # ax.set_aspect('equal') ax = smb.plotvol3(ax=ax, dim=dim) if dim is not None: dim = list(np.ndarray.flatten(np.array(dim))) if len(dim) == 2: dim = dim * 3 elif len(dim) != 6: raise ValueError(f"dim must have 2 or 6 elements, got {dim}. See docstring for details.") ax.set_xlim(dim[0:2]) ax.set_ylim(dim[2:4]) ax.set_zlim(dim[4:]) self.ax = ax
# TODO set flag for 2d or 3d axes, flag errors on the methods called later
[docs] def trplot( self, end: Union[SO3Array, SE3Array], start: Optional[Union[SO3Array, SE3Array]] = None, **kwargs, ): """ Define the transform to animate :param end: the final pose SE(3) or SO(3) to display as a coordinate frame :type end: ndarray(4,4) or ndarray(3,3) :param start: the initial pose SE(3) or SO(3) to display as a coordinate frame, defaults to null :type start: ndarray(4,4) or ndarray(3,3) :param start: an Is polymorphic with ``base.trplot`` and accepts the same parameters. This sets up the animation but doesn't execute it. :seealso: :func:`run` """ self.trajectory = None if not isinstance(end, (np.ndarray, np.generic)) and isinstance(end, Iterable): try: if len(end) == 1: end = end[0] elif len(end) >= 2: self.trajectory = end except TypeError: # a generator has no len() self.trajectory = end # stash the final value if smb.isrot(end): self.end = smb.r2t(end) else: self.end = end if start is None: self.start = np.identity(4) else: if smb.isrot(start): self.start = smb.r2t(start) else: self.start = start # draw axes at the origin smb.trplot(self.start, ax=self, **kwargs)
[docs] def set_proj_type(self, proj_type: str): self.ax.set_proj_type(proj_type)
[docs] def run( self, movie: Optional[str] = None, axes: Optional[plt.Axes] = None, repeat: bool = False, interval: int = 50, nframes: int = 100, wait: bool = False, **kwargs, ): """ Run the animation :param axes: the axes to plot into, defaults to current axes :type axes: Axes3D reference :param repeat: animate in endless loop [default False] :type repeat: bool :param nframes: number of steps in the animation [default 100] :type nframes: int :param interval: number of milliseconds between frames [default 50] :type interval: int :param movie: name of file to write MP4 movie into, or True :type movie: str, bool :param wait: wait until animation is complete, default False :type wait: bool Animates a 3D coordinate frame moving from the world frame to a frame represented by the SO(3) or SE(3) matrix to the current axes. .. note:: - the ``movie`` option requires the ffmpeg package to be installed: ``conda install -c conda-forge ffmpeg`` - if ``movie=True`` then return an HTML5 video which can be displayed in a notebook using ``HTML()`` - invokes the draw() method of every object in the display list """ def update(frame, animation): # frame is the result of calling next() on a iterator or generator # seemingly the animation framework isn't checking StopException # so there is no way to know when this is no longer called. # we implement a rather hacky heartbeat style timeout if isinstance(frame, float): # passed a single transform, interpolate it T = smb.trinterp(start=self.start, end=self.end, s=frame) elif isinstance(frame, NDArray): # type is SO3Array or SE3Array when Animate.trajectory is not None T = frame else: # [unlikely] other types are converted to np array T = np.array(frame) if T.shape == (3, 3): T = smb.r2t(T) # update the scene animation._draw(T) self.count += 1 # say we're still running if movie is not None: repeat = False self.count = 1 if self.trajectory is not None: if not isinstance(self.trajectory, Iterator): # make it iterable, eg. if a list or tuple self.trajectory = iter(self.trajectory) frames = self.trajectory else: frames = iter(np.linspace(0, 1, nframes)) global _ani fig = self.ax.get_figure() _ani = animation.FuncAnimation( fig=fig, func=update, frames=frames, fargs=(self,), # blit=False, # blit leaves a trail and first frame, set to False interval=interval, repeat=repeat, save_count=nframes, ) if movie is True: plt.close(fig) return _ani.to_html5_video() elif isinstance(movie, str): # Set up formatting for the movie files if os.path.exists(movie): print("overwriting movie", movie) else: print("creating movie", movie) FFwriter = animation.FFMpegWriter( fps=1000 / interval, extra_args=["-vcodec", "libx264"] ) _ani.save(movie, writer=FFwriter) if wait: # wait for the animation to finish. Dig into the timer for this # animation and wait for its callback to be deregistered. while True: plt.pause(0.25) if _ani.event_source is None or len(_ani.event_source.callbacks) == 0: break return _ani
[docs] def __repr__(self) -> str: """ Human readable version of the display list :param self: the animation :type self: Animate :returns: readable version of the display list :rtype: str """ return "Animate(" + ", ".join([x.type for x in self.displaylist]) + ")"
[docs] def __str__(self) -> str: return f"Animate(len={len(self.displaylist)}"
[docs] def artists(self) -> List[plt.Artist]: """ List of artists that need to be updated :param self: the animation :type self: Animate :returns: list of artists :rtype: list """ return [x.h for x in self.displaylist]
def _draw(self, T): for x in self.displaylist: x.draw(T) # ------------------- plot() class _Line: def __init__(self, anim: Animate, h, xs, ys, zs): # form 4xN matrix, columns are first/last point in homogeneous form p = np.vstack([xs, ys, zs]) self.p = np.vstack([p, np.ones((p.shape[1],))]) self.h = h self.type = "line" self.anim = anim def draw(self, T): p = T @ self.p self.h.set_data(p[0, :], p[1, :]) self.h.set_3d_properties(p[2, :])
[docs] def plot(self, x: ArrayLike, y: ArrayLike, z: ArrayLike, *args: List, **kwargs): """ Plot a polyline :param x: list of x-coordinates :type x: array_like :param y: list of y-coordinates :type y: array_like :param z: list of z-coordinates :type z: array_like Other arguments as accepted by the matplotlib method. All arrays must have the same length. :seealso: :func:`matplotlib.pyplot.plot` """ (h,) = self.ax.plot(x, y, z, *args, **kwargs) self.displaylist.append(Animate._Line(self, h, x, y, z)) return h
# ------------------- quiver() class _Quiver: def __init__(self, anim, h): self.type = "quiver" self.anim = anim # for matplotlib 3.1.x # ._segments3d is 3x2x3 # first index: line segment in the collection # second index: 0 = start, 1 = end # third index: x, y, z components # https://stackoverflow.com/questions/48911643/set-uvc-equivilent-for-a-3d-quiver-plot-in-matplotlib # # for matplotlib 3.3.x # ._segments3d is a 3-element list, each element is 2x3 # turn to homogeneous form, with columns per point, alternating start, end if isinstance(h._segments3d, np.ndarray): self.p = np.vstack( [h._segments3d.reshape(6, 3).T, np.ones((1, 6))] ) # result is 4x6 else: self.p = np.vstack( [np.hstack([x.T for x in h._segments3d]), np.ones((1, 6))] ) self.h = h self.type = "arrow" self.anim = anim def draw(self, T): # import ipdb; ipdb.set_trace() p = T @ self.p # reshape it p = p[0:3, :].T.reshape(3, 2, 3) self.h.set_segments(p)
[docs] def quiver( self, x: ArrayLike, y: ArrayLike, z: ArrayLike, u: ArrayLike, v: ArrayLike, w: ArrayLike, *args: List, **kwargs, ): """ Plot a quiver :param x: list of base x-coordinates :type x: array_like :param y: list of base y-coordinates :type y: array_like :param z: list of base z-coordinates :type z: array_like :param u: list of vector x-coordinates :type u: array_like :param v: list of vector y-coordinates :type v: array_like :param w: list of vector z-coordinates :type w: array_like Draws a series of arrows, the bases defined by corresponding elements of (x,y,z) and the vector has components defined by corresponding elements of (u,v,w). Other arguments as accepted by the matplotlib method. :seealso: :func:`matplotlib.pyplot.quiver` """ h = self.ax.quiver(x, y, z, u, v, w, *args, **kwargs) self.displaylist.append(Animate._Quiver(self, h))
# ------------------- text() class _Text: def __init__(self, anim, h, x, y, z): self.type = "text" self.h = h self.p = np.r_[x, y, z, 1] self.anim = anim def draw(self, T): p = T @ self.p # x2, y2, _ = proj3d.proj_transform( # p[0], p[1], p[2], self.anim.ax.get_proj()) # self.h.set_position((x2, y2)) self.h.set_position((p[0], p[1])) self.h.set_3d_properties(z=p[2], zdir="x")
[docs] def text(self, x: float, y: float, z: float, *args: List, **kwargs): """ Plot text :param x: x-coordinate :type x: float :param y: float :type y: float :param z: z-coordinate :type z: float :param kwargs: Other arguments as accepted by the matplotlib method. ``.text(x, y, z, s)`` display the string ``s`` at coordinate (``x``, ``y``, ``z``). :seealso: :func:`~matplotlib.pyplot.text` """ h = self.ax.text3D(x, y, z, *args, **kwargs) self.displaylist.append(Animate._Text(self, h, x, y, z))
# ------------------- scatter()
[docs] def scatter( self, xs: ArrayLike, ys: ArrayLike, zs: ArrayLike, s: float = 0, **kwargs ): h = self.plot(xs, ys, zs, ".", markersize=0, **kwargs) self.displaylist.append(Animate._Line(self, h, xs, ys, zs))
# ------------------- wrappers for Axes primitives
[docs] def set_xlim(self, *args: List, **kwargs): self.ax.set_xlim(*args, **kwargs)
[docs] def set_ylim(self, *args: List, **kwargs): self.ax.set_ylim(*args, **kwargs)
[docs] def set_zlim(self, *args: List, **kwargs): self.ax.set_zlim(*args, **kwargs)
[docs] def set_xlabel(self, *args: List, **kwargs): self.ax.set_xlabel(*args, **kwargs)
[docs] def set_ylabel(self, *args: List, **kwargs): self.ax.set_ylabel(*args, **kwargs)
[docs] def set_zlabel(self, *args: List, **kwargs): self.ax.set_zlabel(*args, **kwargs)
[docs]class Animate2: """ Animate objects for matplotlib 2d An instance of this class behaves like an Axes3D and supports proxies for - ``plot`` - ``quiver`` - ``text`` which renders them and also places corresponding objects into a display list. These objects are ``Line``, ``Quiver`` and ``Text``. Only these primitives will be animated. The objects are all drawn relative to the origin, and will be transformed according to the transform that is being animated. Example:: anim = animate.Animate(dims=[0,2]) # set up the 3D axes anim.trplot(T, frame='A', color='green') # draw the frame anim.run(loop=True) # animate it """
[docs] def __init__( self, axes: Optional[plt.Axes] = None, dims: Optional[ArrayLike] = None, labels: Tuple[str, str] = ("X", "Y"), **kwargs, ): """ Construct an Animate object :param axes: the axes to plot into, defaults to current axes :type axes: Axes3D reference :param dims: dimension of plot volume as [xmin, xmax, ymin, ymax]. If dims is [min, max] those limits are applied to the x- and y-axes. :type dims: array_like(4) or array_like(2) :param projection: 3D projection: ortho [default] or persp :type projection: str :param labels: labels for the axes, defaults to X, Y and Z :type labels: 3-tuple of strings Will setup to plot into an existing or a new Axes3D instance. """ self.trajectory = None self.displaylist = [] if axes is None: # create an axes fig = plt.gcf() if fig.axes is None: # no axes in the figure, create a 3D axes axes = fig.add_subplot(111) axes.set_xlabel(labels[0]) axes.set_ylabel(labels[1]) axes.autoscale(enable=True, axis="both") else: # reuse an existing axis axes = plt.gca() if dims is not None: if len(dims) == 2: dims = dims * 2 axes.set_xlim(dims[0:2]) axes.set_ylim(dims[2:4]) # ax.set_aspect('equal') self.ax = axes
# set flag for 2d or 3d axes, flag errors on the methods called later
[docs] def trplot2( self, end: Union[SO2Array, SE2Array], start: Optional[Union[SO2Array, SE2Array]] = None, **kwargs, ): """ Define the transform to animate :param end: the final pose SE(2) or SO(2) to display as a coordinate frame :type end: ndarray(3,3) or ndarray(2,2) :param start: the initial pose SE(2) or SO(2) to display as a coordinate frame, defaults to null :type start: ndarray(3,3) or ndarray(2,2) Is polymorphic with ``base.trplot`` and accepts the same parameters. This sets up the animation but doesn't execute it. :seealso: :func:`run` """ if not isinstance(end, (np.ndarray, np.generic)) and isinstance(end, Iterable): if len(end) == 1: end = end[0] elif len(end) >= 2: self.trajectory = end # stash the final value if smb.isrot2(end): self.end = smb.r2t(end) else: self.end = end if start is None: self.start = np.identity(3) else: if smb.isrot2(start): self.start = smb.r2t(start) else: self.start = start # draw axes at the origin smb.trplot2(self.start, ax=self, block=False, **kwargs)
[docs] def run( self, movie: Optional[str] = None, axes: Optional[plt.Axes] = None, repeat: bool = False, interval: int = 50, nframes: int = 100, wait: bool = False, **kwargs ): """ Run the animation :param axes: the axes to plot into, defaults to current axes :type axes: Axes reference :param nframes: number of steps in the animation [defaault 100] :type nframes: int :param repeat: animate in endless loop [default False] :type repeat: bool :param interval: number of milliseconds between frames [default 50] :type interval: int :param movie: name of file to write MP4 movie into or True :type movie: str, bool :returns: Matplotlib animation object :rtype: Matplotlib animation object Animates a 3D coordinate frame moving from the world frame to a frame represented by the SO(2) or SE(2) matrix to the current axes. .. note:: - the ``movie`` option requires the ffmpeg package to be installed: ``conda install -c conda-forge ffmpeg`` - if ``movie=True`` then return an HTML5 video which can be displayed in a notebook using ``HTML()`` - invokes the draw() method of every object in the display list """ def update(frame, animation): # frame is the result of calling next() on a iterator or generator # seemingly the animation framework isn't checking StopException # so there is no way to know when this is no longer called. # we implement a rather hacky heartbeat style timeout if isinstance(frame, float): # passed a single transform, interpolate it T = smb.trinterp2(start=self.start, end=self.end, s=frame) else: # assume it is an SO(2) or SE(2) T = frame # ensure result is SE(2) if T.shape == (2, 2): T = smb.r2t(T) # update the scene animation._draw(T) self.count += 1 # say we're still running if movie is not None: repeat = False self.count = 1 if self.trajectory is not None: if not isinstance(self.trajectory, Iterator): # make it iterable, eg. if a list or tuple self.trajectory = iter(self.trajectory) frames = self.trajectory else: frames = iter(np.linspace(0, 1, nframes)) global _ani fig = self.ax.get_figure() _ani = animation.FuncAnimation( fig=fig, func=update, frames=frames, fargs=(self,), # blit=False, interval=interval, repeat=repeat, save_count=nframes, ) if movie is True: plt.close(fig) return _ani.to_html5_video() elif isinstance(movie, str): # Set up formatting for the movie files if os.path.exists(movie): print("overwriting movie", movie) else: print("creating movie", movie) FFwriter = animation.FFMpegWriter(fps=1000 / interval, extra_args=["-vcodec", "libx264"]) _ani.save(movie, writer=FFwriter) if wait: # wait for the animation to finish. Dig into the timer for this # animation and wait for its callback to be deregistered. while True: plt.pause(0.25) if _ani.event_source is None or len(_ani.event_source.callbacks) == 0: break return _ani
[docs] def __repr__(self): """ Human readable version of the display list :param self: the animation :type self: Animate :returns: readable version of the display list :rtype: str """ return "Animate2(" + ", ".join([x.type for x in self.displaylist]) + ")"
[docs] def __str__(self): return f"Animate2(len={len(self.displaylist)}"
[docs] def artists(self): """ List of artists that need to be updated :param self: the animation :type self: Animate :returns: list of artists :rtype: list """ return [x.h for x in self.displaylist]
def _draw(self, T): for x in self.displaylist: x.draw(T) # ------------------- plot()
[docs] def set_aspect(self, *args, **kwargs): self.ax.set_aspect(*args, **kwargs)
[docs] def autoscale(self, *args, **kwargs): # self.ax.autoscale(*args, **kwargs) pass
class _Line: def __init__(self, anim, h, xs, ys): # form 3xN matrix, columns are first/last point in homogeneous form p = np.vstack([xs, ys]) self.p = np.vstack([p, np.ones((p.shape[1],))]) self.h = h self.type = "line" self.anim = anim def draw(self, T): p = T @ self.p self.h.set_data(p[0, :], p[1, :])
[docs] def plot(self, x, y, *args, **kwargs): """ Plot a polyline :param x: list of x-coordinates :type x: array_like :param y: list of y-coordinates :type y: array_like Other arguments as accepted by the matplotlib method. All arrays must have the same length. :seealso: :func:`matplotlib.pyplot.plot` """ (h,) = self.ax.plot(x, y, *args, **kwargs) self.displaylist.append(Animate2._Line(self, h, x, y)) return h
# ------------------- quiver() class _Quiver: def __init__(self, anim, h, x, y, u, v): self.type = "quiver" self.anim = anim self.h = h self.type = "arrow" self.anim = anim self.p = np.c_[u - x, v - y].T def draw(self, T): R, t = smb.tr2rt(T) p = R @ self.p # specific to a single Quiver self.h.set_offsets(t) # shift the origin self.h.set_UVC(p[0], p[1])
[docs] def quiver(self, x, y, u, v, *args, **kwargs): """ Plot a quiver :param x: list of base x-coordinates :type x: array_like :param y: list of base y-coordinates :type y: array_like :param u: list of vector x-coordinates :type u: array_like :param v: list of vector y-coordinates :type v: array_like Draws a series of arrows, the bases defined by corresponding elements of (x,y,z) and the vector has components defined by corresponding elements of (u,v,w). Other arguments as accepted by the matplotlib method. :seealso: :func:`matplotlib.pyplot.quiver` """ h = self.ax.quiver(x, y, u, v, *args, **kwargs) self.displaylist.append(Animate2._Quiver(self, h, x, y, u, v))
# ------------------- text() class _Text: def __init__(self, anim, h, x, y): self.type = "text" self.h = h self.p = np.r_[x, y, 1] self.anim = anim def draw(self, T): p = T @ self.p # x2, y2, _ = proj3d.proj_transform( # p[0], p[1], p[2], self.anim.ax.get_proj()) # self.h.set_position((x2, y2)) self.h.set_position((p[0], p[1]))
[docs] def text(self, x, y, *args, **kwargs): """ Plot text :param x: x-coordinate :type x: float :param y: float :type y: float :param z: z-coordinate :type z: float :param kwargs: Other arguments as accepted by the matplotlib method. ``.text(x, y, s)`` display the string ``s`` at coordinate (``x``, ``y``). :seealso: :func:`matplotlib.pyplot.text` """ h = self.ax.text(x, y, *args, **kwargs) self.displaylist.append(Animate2._Text(self, h, x, y))
# ------------------- scatter()
[docs] def scatter(self, x, y, s=0, **kwargs): h = self.plot(x, y, ".", markersize=0, **kwargs) self.displaylist.append(Animate2._Line(self, h, x, y))
# ------------------- wrappers for Axes primitives
[docs] def set_xlim(self, *args, **kwargs): self.ax.set_xlim(*args, **kwargs)
[docs] def set_ylim(self, *args, **kwargs): self.ax.set_ylim(*args, **kwargs)
[docs] def set_xlabel(self, *args, **kwargs): self.ax.set_xlabel(*args, **kwargs)
[docs] def set_ylabel(self, *args, **kwargs): self.ax.set_ylabel(*args, **kwargs)
if __name__ == "__main__": # from spatialmath import UnitQuaternion # from spatialmath.base import tranimate, r2t # J = np.array([[2, -1, 0], [-1, 4, 0], [0, 0, 3]]) # dt = 0.05 # def attitude(): # attitude = UnitQuaternion() # w = 0.2 * np.r_[1, 2, 2].T # for t in np.arange(0, 3, dt): # wd = -np.linalg.inv(J) @ (np.cross(w, J @ w)) # w += wd * dt # attitude.increment(w * dt) # yield attitude.R # plt.figure() # plotvol3(2) # tranimate(attitude()) from spatialmath import base # T = smb.rpy2r(0.3, 0.4, 0.5) # # smb.tranimate(T, wait=True) # s = smb.tranimate(T, movie=True) # with open("zz.html", "w") as f: # print(f"<html>{s}</html>", file=f) T = smb.rot2(2) # smb.tranimate2(T, wait=True) s = smb.tranimate2(T, movie=True) with open("zz.html", "w") as f: print(f"<html>{s}</html>", file=f)