Add axes for controls by modifying gridspec of figure

Makes control placement compatible with tight_layout() method or
constrained_layout option. Falls back on previous absolute placement
method if the gridspec version fails for any reason.

Author: johnomotani

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,40 @@ 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
+
+            # get gridspec to adjust
+            gs = (plt.gca().get_gridspec().get_topmost_subplotspec()
+                    .get_gridspec())
+            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 +103,20 @@ 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
+                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 +158,19 @@ 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
+                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