Merge pull request #6 from boutproject/tight_layout

tight_layout() compatibility

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
@@ -48,6 +49,7 @@ def _animate(self, blocks, timeline):
         self.blocks = blocks
         self._has_slider = False
         self._pause = False
+        self._controls_gridspec_object = None
 
         def update_all(i):
             updates = []
@@ -61,6 +63,41 @@ def update_all(i):
         return FuncAnimation(self.fig, update_all, frames=self.timeline._len,
                              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
 
@@ -70,11 +107,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
 
@@ -116,11 +165,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