Source code for proplot.gridspec

#!/usr/bin/env python3
"""
The new gridspec classes.
"""
import numpy as np
import matplotlib.axes as maxes
import matplotlib.gridspec as mgridspec
from .utils import units
from .config import rc
from .internals import ic  # noqa: F401
from .internals import _not_none

__all__ = ['GridSpec', 'SubplotSpec']


def _default_space(key, share=0, pad=None):
    """
    Return suitable default spacing given a shared axes setting.
    """
    # Pull out sizes
    outerpad = rc['subplots.pad']  # TODO: rename these to outerpad, innerpad
    innerpad = rc['subplots.axpad']
    xtick = rc['xtick.major.size']
    ytick = rc['ytick.major.size']
    xtickpad = rc['xtick.major.pad']
    ytickpad = rc['ytick.major.pad']
    xticklabel = rc._scale_font(rc['xtick.labelsize'])
    yticklabel = 3 * rc._scale_font(rc['ytick.labelsize'])
    label = rc._scale_font(rc['axes.labelsize'])
    title = rc._scale_font(rc['axes.titlesize'])
    titlepad = rc['axes.titlepad']

    # Get suitable size for various spaces
    if key == 'left':
        space = units(_not_none(pad, outerpad)) + (
            ytick + yticklabel + ytickpad + label
        ) / 72

    elif key == 'right':
        space = units(_not_none(pad, outerpad))

    elif key == 'bottom':
        space = units(_not_none(pad, outerpad)) + (
            xtick + xticklabel + xtickpad + label
        ) / 72

    elif key == 'top':
        space = units(_not_none(pad, outerpad)) + (title + titlepad) / 72

    elif key == 'wspace':
        space = units(_not_none(pad, innerpad)) + ytick / 72
        if share < 3:
            space += (yticklabel + ytickpad) / 72
        if share < 1:
            space += label / 72

    elif key == 'hspace':
        space = units(_not_none(pad, innerpad)) + (title + titlepad + xtick) / 72
        if share < 3:
            space += (xticklabel + xtickpad) / 72
        if share < 0:
            space += label / 72

    else:
        raise KeyError(f'Invalid space key {key!r}.')

    return space


def _calc_geometry(**kwargs):
    """
    Save arguments passed to `~proplot.ui.subplots`, calculates
    gridspec settings and figure size necessary for requested geometry, and
    returns keyword args necessary to reconstruct and modify this
    configuration. Note that `wspace`, `hspace`, `left`, `right`, `top`, and
    `bottom` always have fixed physical units, then we scale figure width,
    figure height, and width and height ratios to accommodate spaces.
    """
    # NOTE: In the future this will be replaced with a GeometryConfigurator
    # class that acts as the interface between the figure and the gridspec.
    # Dimensions and geometry
    nrows, ncols = kwargs['nrows'], kwargs['ncols']
    aspect, xref, yref = kwargs['aspect'], kwargs['xref'], kwargs['yref']
    width, height = kwargs['width'], kwargs['height']
    axwidth, axheight = kwargs['axwidth'], kwargs['axheight']

    # Gridspec settings
    wspace, hspace = kwargs['wspace'], kwargs['hspace']
    wratios, hratios = kwargs['wratios'], kwargs['hratios']
    left, bottom = kwargs['left'], kwargs['bottom']
    right, top = kwargs['right'], kwargs['top']

    # Panel string toggles, lists containing empty strings '' (indicating a
    # main axes), or one of 'l', 'r', 'b', 't' (indicating axes panels) or
    # 'f' (indicating figure panels)
    wpanels, hpanels = kwargs['wpanels'], kwargs['hpanels']

    # Checks, important now that we modify gridspec geometry
    if len(hratios) != nrows:
        raise ValueError(
            f'Expected {nrows} width ratios for {nrows} rows, '
            f'got {len(hratios)}.'
        )
    if len(wratios) != ncols:
        raise ValueError(
            f'Expected {ncols} width ratios for {ncols} columns, '
            f'got {len(wratios)}.'
        )
    if len(hspace) != nrows - 1:
        raise ValueError(
            f'Expected {nrows - 1} hspaces for {nrows} rows, '
            f'got {len(hspace)}.'
        )
    if len(wspace) != ncols - 1:
        raise ValueError(
            f'Expected {ncols - 1} wspaces for {ncols} columns, '
            f'got {len(wspace)}.'
        )
    if len(hpanels) != nrows:
        raise ValueError(
            f'Expected {nrows} hpanel toggles for {nrows} rows, '
            f'got {len(hpanels)}.'
        )
    if len(wpanels) != ncols:
        raise ValueError(
            f'Expected {ncols} wpanel toggles for {ncols} columns, '
            f'got {len(wpanels)}.'
        )

    # Get indices corresponding to main axes or main axes space slots
    idxs_ratios, idxs_space = [], []
    for panels in (hpanels, wpanels):
        # Ratio indices
        mask = np.array([bool(s) for s in panels])
        ratio_idxs, = np.where(~mask)
        idxs_ratios.append(ratio_idxs)
        # Space indices
        space_idxs = []
        for idx in ratio_idxs[:-1]:  # exclude last axes slot
            offset = 1
            while panels[idx + offset] not in 'rbf':  # main space next to this
                offset += 1
            space_idxs.append(idx + offset - 1)
        idxs_space.append(space_idxs)

    # Separate the panel and axes ratios
    hratios_main = [hratios[idx] for idx in idxs_ratios[0]]
    wratios_main = [wratios[idx] for idx in idxs_ratios[1]]
    hratios_panels = [
        ratio for idx, ratio in enumerate(hratios)
        if idx not in idxs_ratios[0]
    ]
    wratios_panels = [
        ratio for idx, ratio in enumerate(wratios)
        if idx not in idxs_ratios[1]
    ]
    hspace_main = [hspace[idx] for idx in idxs_space[0]]
    wspace_main = [wspace[idx] for idx in idxs_space[1]]

    # Reduced geometry
    nrows_main = len(hratios_main)
    ncols_main = len(wratios_main)

    # Get reference properties, account for panel slots in space and ratios
    # TODO: Shouldn't panel space be included in these calculations?
    (x1, x2), (y1, y2) = xref, yref
    dx, dy = x2 - x1 + 1, y2 - y1 + 1
    rwspace = sum(wspace_main[x1:x2])
    rhspace = sum(hspace_main[y1:y2])
    rwratio = (
        ncols_main * sum(wratios_main[x1:x2 + 1]) / (dx * sum(wratios_main))
    )
    rhratio = (
        nrows_main * sum(hratios_main[y1:y2 + 1]) / (dy * sum(hratios_main))
    )
    if rwratio == 0 or rhratio == 0:
        raise RuntimeError(
            f'Something went wrong, got wratio={rwratio!r} '
            f'and hratio={rhratio!r} for reference axes.'
        )
    if np.iterable(aspect):
        aspect = aspect[0] / aspect[1]

    # Determine figure and axes dims from input in width or height dimenion.
    # For e.g. common use case [[1,1,2,2],[0,3,3,0]], make sure we still scale
    # the reference axes like square even though takes two columns of gridspec!
    auto_width = (width is None and height is not None)
    auto_height = (height is None and width is not None)
    if width is None and height is None:  # get stuff directly from axes
        if axwidth is None and axheight is None:
            axwidth = units(rc['subplots.axwidth'])
        if axheight is not None:
            auto_width = True
            axheight_all = (nrows_main * (axheight - rhspace)) / (dy * rhratio)
            height = axheight_all + top + bottom + \
                sum(hspace) + sum(hratios_panels)
        if axwidth is not None:
            auto_height = True
            axwidth_all = (ncols_main * (axwidth - rwspace)) / (dx * rwratio)
            width = axwidth_all + left + right + \
                sum(wspace) + sum(wratios_panels)
        if axwidth is not None and axheight is not None:
            auto_width = auto_height = False
    else:
        if height is not None:
            axheight_all = height - top - bottom - \
                sum(hspace) - sum(hratios_panels)
            axheight = (axheight_all * dy * rhratio) / nrows_main + rhspace
        if width is not None:
            axwidth_all = width - left - right - \
                sum(wspace) - sum(wratios_panels)
            axwidth = (axwidth_all * dx * rwratio) / ncols_main + rwspace

    # Automatically figure dim that was not specified above
    if auto_height:
        axheight = axwidth / aspect
        axheight_all = (nrows_main * (axheight - rhspace)) / (dy * rhratio)
        height = axheight_all + top + bottom + \
            sum(hspace) + sum(hratios_panels)
    elif auto_width:
        axwidth = axheight * aspect
        axwidth_all = (ncols_main * (axwidth - rwspace)) / (dx * rwratio)
        width = axwidth_all + left + right + sum(wspace) + sum(wratios_panels)
    if axwidth_all < 0:
        raise ValueError(
            f'Not enough room for axes (would have width {axwidth_all}). '
            'Try using tight=False, increasing figure width, or decreasing '
            "'left', 'right', or 'wspace' spaces."
        )
    if axheight_all < 0:
        raise ValueError(
            f'Not enough room for axes (would have height {axheight_all}). '
            'Try using tight=False, increasing figure height, or decreasing '
            "'top', 'bottom', or 'hspace' spaces."
        )

    # Reconstruct the ratios array with physical units for subplot slots
    # The panel slots are unchanged because panels have fixed widths
    wratios_main = axwidth_all * np.array(wratios_main) / sum(wratios_main)
    hratios_main = axheight_all * np.array(hratios_main) / sum(hratios_main)
    for idx, ratio in zip(idxs_ratios[0], hratios_main):
        hratios[idx] = ratio
    for idx, ratio in zip(idxs_ratios[1], wratios_main):
        wratios[idx] = ratio

    # Convert margins to figure-relative coordinates
    left = left / width
    bottom = bottom / height
    right = 1 - right / width
    top = 1 - top / height

    # Return gridspec keyword args
    gridspec_kw = {
        'ncols': ncols, 'nrows': nrows,
        'wspace': wspace, 'hspace': hspace,
        'width_ratios': wratios, 'height_ratios': hratios,
        'left': left, 'bottom': bottom, 'right': right, 'top': top,
    }

    return (width, height), gridspec_kw, kwargs


[docs]class SubplotSpec(mgridspec.SubplotSpec): """ Matplotlib `~matplotlib.gridspec.SubplotSpec` subclass that adds some helpful methods. """ def __repr__(self): nrows, ncols, row1, row2, col1, col2 = self.get_rows_columns() return f'SubplotSpec({nrows}, {ncols}; {row1}:{row2}, {col1}:{col2})'
[docs] def get_active_geometry(self): """ Return the number of rows, number of columns, and 1d subplot location indices, ignoring rows and columns allocated for spaces. """ nrows, ncols, row1, row2, col1, col2 = self.get_active_rows_columns() num1 = row1 * ncols + col1 num2 = row2 * ncols + col2 return nrows, ncols, num1, num2
[docs] def get_active_rows_columns(self): """ Return the number of rows, number of columns, first subplot row, last subplot row, first subplot column, and last subplot column, ignoring rows and columns allocated for spaces. """ gridspec = self.get_gridspec() nrows, ncols = gridspec.get_geometry() row1, col1 = divmod(self.num1, ncols) if self.num2 is not None: row2, col2 = divmod(self.num2, ncols) else: row2 = row1 col2 = col1 return ( nrows // 2, ncols // 2, row1 // 2, row2 // 2, col1 // 2, col2 // 2 )
[docs] def get_geometry(self): """ Return the number of rows, number of columns, and 1d subplot location indices. """ gridspec = self.get_gridspec() rows, cols = gridspec.get_geometry() return rows, cols, self.num1, self.num2
[docs] def get_rows_columns(self): """ Return the number of rows, number of columns, first subplot row, last subplot row, first subplot column, and last subplot column. """ gridspec = self.get_gridspec() nrows, ncols = gridspec.get_geometry() row_start, col_start = divmod(self.num1, ncols) row_stop, col_stop = divmod(self.num2, ncols) return nrows, ncols, row_start, row_stop, col_start, col_stop
@classmethod def _from_subplotspec(cls, subplotspec): """ Translate a matplotlib subplotspec to proplot subplotspec. """ return cls(subplotspec._gridspec, subplotspec.num1, subplotspec.num2)
[docs]class GridSpec(mgridspec.GridSpec): """ Matplotlib `~matplotlib.gridspec.GridSpec` subclass that allows for grids with variable spacing between successive rows and columns of axes. This is done by drawing ``nrows * 2 + 1`` and ``ncols * 2 + 1`` `~matplotlib.gridspec.GridSpec` rows and columns, setting `wspace` and `hspace` to ``0``, and masking out every other row and column of the `~matplotlib.gridspec.GridSpec`, so they act as "spaces". These "spaces" are then allowed to vary in width using the builtin `width_ratios` and `height_ratios` properties. Note ---- In a future version, this class will natively support variable spacing between successive rows and columns without the obfuscation. It will also support specifying spaces in physical units via `~proplot.utils.units`. """ def __repr__(self): # do not show width and height ratios nrows, ncols = self.get_geometry() return f'GridSpec({nrows}, {ncols})' def __init__(self, figure, nrows=1, ncols=1, **kwargs): """ Parameters ---------- figure : `~proplot.figure.Figure` The figure instance filled by this gridspec. Unlike `~matplotlib.gridspec.GridSpec`, this argument is required. nrows, ncols : int, optional The number of rows and columns on the subplot grid. hspace, wspace : float or list of float The vertical and horizontal spacing between rows and columns of subplots, respectively. In `~proplot.ui.subplots`, ``wspace`` and ``hspace`` are in physical units. When calling `GridSpec` directly, values are scaled relative to the average subplot height or width. If float, the spacing is identical between all rows and columns. If list of float, the length of the lists must equal ``nrows-1`` and ``ncols-1``, respectively. height_ratios, width_ratios : list of float Ratios for the relative heights and widths for rows and columns of subplots, respectively. For example, ``width_ratios=(1,2)`` scales a 2-column gridspec so that the second column is twice as wide as the first column. left, right, top, bottom : float or str Passed to `~matplotlib.gridspec.GridSpec`, denotes the margin positions in figure-relative coordinates. Other parameters ---------------- **kwargs Passed to `~matplotlib.gridspec.GridSpec`. """ self._nrows = nrows * 2 - 1 # used with get_geometry self._ncols = ncols * 2 - 1 self._nrows_active = nrows self._ncols_active = ncols wratios, hratios, kwargs = self._spaces_as_ratios(**kwargs) super().__init__( self._nrows, self._ncols, hspace=0, wspace=0, # replaced with "hidden" slots width_ratios=wratios, height_ratios=hratios, figure=figure, **kwargs )
[docs] def __getitem__(self, key): """ Magic obfuscation that renders `~matplotlib.gridspec.GridSpec` rows and columns designated as 'spaces' inaccessible. """ nrows, ncols = self.get_geometry() nrows_active, ncols_active = self.get_active_geometry() if not isinstance(key, tuple): # usage gridspec[1,2] num1, num2 = self._normalize(key, nrows_active * ncols_active) else: if len(key) == 2: k1, k2 = key else: raise ValueError(f'Invalid index {key!r}.') num1 = self._normalize(k1, nrows_active) num2 = self._normalize(k2, ncols_active) num1, num2 = np.ravel_multi_index((num1, num2), (nrows, ncols)) num1 = self._positem(num1) num2 = self._positem(num2) return SubplotSpec(self, num1, num2)
@staticmethod def _positem(size): """ Account for negative indices. """ if size < 0: # want -1 to stay -1, -2 becomes -3, etc. return 2 * (size + 1) - 1 else: return size * 2 @staticmethod def _normalize(key, size): """ Transform gridspec index into standardized form. """ if isinstance(key, slice): start, stop, _ = key.indices(size) if stop > start: return start, stop - 1 else: if key < 0: key += size if 0 <= key < size: return key, key raise IndexError(f'Invalid index: {key} with size {size}.') def _spaces_as_ratios( self, hspace=None, wspace=None, # spacing between axes height_ratios=None, width_ratios=None, **kwargs ): """ For keyword argument usage, see `GridSpec`. """ # Parse flexible input nrows, ncols = self.get_active_geometry() hratios = np.atleast_1d(_not_none(height_ratios, 1)) wratios = np.atleast_1d(_not_none(width_ratios, 1)) hspace = np.atleast_1d(_not_none(hspace, np.mean(hratios) * 0.10)) # relative wspace = np.atleast_1d(_not_none(wspace, np.mean(wratios) * 0.10)) if len(wspace) == 1: wspace = np.repeat(wspace, (ncols - 1,)) # note: may be length 0 if len(hspace) == 1: hspace = np.repeat(hspace, (nrows - 1,)) if len(wratios) == 1: wratios = np.repeat(wratios, (ncols,)) if len(hratios) == 1: hratios = np.repeat(hratios, (nrows,)) # Verify input ratios and spacings # Translate height/width spacings, implement as extra columns/rows if len(hratios) != nrows: raise ValueError(f'Got {nrows} rows, but {len(hratios)} hratios.') if len(wratios) != ncols: raise ValueError( f'Got {ncols} columns, but {len(wratios)} wratios.' ) if len(wspace) != ncols - 1: raise ValueError( f'Require {ncols-1} width spacings for {ncols} columns, ' f'got {len(wspace)}.' ) if len(hspace) != nrows - 1: raise ValueError( f'Require {nrows-1} height spacings for {nrows} rows, ' f'got {len(hspace)}.' ) # Assign spacing as ratios nrows, ncols = self.get_geometry() wratios_final = [None] * ncols wratios_final[::2] = [*wratios] if ncols > 1: wratios_final[1::2] = [*wspace] hratios_final = [None] * nrows hratios_final[::2] = [*hratios] if nrows > 1: hratios_final[1::2] = [*hspace] return wratios_final, hratios_final, kwargs # bring extra kwargs back
[docs] def get_margins(self): """ Returns left, bottom, right, top values. Not sure why this method doesn't already exist on `~matplotlib.gridspec.GridSpec`. """ return self.left, self.bottom, self.right, self.top
[docs] def get_hspace(self): """ Returns row ratios allocated for spaces. """ return self.get_height_ratios()[1::2]
[docs] def get_wspace(self): """ Returns column ratios allocated for spaces. """ return self.get_width_ratios()[1::2]
[docs] def get_active_height_ratios(self): """ Returns height ratios excluding slots allocated for spaces. """ return self.get_height_ratios()[::2]
[docs] def get_active_width_ratios(self): """ Returns width ratios excluding slots allocated for spaces. """ return self.get_width_ratios()[::2]
[docs] def get_active_geometry(self): """ Returns the number of active rows and columns, i.e. the rows and columns that aren't skipped by `~GridSpec.__getitem__`. """ return self._nrows_active, self._ncols_active
[docs] def update(self, **kwargs): """ Update the gridspec with arbitrary initialization keyword arguments then *apply* those updates to every figure using this gridspec. The default `~matplotlib.gridspec.GridSpec.update` tries to update positions for axes on all active figures -- but this can fail after successive figure edits if it has been removed from the figure manager. ProPlot insists one gridspec per figure. Parameters ---------- **kwargs Valid initialization keyword arguments. See `GridSpec`. """ # Convert spaces to ratios wratios, hratios, kwargs = self._spaces_as_ratios(**kwargs) self.set_width_ratios(wratios) self.set_height_ratios(hratios) # Validate args nrows = kwargs.pop('nrows', None) ncols = kwargs.pop('ncols', None) nrows_current, ncols_current = self.get_active_geometry() if ( nrows is not None and nrows != nrows_current or ncols is not None and ncols != ncols_current ): raise ValueError( f'Input geometry {(nrows, ncols)} does not match ' f'current geometry {(nrows_current, ncols_current)}.' ) self.left = kwargs.pop('left', None) self.right = kwargs.pop('right', None) self.bottom = kwargs.pop('bottom', None) self.top = kwargs.pop('top', None) if kwargs: raise ValueError(f'Unknown keyword arg(s): {kwargs}.') # Apply to figure and all axes fig = self.figure fig.subplotpars.update(self.left, self.bottom, self.right, self.top) for ax in fig.axes: if not isinstance(ax, maxes.SubplotBase): continue subplotspec = ax.get_subplotspec().get_topmost_subplotspec() if subplotspec.get_gridspec() is not self: continue ax.update_params() ax.set_position(ax.figbox) fig.stale = True
class _GridSpecFromSubplotSpec(mgridspec.GridSpecFromSubplotSpec): """ Subclass that generates `SubplotSpec` objects. Avoids a recent deprecation of `~matplotlib.gridspec.SubplotSpec.get_rows_columns`. This is currently only used to draw colorbars with partial span along subplot edges but will not be needed once proplot implements "edge stacks." """ def __getitem__(self, key): subplotspec = super().__getitem__(key) return SubplotSpec._from_subplotspec(subplotspec)