47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
"""Easy matplotlib animation. From https://github.com/jwkvam/celluloid."""
|
|
|
|
from collections import defaultdict
|
|
from typing import Dict, List
|
|
|
|
from matplotlib.animation import ArtistAnimation
|
|
from matplotlib.artist import Artist
|
|
from matplotlib.figure import Figure
|
|
|
|
__version__ = "0.2.0"
|
|
|
|
|
|
class Camera:
|
|
"""Make animations easier."""
|
|
def __init__(self, figure: Figure) -> None:
|
|
"""Create camera from matplotlib figure."""
|
|
self._figure = figure
|
|
# need to keep track off artists for each axis
|
|
self._offsets: Dict[str, Dict[int, int]] = {
|
|
k: defaultdict(int)
|
|
for k in
|
|
["collections", "patches", "lines", "texts", "artists", "images"]
|
|
}
|
|
self._photos: List[List[Artist]] = []
|
|
|
|
def snap(self) -> List[Artist]:
|
|
"""Capture current state of the figure."""
|
|
frame_artists: List[Artist] = []
|
|
for i, axis in enumerate(self._figure.axes):
|
|
if axis.legend_ is not None:
|
|
axis.add_artist(axis.legend_)
|
|
for name in self._offsets:
|
|
new_artists = getattr(axis, name)[self._offsets[name][i]:]
|
|
frame_artists += new_artists
|
|
self._offsets[name][i] += len(new_artists)
|
|
self._photos.append(frame_artists)
|
|
return frame_artists
|
|
|
|
def animate(self, *args, **kwargs) -> ArtistAnimation:
|
|
"""Animate the snapshots taken.
|
|
Uses matplotlib.animation.ArtistAnimation
|
|
Returns
|
|
-------
|
|
ArtistAnimation
|
|
"""
|
|
return ArtistAnimation(self._figure, self._photos, *args, **kwargs)
|