axes_grid1 toolkits now supports gridspec

leejjoon · leejjoon · commit 0a2a0f9c704b · 2011-04-22T00:51:02.000+09:00
diff --git a/lib/mpl_toolkits/axes_grid1/axes_divider.py b/lib/mpl_toolkits/axes_grid1/axes_divider.py
@@ -333,6 +333,7 @@ def __call__(self, axes, renderer):
                                          renderer)
 
 
+from matplotlib.gridspec import SubplotSpec, GridSpec
 
 class SubplotDivider(Divider):
     """
@@ -357,28 +358,48 @@ def __init__(self, fig, *args, **kwargs):
         self.figure = fig
 
         if len(args)==1:
-            s = str(args[0])
-            if len(s) != 3:
-                raise ValueError('Argument to subplot must be a 3 digits long')
-            rows, cols, num = map(int, s)
+            if isinstance(args[0], SubplotSpec):
+                self._subplotspec = args[0]
+            else:
+                try:
+                    s = str(int(args[0]))
+                    rows, cols, num = map(int, s)
+                except ValueError:
+                    raise ValueError(
+                         'Single argument to subplot must be a 3-digit integer')
+                self._subplotspec = GridSpec(rows, cols)[num-1]
+                # num - 1 for converting from MATLAB to python indexing
         elif len(args)==3:
             rows, cols, num = args
+            rows = int(rows)
+            cols = int(cols)
+            if isinstance(num, tuple) and len(num) == 2:
+                num = [int(n) for n in num]
+                self._subplotspec = GridSpec(rows, cols)[num[0]-1:num[1]]
+            else:
+                self._subplotspec = GridSpec(rows, cols)[int(num)-1]
+                # num - 1 for converting from MATLAB to python indexing
         else:
-            raise ValueError(  'Illegal argument to subplot')
+            raise ValueError('Illegal argument(s) to subplot: %s' % (args,))
+
 
+        # total = rows*cols
+        # num -= 1    # convert from matlab to python indexing
+        #             # ie num in range(0,total)
+        # if num >= total:
+        #     raise ValueError( 'Subplot number exceeds total subplots')
+        # self._rows = rows
+        # self._cols = cols
+        # self._num = num
 
-        total = rows*cols
-        num -= 1    # convert from matlab to python indexing
-                    # ie num in range(0,total)
-        if num >= total:
-            raise ValueError( 'Subplot number exceeds total subplots')
-        self._rows = rows
-        self._cols = cols
-        self._num = num
+        # self.update_params()
 
+
+        # sets self.fixbox
         self.update_params()
 
         pos = self.figbox.bounds
+
         horizontal = kwargs.pop("horizontal", [])
         vertical = kwargs.pop("vertical", [])
         aspect = kwargs.pop("aspect", None)
@@ -393,40 +414,67 @@ def __init__(self, fig, *args, **kwargs):
 
     def get_position(self):
         "return the bounds of the subplot box"
-        self.update_params()
+
+        self.update_params() # update self.figbox
         return self.figbox.bounds
 
 
+    # def update_params(self):
+    #     'update the subplot position from fig.subplotpars'
+
+    #     rows = self._rows
+    #     cols = self._cols
+    #     num = self._num
+
+    #     pars = self.figure.subplotpars
+    #     left = pars.left
+    #     right = pars.right
+    #     bottom = pars.bottom
+    #     top = pars.top
+    #     wspace = pars.wspace
+    #     hspace = pars.hspace
+    #     totWidth = right-left
+    #     totHeight = top-bottom
+
+    #     figH = totHeight/(rows + hspace*(rows-1))
+    #     sepH = hspace*figH
+
+    #     figW = totWidth/(cols + wspace*(cols-1))
+    #     sepW = wspace*figW
+
+    #     rowNum, colNum =  divmod(num, cols)
+
+    #     figBottom = top - (rowNum+1)*figH - rowNum*sepH
+    #     figLeft = left + colNum*(figW + sepW)
+
+    #     self.figbox = mtransforms.Bbox.from_bounds(figLeft, figBottom,
+    #                                                figW, figH)
+
     def update_params(self):
         'update the subplot position from fig.subplotpars'
 
-        rows = self._rows
-        cols = self._cols
-        num = self._num
+        self.figbox = self.get_subplotspec().get_position(self.figure)
 
-        pars = self.figure.subplotpars
-        left = pars.left
-        right = pars.right
-        bottom = pars.bottom
-        top = pars.top
-        wspace = pars.wspace
-        hspace = pars.hspace
-        totWidth = right-left
-        totHeight = top-bottom
+    def get_geometry(self):
+        'get the subplot geometry, eg 2,2,3'
+        rows, cols, num1, num2 = self.get_subplotspec().get_geometry()
+        return rows, cols, num1+1 # for compatibility
 
-        figH = totHeight/(rows + hspace*(rows-1))
-        sepH = hspace*figH
-
-        figW = totWidth/(cols + wspace*(cols-1))
-        sepW = wspace*figW
+    # COVERAGE NOTE: Never used internally or from examples
+    def change_geometry(self, numrows, numcols, num):
+        'change subplot geometry, eg. from 1,1,1 to 2,2,3'
+        self._subplotspec = GridSpec(numrows, numcols)[num-1]
+        self.update_params()
+        self.set_position(self.figbox)
 
-        rowNum, colNum =  divmod(num, cols)
+    def get_subplotspec(self):
+        'get the SubplotSpec instance'
+        return self._subplotspec
 
-        figBottom = top - (rowNum+1)*figH - rowNum*sepH
-        figLeft = left + colNum*(figW + sepW)
+    def set_subplotspec(self, subplotspec):
+        'set the SubplotSpec instance'
+        self._subplotspec = subplotspec
 
-        self.figbox = mtransforms.Bbox.from_bounds(figLeft, figBottom,
-                                                   figW, figH)
 
 
 class AxesDivider(Divider):
diff --git a/lib/mpl_toolkits/axes_grid1/axes_grid.py b/lib/mpl_toolkits/axes_grid1/axes_grid.py
@@ -9,6 +9,8 @@
 import matplotlib.lines as mlines
 import matplotlib.ticker as ticker
 
+from matplotlib.gridspec import SubplotSpec, GridSpec
+
 from axes_divider import Size, SubplotDivider, LocatableAxes, Divider
 
 #import numpy as np
@@ -235,6 +237,9 @@ def __init__(self, fig,
         if cbook.is_string_like(rect) or cbook.is_numlike(rect):
             self._divider = SubplotDivider(fig, rect, horizontal=h, vertical=v,
                                            aspect=False)
+        elif isinstance(rect, SubplotSpec):
+            self._divider = SubplotDivider(fig, rect, horizontal=h, vertical=v,
+                                           aspect=False)
         elif len(rect) == 3:
             kw = dict(horizontal=h, vertical=v, aspect=False)
             self._divider = SubplotDivider(fig, *rect, **kw)
@@ -413,6 +418,9 @@ def set_label_mode(self, mode):
             ax = self.axes_llc
             _tick_only(ax, bottom_on=False, left_on=False)
 
+    def get_divider(self):
+        return self._divider
+        
     def set_axes_locator(self, locator):
         self._divider.set_locator(locator)
 
@@ -543,6 +551,9 @@ def __init__(self, fig,
         if cbook.is_string_like(rect) or cbook.is_numlike(rect):
             self._divider = SubplotDivider(fig, rect, horizontal=h, vertical=v,
                                            aspect=aspect)
+        elif isinstance(rect, SubplotSpec):
+            self._divider = SubplotDivider(fig, rect, horizontal=h, vertical=v,
+                                           aspect=aspect)
         elif len(rect) == 3:
             kw = dict(horizontal=h, vertical=v, aspect=aspect)
             self._divider = SubplotDivider(fig, *rect, **kw)
