From cc9feb50b6f1a64b698f4e3d0673017adbe83698 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Mon, 31 Oct 2022 10:01:59 +0100 Subject: [PATCH] Shorten and explain more calculations in axes_divider. Make get_horizontal_sizes, get_vertical_sizes return arrays (which make subsequent computations easier). Move Divider._calc_k, Divider._calc_offsets closer to their use site (i.e. Divider._locate). Inline and shorten offsetbox._determine_karray, offsetbox._calc_offsets into offsetbox._locate which should make the logic easier to follow. --- lib/mpl_toolkits/axes_grid1/axes_divider.py | 98 ++++++++------------- 1 file changed, 37 insertions(+), 61 deletions(-) diff --git a/lib/mpl_toolkits/axes_grid1/axes_divider.py b/lib/mpl_toolkits/axes_grid1/axes_divider.py index 2d092674623d..aa3dd5415f56 100644 --- a/lib/mpl_toolkits/axes_grid1/axes_divider.py +++ b/lib/mpl_toolkits/axes_grid1/axes_divider.py @@ -55,32 +55,10 @@ def __init__(self, fig, pos, horizontal, vertical, self._locator = None def get_horizontal_sizes(self, renderer): - return [s.get_size(renderer) for s in self.get_horizontal()] + return np.array([s.get_size(renderer) for s in self.get_horizontal()]) def get_vertical_sizes(self, renderer): - return [s.get_size(renderer) for s in self.get_vertical()] - - @staticmethod - def _calc_k(l, total_size): - - rs_sum, as_sum = 0., 0. - - for _rs, _as in l: - rs_sum += _rs - as_sum += _as - - if rs_sum != 0.: - k = (total_size - as_sum) / rs_sum - return k - else: - return 0. - - @staticmethod - def _calc_offsets(l, k): - offsets = [0.] - for _rs, _as in l: - offsets.append(offsets[-1] + _rs*k + _as) - return offsets + return np.array([s.get_size(renderer) for s in self.get_vertical()]) def set_position(self, pos): """ @@ -171,6 +149,19 @@ def get_position_runtime(self, ax, renderer): else: return self._locator(ax, renderer).bounds + @staticmethod + def _calc_k(sizes, total): + # sizes is a (n, 2) array of (rel_size, abs_size); this method finds + # the k factor such that sum(rel_size * k + abs_size) == total. + rel_sum, abs_sum = sizes.sum(0) + return (total - abs_sum) / rel_sum if rel_sum else 0 + + @staticmethod + def _calc_offsets(sizes, k): + # Apply k factors to (n, 2) sizes array of (rel_size, abs_size); return + # the resulting cumulative offset positions. + return np.cumsum([0, *(sizes @ [k, 1])]) + def locate(self, nx, ny, nx1=None, ny1=None, axes=None, renderer=None): """ Parameters @@ -542,48 +533,33 @@ def get_subplotspec(self): # Helper for HBoxDivider/VBoxDivider. # The variable names are written for a horizontal layout, but the calculations -# work identically for vertical layouts (and likewise for the helpers below). -def _determine_karray(summed_widths, equal_heights, total_width, max_height): +# work identically for vertical layouts. +def _locate(x, y, w, h, summed_widths, equal_heights, fig_w, fig_h, anchor): + + total_width = fig_w * w + max_height = fig_h * h + + # Determine the k factors. n = len(equal_heights) - eq_rs, eq_as = np.asarray(equal_heights).T - sm_rs, sm_as = np.asarray(summed_widths).T - A = np.zeros((n + 1, n + 1)) - B = np.zeros(n + 1) - np.fill_diagonal(A[:n, :n], eq_rs) + eq_rels, eq_abss = equal_heights.T + sm_rels, sm_abss = summed_widths.T + A = np.diag([*eq_rels, 0]) A[:n, -1] = -1 - A[-1, :-1] = sm_rs - B[:n] = -eq_as - B[-1] = total_width - 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_width: fixed total width - # (foo_r_i * k_i + foo_a_i will end up being the size of foo.) - karray_and_height = np.linalg.solve(A, B) - karray = karray_and_height[:-1] - height = karray_and_height[-1] + A[-1, :-1] = sm_rels + B = [*(-eq_abss), total_width - sm_abss.sum()] + # A @ K = B: This finds factors {k_0, ..., k_{N-1}, H} so that + # eq_rel_i * k_i + eq_abs_i = H for all i: all axes have the same height + # sum(sm_rel_i * k_i + sm_abs_i) = total_width: fixed total width + # (foo_rel_i * k_i + foo_abs_i will end up being the size of foo.) + *karray, height = np.linalg.solve(A, B) if height > max_height: # Additionally, upper-bound the height. - karray = (max_height - eq_as) / eq_rs - return karray - - -# Helper for HBoxDivider/VBoxDivider (see above re: variable naming). -def _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 - - -# Helper for HBoxDivider/VBoxDivider (see above re: variable naming). -def _locate(x, y, w, h, summed_widths, equal_heights, fig_w, fig_h, anchor): - karray = _determine_karray( - summed_widths, equal_heights, - total_width=fig_w * w, max_height=fig_h * h) - ox = _calc_offsets(summed_widths, karray) + karray = (max_height - eq_abss) / eq_rels + # Compute the offsets corresponding to these factors. + ox = np.cumsum([0, *(sm_rels * karray + sm_abss)]) ww = (ox[-1] - ox[0]) / fig_w - h0_r, h0_a = equal_heights[0] - hh = (karray[0]*h0_r + h0_a) / fig_h + h0_rel, h0_abs = equal_heights[0] + hh = (karray[0]*h0_rel + h0_abs) / fig_h pb = mtransforms.Bbox.from_bounds(x, y, w, h) pb1 = mtransforms.Bbox.from_bounds(x, y, ww, hh) x0, y0 = pb1.anchored(anchor, pb).p0