"""

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)