diff --git a/animatplot/animation.py b/animatplot/animation.py index ce19522..0be0af0 100644 --- a/animatplot/animation.py +++ b/animatplot/animation.py @@ -1,4 +1,5 @@ from matplotlib.animation import FuncAnimation, PillowWriter +from matplotlib.gridspec import GridSpec from matplotlib.widgets import Button, Slider import matplotlib.pyplot as plt import numpy as np @@ -42,6 +43,7 @@ def __init__(self, blocks, timeline=None, fig=None): self.fig = plt.gcf() if fig is None else fig self._has_slider = False self._pause = False + self._controls_gridspec_object = None def animate(i): updates = [] @@ -58,6 +60,41 @@ def animate(i): interval=1000/self.timeline.fps ) + @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.-controls_height)/nrows + for i in range(nrows)] + else: + height_ratios = [r*(1.-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=[.07, .65, .28], wspace=0., hspace=0.) + gs.update() + + return self._controls_gridspec_object + def toggle(self, ax=None): """Creates a play/pause button to start/stop the animation @@ -67,11 +104,23 @@ def toggle(self, ax=None): The matplotlib axes to attach the button to. """ if ax is None: - adjust_plot = {'bottom': .2} - rect = [.78, .03, .1, .07] - - plt.subplots_adjust(**adjust_plot) - self.button_ax = plt.axes(rect) + 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=[.05, .5, .45], wspace=0., hspace=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': .2} + rect = [.78, .03, .1, .07] + + plt.subplots_adjust(**adjust_plot) + self.button_ax = plt.axes(rect) else: self.button_ax = ax @@ -113,11 +162,22 @@ def timeline_slider(self, text='Time', ax=None, valfmt=None, color=None): The color of the slider. """ if ax is None: - adjust_plot = {'bottom': .2} - rect = [.18, .05, .5, .03] - - plt.subplots_adjust(**adjust_plot) - self.slider_ax = plt.axes(rect) + try: + slider_subplotspec = self._controls_gridspec[0, 1] + slider_gridspec = slider_subplotspec.subgridspec( + 3, 1, height_ratios=[.2, .2, .6], wspace=0., + hspace=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': .2} + rect = [.18, .05, .5, .03] + plt.subplots_adjust(**adjust_plot) + self.slider_ax = plt.axes(rect) else: self.slider_ax = ax