Don't make VBoxDivider inherit from HBoxDivider.

anntzer · anntzer · commit 1b0b8d41cf51 · 2021-04-25T12:16:52.000+02:00
They share some common calculation helpers, but a VBoxDivider is not a
specialization of an HBoxDivider, so just move the helpers out to free
functions (they're staticmethods or nearly so anyways).

Also rename "equivalent" to "equal" and "appended" to "summed", which I
think actually express the constraints better, and are shorter too.

Also document their usage, and PEP8-ify a couple of variables.
diff --git a/lib/mpl_toolkits/axes_grid1/axes_divider.py b/lib/mpl_toolkits/axes_grid1/axes_divider.py
@@ -581,39 +581,65 @@ def get_subplotspec(self):
             return None
 
 
-class HBoxDivider(SubplotDivider):
+def _hvboxdiv_determine_karray(
+        equal_sizes, summed_sizes, max_equal_size, total_summed_size):
+    n = len(equal_sizes)
+    eq_rs, eq_as = np.asarray(equal_sizes).T
+    sm_rs, sm_as = np.asarray(summed_sizes).T
+    A = np.zeros((n + 1, n + 1))
+    B = np.zeros(n + 1)
+    np.fill_diagonal(A[:n, :n], eq_rs)
+    A[:n, -1] = -1
+    A[-1, :-1] = sm_rs
+    B[:n] = -eq_as
+    B[-1] = total_summed_size - sum(sm_as)
+    # A @ K = B: This solves for {k_0, ..., k_{N-1}, H} so that
+    #   eq_r_i * k_i + eq_a_i = H for all i: all axes have the same height
+    #   sum(sm_r_i * k_i + sm_a_i) = total_summed_size: fixed total width
+    # (foo_r_i * k_i + foo_a_i will end up being the size of foo.)
+    karray_H = np.linalg.solve(A, B)
+    karray = karray_H[:-1]
+    H = karray_H[-1]
+    if H > max_equal_size:  # Additionally, upper-bound the height.
+        karray = (max_equal_size - eq_as) / eq_rs
+    return karray
+
+
+def _hvboxdiv_calc_offsets(summed_sizes, karray):
+    offsets = [0.]
+    for (r, a), k in zip(summed_sizes, karray):
+        offsets.append(offsets[-1] + r*k + a)
+    return offsets
+
+
+# The variable names are written for a horizontal layout, but the calculations
+# work identically for vertical layouts.
+def _hvboxdiv_locate(x, y, w, h, equal_ys, summed_xs, fig_w, fig_h, anchor):
+    karray = _hvboxdiv_determine_karray(
+        equal_ys, summed_xs,
+        max_equal_size=fig_h * h, total_summed_size=fig_w * w)
+    ox = _hvboxdiv_calc_offsets(summed_xs, karray)
+
+    ww = (ox[-1] - ox[0]) / fig_w
+    ref_h = equal_ys[0]
+    hh = (karray[0]*ref_h[0] + ref_h[1]) / fig_h
+    pb = mtransforms.Bbox.from_bounds(x, y, w, h)
+    pb1 = mtransforms.Bbox.from_bounds(x, y, ww, hh)
+    pb1_anchored = pb1.anchored(anchor, pb)
+    x0, y0 = pb1_anchored.x0, pb1_anchored.y0
+
+    return x0, y0, ox, hh
 
-    @staticmethod
-    def _determine_karray(equivalent_sizes, appended_sizes,
-                          max_equivalent_size,
-                          total_appended_size):
-        n = len(equivalent_sizes)
-        eq_rs, eq_as = np.asarray(equivalent_sizes).T
-        ap_rs, ap_as = np.asarray(appended_sizes).T
-        A = np.zeros((n + 1, n + 1))
-        B = np.zeros(n + 1)
-        np.fill_diagonal(A[:n, :n], eq_rs)
-        A[:n, -1] = -1
-        A[-1, :-1] = ap_rs
-        B[:n] = -eq_as
-        B[-1] = total_appended_size - sum(ap_as)
-        # A @ K = B: This solves for {k_0, ..., k_{N-1}, H} so that
-        #   eq_r_i * k_i + eq_a_i = H for all i: all axes have the same height
-        #   sum(ap_r_i * k_i + ap_a_i) = total_appended_size: fixed total width
-        # (foo_r_i * k_i + foo_a_i will end up being the size of foo.)
-        karray_H = np.linalg.solve(A, B)
-        karray = karray_H[:-1]
-        H = karray_H[-1]
-        if H > max_equivalent_size:  # Additionally, upper-bound the height.
-            karray = (max_equivalent_size - eq_as) / eq_rs
-        return karray
 
-    @staticmethod
-    def _calc_offsets(appended_sizes, karray):
-        offsets = [0.]
-        for (r, a), k in zip(appended_sizes, karray):
-            offsets.append(offsets[-1] + r*k + a)
-        return offsets
+class HBoxDivider(SubplotDivider):
+    """
+    A `SubplotDivider` for laying out axes horizontally, while ensuring that
+    they have equal heights.
+
+    Examples
+    --------
+    .. plot:: gallery/axes_grid1/demo_axes_hbox_divider.py
+    """
 
     def new_locator(self, nx, nx1=None):
         """
@@ -629,52 +655,26 @@ def new_locator(self, nx, nx1=None):
         """
         return AxesLocator(self, nx, 0, nx1, None)
 
-    def _locate(self, x, y, w, h,
-                y_equivalent_sizes, x_appended_sizes,
-                figW, figH):
-        equivalent_sizes = y_equivalent_sizes
-        appended_sizes = x_appended_sizes
-
-        max_equivalent_size = figH * h
-        total_appended_size = figW * w
-        karray = self._determine_karray(equivalent_sizes, appended_sizes,
-                                        max_equivalent_size,
-                                        total_appended_size)
-
-        ox = self._calc_offsets(appended_sizes, karray)
-
-        ww = (ox[-1] - ox[0]) / figW
-        ref_h = equivalent_sizes[0]
-        hh = (karray[0]*ref_h[0] + ref_h[1]) / figH
-        pb = mtransforms.Bbox.from_bounds(x, y, w, h)
-        pb1 = mtransforms.Bbox.from_bounds(x, y, ww, hh)
-        pb1_anchored = pb1.anchored(self.get_anchor(), pb)
-        x0, y0 = pb1_anchored.x0, pb1_anchored.y0
-
-        return x0, y0, ox, hh
-
     def locate(self, nx, ny, nx1=None, ny1=None, axes=None, renderer=None):
         # docstring inherited
-        figW, figH = self._fig.get_size_inches()
+        fig_w, fig_h = self._fig.get_size_inches()
         x, y, w, h = self.get_position_runtime(axes, renderer)
-
-        y_equivalent_sizes = self.get_vertical_sizes(renderer)
-        x_appended_sizes = self.get_horizontal_sizes(renderer)
-        x0, y0, ox, hh = self._locate(x, y, w, h,
-                                      y_equivalent_sizes, x_appended_sizes,
-                                      figW, figH)
+        y_equal_sizes = self.get_vertical_sizes(renderer)
+        x_summed_sizes = self.get_horizontal_sizes(renderer)
+        x0, y0, ox, hh = _hvboxdiv_locate(
+            x, y, w, h, y_equal_sizes, x_summed_sizes, fig_w, fig_h,
+            self.get_anchor())
         if nx1 is None:
             nx1 = nx + 1
-
-        x1, w1 = x0 + ox[nx] / figW, (ox[nx1] - ox[nx]) / figW
+        x1, w1 = x0 + ox[nx] / fig_w, (ox[nx1] - ox[nx]) / fig_w
         y1, h1 = y0, hh
-
         return mtransforms.Bbox.from_bounds(x1, y1, w1, h1)
 
 
-class VBoxDivider(HBoxDivider):
+class VBoxDivider(SubplotDivider):
     """
-    The Divider class whose rectangle area is specified as a subplot geometry.
+    A `SubplotDivider` for laying out axes vertically, while ensuring that they
+    have equal widths.
     """
 
     def new_locator(self, ny, ny1=None):
@@ -693,20 +693,17 @@ def new_locator(self, ny, ny1=None):
 
     def locate(self, nx, ny, nx1=None, ny1=None, axes=None, renderer=None):
         # docstring inherited
-        figW, figH = self._fig.get_size_inches()
+        fig_w, fig_h = self._fig.get_size_inches()
         x, y, w, h = self.get_position_runtime(axes, renderer)
-
-        x_equivalent_sizes = self.get_horizontal_sizes(renderer)
-        y_appended_sizes = self.get_vertical_sizes(renderer)
-        y0, x0, oy, ww = self._locate(y, x, h, w,
-                                      x_equivalent_sizes, y_appended_sizes,
-                                      figH, figW)
+        x_equal_sizes = self.get_horizontal_sizes(renderer)
+        y_summed_sizes = self.get_vertical_sizes(renderer)
+        y0, x0, oy, ww = _hvboxdiv_locate(
+            y, x, h, w, x_equal_sizes, y_summed_sizes, fig_h, fig_w,
+            self.get_anchor())
         if ny1 is None:
             ny1 = ny + 1
-
         x1, w1 = x0, ww
-        y1, h1 = y0 + oy[ny] / figH, (oy[ny1] - oy[ny]) / figH
-
+        y1, h1 = y0 + oy[ny] / fig_h, (oy[ny1] - oy[ny]) / fig_h
         return mtransforms.Bbox.from_bounds(x1, y1, w1, h1)
 
 
