548 lines
21 KiB
Python
548 lines
21 KiB
Python
"""
|
|
A layoutgrid is a nrows by ncols set of boxes, meant to be used by
|
|
`._constrained_layout`, each box is analogous to a subplotspec element of
|
|
a gridspec.
|
|
|
|
Each box is defined by left[ncols], right[ncols], bottom[nrows] and top[nrows],
|
|
and by two editable margins for each side. The main margin gets its value
|
|
set by the size of ticklabels, titles, etc on each axes that is in the figure.
|
|
The outer margin is the padding around the axes, and space for any
|
|
colorbars.
|
|
|
|
The "inner" widths and heights of these boxes are then constrained to be the
|
|
same (relative the values of `width_ratios[ncols]` and `height_ratios[nrows]`).
|
|
|
|
The layoutgrid is then constrained to be contained within a parent layoutgrid,
|
|
its column(s) and row(s) specified when it is created.
|
|
"""
|
|
|
|
import itertools
|
|
import kiwisolver as kiwi
|
|
import logging
|
|
import numpy as np
|
|
|
|
import matplotlib as mpl
|
|
import matplotlib.patches as mpatches
|
|
from matplotlib.transforms import Bbox
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
class LayoutGrid:
|
|
"""
|
|
Analogous to a gridspec, and contained in another LayoutGrid.
|
|
"""
|
|
|
|
def __init__(self, parent=None, parent_pos=(0, 0),
|
|
parent_inner=False, name='', ncols=1, nrows=1,
|
|
h_pad=None, w_pad=None, width_ratios=None,
|
|
height_ratios=None):
|
|
Variable = kiwi.Variable
|
|
self.parent_pos = parent_pos
|
|
self.parent_inner = parent_inner
|
|
self.name = name + seq_id()
|
|
if isinstance(parent, LayoutGrid):
|
|
self.name = f'{parent.name}.{self.name}'
|
|
self.nrows = nrows
|
|
self.ncols = ncols
|
|
self.height_ratios = np.atleast_1d(height_ratios)
|
|
if height_ratios is None:
|
|
self.height_ratios = np.ones(nrows)
|
|
self.width_ratios = np.atleast_1d(width_ratios)
|
|
if width_ratios is None:
|
|
self.width_ratios = np.ones(ncols)
|
|
|
|
sn = self.name + '_'
|
|
if not isinstance(parent, LayoutGrid):
|
|
# parent can be a rect if not a LayoutGrid
|
|
# allows specifying a rectangle to contain the layout.
|
|
self.solver = kiwi.Solver()
|
|
else:
|
|
parent.add_child(self, *parent_pos)
|
|
self.solver = parent.solver
|
|
# keep track of artist associated w/ this layout. Can be none
|
|
self.artists = np.empty((nrows, ncols), dtype=object)
|
|
self.children = np.empty((nrows, ncols), dtype=object)
|
|
|
|
self.margins = {}
|
|
self.margin_vals = {}
|
|
# all the boxes in each column share the same left/right margins:
|
|
for todo in ['left', 'right', 'leftcb', 'rightcb']:
|
|
# track the value so we can change only if a margin is larger
|
|
# than the current value
|
|
self.margin_vals[todo] = np.zeros(ncols)
|
|
|
|
sol = self.solver
|
|
|
|
self.lefts = [Variable(f'{sn}lefts[{i}]') for i in range(ncols)]
|
|
self.rights = [Variable(f'{sn}rights[{i}]') for i in range(ncols)]
|
|
for todo in ['left', 'right', 'leftcb', 'rightcb']:
|
|
self.margins[todo] = [Variable(f'{sn}margins[{todo}][{i}]')
|
|
for i in range(ncols)]
|
|
for i in range(ncols):
|
|
sol.addEditVariable(self.margins[todo][i], 'strong')
|
|
|
|
for todo in ['bottom', 'top', 'bottomcb', 'topcb']:
|
|
self.margins[todo] = np.empty((nrows), dtype=object)
|
|
self.margin_vals[todo] = np.zeros(nrows)
|
|
|
|
self.bottoms = [Variable(f'{sn}bottoms[{i}]') for i in range(nrows)]
|
|
self.tops = [Variable(f'{sn}tops[{i}]') for i in range(nrows)]
|
|
for todo in ['bottom', 'top', 'bottomcb', 'topcb']:
|
|
self.margins[todo] = [Variable(f'{sn}margins[{todo}][{i}]')
|
|
for i in range(nrows)]
|
|
for i in range(nrows):
|
|
sol.addEditVariable(self.margins[todo][i], 'strong')
|
|
|
|
# set these margins to zero by default. They will be edited as
|
|
# children are filled.
|
|
self.reset_margins()
|
|
self.add_constraints(parent)
|
|
|
|
self.h_pad = h_pad
|
|
self.w_pad = w_pad
|
|
|
|
def __repr__(self):
|
|
str = f'LayoutBox: {self.name:25s} {self.nrows}x{self.ncols},\n'
|
|
for i in range(self.nrows):
|
|
for j in range(self.ncols):
|
|
str += f'{i}, {j}: '\
|
|
f'L{self.lefts[j].value():1.3f}, ' \
|
|
f'B{self.bottoms[i].value():1.3f}, ' \
|
|
f'R{self.rights[j].value():1.3f}, ' \
|
|
f'T{self.tops[i].value():1.3f}, ' \
|
|
f'ML{self.margins["left"][j].value():1.3f}, ' \
|
|
f'MR{self.margins["right"][j].value():1.3f}, ' \
|
|
f'MB{self.margins["bottom"][i].value():1.3f}, ' \
|
|
f'MT{self.margins["top"][i].value():1.3f}, \n'
|
|
return str
|
|
|
|
def reset_margins(self):
|
|
"""
|
|
Reset all the margins to zero. Must do this after changing
|
|
figure size, for instance, because the relative size of the
|
|
axes labels etc changes.
|
|
"""
|
|
for todo in ['left', 'right', 'bottom', 'top',
|
|
'leftcb', 'rightcb', 'bottomcb', 'topcb']:
|
|
self.edit_margins(todo, 0.0)
|
|
|
|
def add_constraints(self, parent):
|
|
# define self-consistent constraints
|
|
self.hard_constraints()
|
|
# define relationship with parent layoutgrid:
|
|
self.parent_constraints(parent)
|
|
# define relative widths of the grid cells to each other
|
|
# and stack horizontally and vertically.
|
|
self.grid_constraints()
|
|
|
|
def hard_constraints(self):
|
|
"""
|
|
These are the redundant constraints, plus ones that make the
|
|
rest of the code easier.
|
|
"""
|
|
for i in range(self.ncols):
|
|
hc = [self.rights[i] >= self.lefts[i],
|
|
(self.rights[i] - self.margins['right'][i] -
|
|
self.margins['rightcb'][i] >=
|
|
self.lefts[i] - self.margins['left'][i] -
|
|
self.margins['leftcb'][i])
|
|
]
|
|
for c in hc:
|
|
self.solver.addConstraint(c | 'required')
|
|
|
|
for i in range(self.nrows):
|
|
hc = [self.tops[i] >= self.bottoms[i],
|
|
(self.tops[i] - self.margins['top'][i] -
|
|
self.margins['topcb'][i] >=
|
|
self.bottoms[i] - self.margins['bottom'][i] -
|
|
self.margins['bottomcb'][i])
|
|
]
|
|
for c in hc:
|
|
self.solver.addConstraint(c | 'required')
|
|
|
|
def add_child(self, child, i=0, j=0):
|
|
# np.ix_ returns the cross product of i and j indices
|
|
self.children[np.ix_(np.atleast_1d(i), np.atleast_1d(j))] = child
|
|
|
|
def parent_constraints(self, parent):
|
|
# constraints that are due to the parent...
|
|
# i.e. the first column's left is equal to the
|
|
# parent's left, the last column right equal to the
|
|
# parent's right...
|
|
if not isinstance(parent, LayoutGrid):
|
|
# specify a rectangle in figure coordinates
|
|
hc = [self.lefts[0] == parent[0],
|
|
self.rights[-1] == parent[0] + parent[2],
|
|
# top and bottom reversed order...
|
|
self.tops[0] == parent[1] + parent[3],
|
|
self.bottoms[-1] == parent[1]]
|
|
else:
|
|
rows, cols = self.parent_pos
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
left = parent.lefts[cols[0]]
|
|
right = parent.rights[cols[-1]]
|
|
top = parent.tops[rows[0]]
|
|
bottom = parent.bottoms[rows[-1]]
|
|
if self.parent_inner:
|
|
# the layout grid is contained inside the inner
|
|
# grid of the parent.
|
|
left += parent.margins['left'][cols[0]]
|
|
left += parent.margins['leftcb'][cols[0]]
|
|
right -= parent.margins['right'][cols[-1]]
|
|
right -= parent.margins['rightcb'][cols[-1]]
|
|
top -= parent.margins['top'][rows[0]]
|
|
top -= parent.margins['topcb'][rows[0]]
|
|
bottom += parent.margins['bottom'][rows[-1]]
|
|
bottom += parent.margins['bottomcb'][rows[-1]]
|
|
hc = [self.lefts[0] == left,
|
|
self.rights[-1] == right,
|
|
# from top to bottom
|
|
self.tops[0] == top,
|
|
self.bottoms[-1] == bottom]
|
|
for c in hc:
|
|
self.solver.addConstraint(c | 'required')
|
|
|
|
def grid_constraints(self):
|
|
# constrain the ratio of the inner part of the grids
|
|
# to be the same (relative to width_ratios)
|
|
|
|
# constrain widths:
|
|
w = (self.rights[0] - self.margins['right'][0] -
|
|
self.margins['rightcb'][0])
|
|
w = (w - self.lefts[0] - self.margins['left'][0] -
|
|
self.margins['leftcb'][0])
|
|
w0 = w / self.width_ratios[0]
|
|
# from left to right
|
|
for i in range(1, self.ncols):
|
|
w = (self.rights[i] - self.margins['right'][i] -
|
|
self.margins['rightcb'][i])
|
|
w = (w - self.lefts[i] - self.margins['left'][i] -
|
|
self.margins['leftcb'][i])
|
|
c = (w == w0 * self.width_ratios[i])
|
|
self.solver.addConstraint(c | 'strong')
|
|
# constrain the grid cells to be directly next to each other.
|
|
c = (self.rights[i - 1] == self.lefts[i])
|
|
self.solver.addConstraint(c | 'strong')
|
|
|
|
# constrain heights:
|
|
h = self.tops[0] - self.margins['top'][0] - self.margins['topcb'][0]
|
|
h = (h - self.bottoms[0] - self.margins['bottom'][0] -
|
|
self.margins['bottomcb'][0])
|
|
h0 = h / self.height_ratios[0]
|
|
# from top to bottom:
|
|
for i in range(1, self.nrows):
|
|
h = (self.tops[i] - self.margins['top'][i] -
|
|
self.margins['topcb'][i])
|
|
h = (h - self.bottoms[i] - self.margins['bottom'][i] -
|
|
self.margins['bottomcb'][i])
|
|
c = (h == h0 * self.height_ratios[i])
|
|
self.solver.addConstraint(c | 'strong')
|
|
# constrain the grid cells to be directly above each other.
|
|
c = (self.bottoms[i - 1] == self.tops[i])
|
|
self.solver.addConstraint(c | 'strong')
|
|
|
|
# Margin editing: The margins are variable and meant to
|
|
# contain things of a fixed size like axes labels, tick labels, titles
|
|
# etc
|
|
def edit_margin(self, todo, size, cell):
|
|
"""
|
|
Change the size of the margin for one cell.
|
|
|
|
Parameters
|
|
----------
|
|
todo : string (one of 'left', 'right', 'bottom', 'top')
|
|
margin to alter.
|
|
|
|
size : float
|
|
Size of the margin. If it is larger than the existing minimum it
|
|
updates the margin size. Fraction of figure size.
|
|
|
|
cell : int
|
|
Cell column or row to edit.
|
|
"""
|
|
self.solver.suggestValue(self.margins[todo][cell], size)
|
|
self.margin_vals[todo][cell] = size
|
|
|
|
def edit_margin_min(self, todo, size, cell=0):
|
|
"""
|
|
Change the minimum size of the margin for one cell.
|
|
|
|
Parameters
|
|
----------
|
|
todo : string (one of 'left', 'right', 'bottom', 'top')
|
|
margin to alter.
|
|
|
|
size : float
|
|
Minimum size of the margin . If it is larger than the
|
|
existing minimum it updates the margin size. Fraction of
|
|
figure size.
|
|
|
|
cell : int
|
|
Cell column or row to edit.
|
|
"""
|
|
|
|
if size > self.margin_vals[todo][cell]:
|
|
self.edit_margin(todo, size, cell)
|
|
|
|
def edit_margins(self, todo, size):
|
|
"""
|
|
Change the size of all the margin of all the cells in the layout grid.
|
|
|
|
Parameters
|
|
----------
|
|
todo : string (one of 'left', 'right', 'bottom', 'top')
|
|
margin to alter.
|
|
|
|
size : float
|
|
Size to set the margins. Fraction of figure size.
|
|
"""
|
|
|
|
for i in range(len(self.margin_vals[todo])):
|
|
self.edit_margin(todo, size, i)
|
|
|
|
def edit_all_margins_min(self, todo, size):
|
|
"""
|
|
Change the minimum size of all the margin of all
|
|
the cells in the layout grid.
|
|
|
|
Parameters
|
|
----------
|
|
todo : {'left', 'right', 'bottom', 'top'}
|
|
The margin to alter.
|
|
|
|
size : float
|
|
Minimum size of the margin. If it is larger than the
|
|
existing minimum it updates the margin size. Fraction of
|
|
figure size.
|
|
"""
|
|
|
|
for i in range(len(self.margin_vals[todo])):
|
|
self.edit_margin_min(todo, size, i)
|
|
|
|
def edit_outer_margin_mins(self, margin, ss):
|
|
"""
|
|
Edit all four margin minimums in one statement.
|
|
|
|
Parameters
|
|
----------
|
|
margin : dict
|
|
size of margins in a dict with keys 'left', 'right', 'bottom',
|
|
'top'
|
|
|
|
ss : SubplotSpec
|
|
defines the subplotspec these margins should be applied to
|
|
"""
|
|
|
|
self.edit_margin_min('left', margin['left'], ss.colspan.start)
|
|
self.edit_margin_min('leftcb', margin['leftcb'], ss.colspan.start)
|
|
self.edit_margin_min('right', margin['right'], ss.colspan.stop - 1)
|
|
self.edit_margin_min('rightcb', margin['rightcb'], ss.colspan.stop - 1)
|
|
# rows are from the top down:
|
|
self.edit_margin_min('top', margin['top'], ss.rowspan.start)
|
|
self.edit_margin_min('topcb', margin['topcb'], ss.rowspan.start)
|
|
self.edit_margin_min('bottom', margin['bottom'], ss.rowspan.stop - 1)
|
|
self.edit_margin_min('bottomcb', margin['bottomcb'],
|
|
ss.rowspan.stop - 1)
|
|
|
|
def get_margins(self, todo, col):
|
|
"""Return the margin at this position"""
|
|
return self.margin_vals[todo][col]
|
|
|
|
def get_outer_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the outer bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
self.lefts[cols[0]].value(),
|
|
self.bottoms[rows[-1]].value(),
|
|
self.rights[cols[-1]].value(),
|
|
self.tops[rows[0]].value())
|
|
return bbox
|
|
|
|
def get_inner_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the inner bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.lefts[cols[0]].value() +
|
|
self.margins['left'][cols[0]].value() +
|
|
self.margins['leftcb'][cols[0]].value()),
|
|
(self.bottoms[rows[-1]].value() +
|
|
self.margins['bottom'][rows[-1]].value() +
|
|
self.margins['bottomcb'][rows[-1]].value()),
|
|
(self.rights[cols[-1]].value() -
|
|
self.margins['right'][cols[-1]].value() -
|
|
self.margins['rightcb'][cols[-1]].value()),
|
|
(self.tops[rows[0]].value() -
|
|
self.margins['top'][rows[0]].value() -
|
|
self.margins['topcb'][rows[0]].value())
|
|
)
|
|
return bbox
|
|
|
|
def get_bbox_for_cb(self, rows=0, cols=0):
|
|
"""
|
|
Return the bounding box that includes the
|
|
decorations but, *not* the colorbar...
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.lefts[cols[0]].value() +
|
|
self.margins['leftcb'][cols[0]].value()),
|
|
(self.bottoms[rows[-1]].value() +
|
|
self.margins['bottomcb'][rows[-1]].value()),
|
|
(self.rights[cols[-1]].value() -
|
|
self.margins['rightcb'][cols[-1]].value()),
|
|
(self.tops[rows[0]].value() -
|
|
self.margins['topcb'][rows[0]].value())
|
|
)
|
|
return bbox
|
|
|
|
def get_left_margin_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the left margin bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.lefts[cols[0]].value() +
|
|
self.margins['leftcb'][cols[0]].value()),
|
|
(self.bottoms[rows[-1]].value()),
|
|
(self.lefts[cols[0]].value() +
|
|
self.margins['leftcb'][cols[0]].value() +
|
|
self.margins['left'][cols[0]].value()),
|
|
(self.tops[rows[0]].value()))
|
|
return bbox
|
|
|
|
def get_bottom_margin_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the left margin bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.lefts[cols[0]].value()),
|
|
(self.bottoms[rows[-1]].value() +
|
|
self.margins['bottomcb'][rows[-1]].value()),
|
|
(self.rights[cols[-1]].value()),
|
|
(self.bottoms[rows[-1]].value() +
|
|
self.margins['bottom'][rows[-1]].value() +
|
|
self.margins['bottomcb'][rows[-1]].value()
|
|
))
|
|
return bbox
|
|
|
|
def get_right_margin_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the left margin bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.rights[cols[-1]].value() -
|
|
self.margins['right'][cols[-1]].value() -
|
|
self.margins['rightcb'][cols[-1]].value()),
|
|
(self.bottoms[rows[-1]].value()),
|
|
(self.rights[cols[-1]].value() -
|
|
self.margins['rightcb'][cols[-1]].value()),
|
|
(self.tops[rows[0]].value()))
|
|
return bbox
|
|
|
|
def get_top_margin_bbox(self, rows=0, cols=0):
|
|
"""
|
|
Return the left margin bounding box of the subplot specs
|
|
given by rows and cols. rows and cols can be spans.
|
|
"""
|
|
rows = np.atleast_1d(rows)
|
|
cols = np.atleast_1d(cols)
|
|
|
|
bbox = Bbox.from_extents(
|
|
(self.lefts[cols[0]].value()),
|
|
(self.tops[rows[0]].value() -
|
|
self.margins['topcb'][rows[0]].value()),
|
|
(self.rights[cols[-1]].value()),
|
|
(self.tops[rows[0]].value() -
|
|
self.margins['topcb'][rows[0]].value() -
|
|
self.margins['top'][rows[0]].value()))
|
|
return bbox
|
|
|
|
def update_variables(self):
|
|
"""
|
|
Update the variables for the solver attached to this layoutgrid.
|
|
"""
|
|
self.solver.updateVariables()
|
|
|
|
_layoutboxobjnum = itertools.count()
|
|
|
|
|
|
def seq_id():
|
|
"""Generate a short sequential id for layoutbox objects."""
|
|
return '%06d' % next(_layoutboxobjnum)
|
|
|
|
|
|
def plot_children(fig, lg=None, level=0):
|
|
"""Simple plotting to show where boxes are."""
|
|
if lg is None:
|
|
_layoutgrids = fig.get_layout_engine().execute(fig)
|
|
lg = _layoutgrids[fig]
|
|
colors = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
|
|
col = colors[level]
|
|
for i in range(lg.nrows):
|
|
for j in range(lg.ncols):
|
|
bb = lg.get_outer_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bb.p0, bb.width, bb.height, linewidth=1,
|
|
edgecolor='0.7', facecolor='0.7',
|
|
alpha=0.2, transform=fig.transFigure,
|
|
zorder=-3))
|
|
bbi = lg.get_inner_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bbi.p0, bbi.width, bbi.height, linewidth=2,
|
|
edgecolor=col, facecolor='none',
|
|
transform=fig.transFigure, zorder=-2))
|
|
|
|
bbi = lg.get_left_margin_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bbi.p0, bbi.width, bbi.height, linewidth=0,
|
|
edgecolor='none', alpha=0.2,
|
|
facecolor=[0.5, 0.7, 0.5],
|
|
transform=fig.transFigure, zorder=-2))
|
|
bbi = lg.get_right_margin_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bbi.p0, bbi.width, bbi.height, linewidth=0,
|
|
edgecolor='none', alpha=0.2,
|
|
facecolor=[0.7, 0.5, 0.5],
|
|
transform=fig.transFigure, zorder=-2))
|
|
bbi = lg.get_bottom_margin_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bbi.p0, bbi.width, bbi.height, linewidth=0,
|
|
edgecolor='none', alpha=0.2,
|
|
facecolor=[0.5, 0.5, 0.7],
|
|
transform=fig.transFigure, zorder=-2))
|
|
bbi = lg.get_top_margin_bbox(rows=i, cols=j)
|
|
fig.add_artist(
|
|
mpatches.Rectangle(bbi.p0, bbi.width, bbi.height, linewidth=0,
|
|
edgecolor='none', alpha=0.2,
|
|
facecolor=[0.7, 0.2, 0.7],
|
|
transform=fig.transFigure, zorder=-2))
|
|
for ch in lg.children.flat:
|
|
if ch is not None:
|
|
plot_children(fig, ch, level=level+1)
|