from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.gridspec import GridSpec
from matplotlib.widgets import Button, Slider
from matplotlib.text import Text
import matplotlib.pyplot as plt
import numpy as np
from animatplot import Timeline
from animatplot.blocks.base import Block
[docs]
class Animation:
"""The foundation of all animations.
Parameters
----------
blocks : list of animatplot.animations.Block
A list of blocks to be animated
timeline : Timeline or 1D array, optional
If an array is passed in, it will be converted to a Timeline.
If not given, a timeline will be created using the length of the
first block.
fig : matplotlib figure, optional
The figure that the animation is to occur on
Attributes
----------
animation
a matplotlib animation returned from FuncAnimation
"""
[docs]
def __init__(self, blocks, timeline=None, fig=None):
self.fig = plt.gcf() if fig is None else fig
self.animation = self._animate(blocks, timeline)
def _animate(self, blocks, timeline):
if timeline is None:
self.timeline = Timeline(range(len(blocks[0])))
elif not isinstance(timeline, Timeline):
self.timeline = Timeline(timeline)
else:
self.timeline = timeline
_len_time = len(self.timeline)
for block in blocks:
if len(block) != _len_time:
raise ValueError(
"All blocks must animate for the same amount " "of time"
)
self.blocks = blocks
self._has_slider = False
self._pause = False
self._controls_gridspec_object = None
def update_all(i):
updates = []
for block in self.blocks:
updates.append(block._update(self.timeline.index))
if self._has_slider:
self.slider.set_val(self.timeline.index)
self.timeline._update()
return updates
args = self.fig, update_all
kwargs = dict(frames=self.timeline._len, interval=1000 / self.timeline.fps)
try:
return FuncAnimation(*args, **kwargs)
except KeyError:
return FuncAnimation(*args, **kwargs)
@property
def _controls_gridspec(self):
if self._controls_gridspec_object is None:
# make the bottom of the subplots grid lower to fit the controls in
adjust_plot = {"bottom": 0.03}
plt.subplots_adjust(**adjust_plot)
controls_height = 0.2
fig_gridspecs = self.fig._gridspecs
if len(fig_gridspecs) > 1:
raise ValueError("multiple gridspecs found in figure")
gs = fig_gridspecs[0]
nrows, ncols = gs.get_geometry()
height_ratios = gs.get_height_ratios()
# update parameters with a new row
if height_ratios is None:
# if height_ratios is None, all rows on the original gridspec
# are the same height
height_ratios = [(1.0 - controls_height) / nrows for i in range(nrows)]
else:
height_ratios = [r * (1.0 - controls_height) for r in height_ratios]
height_ratios.append(controls_height)
gs._nrows += 1
gs.set_height_ratios(height_ratios)
# make a sub-grid in the bottom row
self._controls_gridspec_object = gs[-1, :].subgridspec(
1, 3, width_ratios=[0.07, 0.65, 0.28], wspace=0.0, hspace=0.0
)
gs.update()
return self._controls_gridspec_object
[docs]
def toggle(self, ax=None):
"""Creates a play/pause button to start/stop the animation
Parameters
----------
ax : matplotlib.axes.Axes, optional
The matplotlib axes to attach the button to.
"""
if ax is None:
try:
button_subplotspec = self._controls_gridspec[0, 2]
button_gridspec = button_subplotspec.subgridspec(
3,
3,
width_ratios=[0.45, 0.45, 0.1],
height_ratios=[0.05, 0.5, 0.45],
wspace=0.0,
hspace=0.0,
)
self.button_ax = self.fig.add_subplot(button_gridspec[1, 1])
except:
# editing the gridspec did not work for some reason, fall back to
# subplots_adjust
print(
"warning: adding play/pause button to gridspec failed, "
"adding in independent axes. tight_layout() will ignore "
"the button."
)
adjust_plot = {"bottom": 0.2}
left, bottom, width, height = (0.78, 0.03, 0.1, 0.07)
rect = (left, bottom, width, height)
plt.subplots_adjust(**adjust_plot)
self.button_ax = plt.axes(rect)
else:
self.button_ax = ax
self.button = Button(self.button_ax, "Pause")
self.button.label2 = self.button_ax.text(
x=0.5,
y=0.5,
s="Play",
verticalalignment="center",
horizontalalignment="center",
transform=self.button_ax.transAxes,
)
self.button.label2.set_visible(False)
def pause(event):
if self._pause:
self.animation.event_source.start()
self.button.label.set_visible(True)
self.button.label2.set_visible(False)
else:
self.animation.event_source.stop()
self.button.label.set_visible(False)
self.button.label2.set_visible(True)
self.fig.canvas.draw()
self._pause ^= True
self.button.on_clicked(pause)
[docs]
def timeline_slider(self, text="Time", ax=None, valfmt=None, color=None):
"""Creates a timeline slider.
Parameters
----------
text : str, optional
The text to display for the slider. Defaults to 'Time'
ax : matplotlib.axes.Axes, optional
The matplotlib axes to attach the slider to.
valfmt : str, optional
a format specifier used to print the time
Defaults to '%s' for datetime64, timedelta64 and '%1.2f' otherwise.
color :
The color of the slider.
"""
if ax is None:
try:
slider_subplotspec = self._controls_gridspec[0, 1]
slider_gridspec = slider_subplotspec.subgridspec(
3, 1, height_ratios=[0.2, 0.2, 0.6], wspace=0.0, hspace=0.0
)
self.slider_ax = self.fig.add_subplot(slider_gridspec[1, 0])
except:
# editing the gridspec did not work for some reason, fall back to
# subplots_adjust
print(
"warning: adding timeline slider to gridspec failed, "
"adding in independent axes. tight_layout() will ignore "
"the slider."
)
adjust_plot = {"bottom": 0.2}
rect = [0.18, 0.05, 0.5, 0.03]
plt.subplots_adjust(**adjust_plot)
self.slider_ax = plt.axes(rect)
else:
self.slider_ax = ax
if valfmt is None:
if np.issubdtype(self.timeline.t.dtype, np.datetime64) or np.issubdtype(
self.timeline.t.dtype, np.timedelta64
):
valfmt = "%s"
else:
valfmt = "%1.2f"
if self.timeline.log:
valfmt = "$10^{%s}$" % valfmt
if ax is None:
# Try to intelligently decide slider width to avoid overlap
renderer = self.fig.canvas.get_renderer()
# Calculate width of widest time value on plot
def text_width(txt):
t_val_text = Text(text=txt, figure=self.fig)
bbox = t_val_text.get_window_extent(renderer=renderer)
extents = self.fig.transFigure.inverted().transform(bbox)
return extents[1][0] - extents[0][0]
text_val_width = max(
text_width(valfmt % (self.timeline[i]))
for i in range(len(self.timeline))
)
label_width = text_width(text)
# Calculate width of slider
default_button_width = 0.1
width = 0.73 - text_val_width - label_width - default_button_width
adjust_plot = {"bottom": 0.2}
left, bottom, height = (0.18, 0.05, 0.03)
rect = (left, bottom, width, height)
plt.subplots_adjust(**adjust_plot)
self.slider_ax = plt.axes(rect)
else:
self.slider_ax = ax
self.slider = Slider(
self.slider_ax,
label=text,
valmin=0,
valmax=self.timeline._len - 1,
valinit=0,
valfmt=(valfmt + self.timeline.units),
valstep=1,
color=color,
)
self._has_slider = True
def set_time(new_slider_val):
# Update slider value and text on each step
self.timeline.index = int(new_slider_val)
self.slider.valtext.set_text(
self.slider.valfmt % (self.timeline[self.timeline.index])
)
if self._pause:
for block in self.blocks:
block._update(self.timeline.index)
self.fig.canvas.draw()
self.slider.on_changed(set_time)
[docs]
def controls(self, timeline_slider_args={}, toggle_args={}):
"""Creates interactive controls for the animation
Creates both a play/pause button, and a time slider at once
Parameters
----------
timeline_slider_args : Dict, optional
A dictionary of arguments to be passed to timeline_slider()
toggle_args : Dict, optional
A dictionary of argyments to be passed to toggle()
"""
self.timeline_slider(**timeline_slider_args)
self.toggle(**toggle_args)
[docs]
def save_gif(self, filename):
"""Saves the animation to a gif
A convenience function. Provided to let the user avoid dealing
with writers - uses PillowWriter.
Parameters
----------
filename : str
the name of the file to be created without the file extension
"""
self.timeline.index -= 1 # required for proper starting point for save
self.animation.save(
filename + ".gif", writer=PillowWriter(fps=self.timeline.fps)
)
[docs]
def save(self, *args, **kwargs):
"""Saves an animation
A wrapper around :meth:`matplotlib.animation.Animation.save`
"""
self.timeline.index -= 1 # required for proper starting point for save
self.animation.save(*args, **kwargs)
[docs]
def add(self, new):
"""
Updates the animation object by adding additional blocks.
The new blocks can be passed as a list, or as part of a second animaion.
If passed as part of a new animation, the timeline of this new
animation object will replace the old one.
Parameters
----------
new : amp.animation.Animation, or list of amp.block.Block objects
Either blocks to add to animation instance, or another animation
instance whose blocks should be combined with this animation.
"""
if isinstance(new, Animation):
new_blocks = new.blocks
new_timeline = new.timeline
else:
if not isinstance(new, list):
new_blocks = [new]
else:
new_blocks = new
new_timeline = self.timeline
for i, block in enumerate(new_blocks):
if not isinstance(block, Block):
raise TypeError(
f"Block number {i} passed is of type "
f"{type(block)}, not of type "
f"animatplot.blocks.Block (or a subclass)"
)
self.blocks.append(block)
self.animation = self._animate(self.blocks, new_timeline)