Improve tight layout algorithm handling of empty gridspec slots

[zxdawn]

Description

When some subplots have colorbar and some don't, the position will be wrong.

Steps to reproduce

import proplot as pplt
import numpy as np

# Colorbars
fig, axs = pplt.subplots(nrows=2, ncols=2)
state = np.random.RandomState(51423)

m = axs.heatmap(state.rand(10, 10), cmap='dusk')
axs[0].colorbar(m[0], loc='r', label='test')
axs[1].colorbar(m[0], loc='r', label='test')
axs[3].colorbar(m[0], loc='l', label='test')
axs[3].colorbar(m[0], loc='r', label='test')
axs[3].colorbar(m[0], loc='r', label='test')

Actual behavior:

Proplot version

Paste the results of import matplotlib; print(matplotlib.__version__); import proplot; print(proplot.version)here.

3.5.0
0.9.5.post105

[lukelbd]

🤨 looks like another matplotlib 3.5 issue. Will follow up

For anyone reading this thread: As a workaround, you can use e.g. pplt.subplots(..., wspace=4) to manually allocate the necessary space between the two subplot columns.

[lukelbd]

Seems this is present in the docs build environment too (matplotlib 3.2.1). I guess it's some really subtle issue with the tight layout algorithm that doesn't come up often. This might have to wait until the next version (want to get 0.10.0 out).

[lukelbd]

Oh wait I can see the issue now... I thought it was some weird draw-time race condition problem but in fact it's not that subtle. It's just that the tight layout algorithm only compares axes extents for subplots in adjacent gridspec slots. If the adjacent slot is empty, rather than going to the next slot over, it simply ignores that slot.

Here's the general idea:

• In this example, there are two "panel" gridspec slots between the two columns of main subplots: one for a right colorbar, (used by the top-left subplot) and the other for a left colorbar (used by the bottom-right subplot). Proplot has to determine the necessary space between these two panel slots to prevent overlapping labels. It does this row-by-row.
• For the first row, proplot sees that there is no subplot in the panel slot to the right of the right colorbar. So, the tight layout algorithm allocates zero space for the top row.
• For the second row, proplot sees that there is no subplot in the panel slot to the left of the left colorbar. So, the tight layout algorithm again allocates zero space for the bottom row.
• In total, the maximum space "needed" between the panel slots is max((0, 0)) = 0. However, the labels are large enough that they actually extend over the empty slot and bleed into the next one.

The solution is: If a gridspec slot is empty, instead of just skipping the tight layout test, I should jump to the next slot over and get the space between those objects. Shouldn't be too hard to fix.

I could also consider adding some rc setting and keyword argument that makes the tight layout algorithm check extents between all subplots in a given column rather than this row-by-row calculation. In this example, that behavior would push the top-left subplot colorbar "label" to the left of the bottom-right subplot colorbar "label".

[lukelbd]

Done! Here's how this looks now:

import proplot as pplt
import numpy as np
fig, axs = pplt.subplots(nrows=2, ncols=2)
state = np.random.RandomState(51423)
m = axs.heatmap(state.rand(10, 10), cmap='dusk')
axs[0].colorbar(m[0], loc='r', label='test')
axs[3].colorbar(m[0], loc='l', label='test')

[lukelbd]

Also decided to add that option I mentioned. Now similar to wequal, hequal, and equal there are wgroup, hgroup, and group options that can be passed to subplots (or Figure or GridSpec). This toggles whether all rows and columns are considered or whether just adjacent "groups" are considered. I've also added rc['subplots.equalspace'] and rc['subplots.groupspace'] settings that lets you toggle these persistently.

Here is a demonstration from your example, with group turned on (the default behavior) then turned off:

import proplot as pplt
import numpy as np
for group in (True, False):
    fig, axs = pplt.subplots(nrows=2, ncols=2, group=group)
    state = np.random.RandomState(51423)
    m = axs.heatmap(state.rand(10, 10), cmap='dusk')
    axs[0].colorbar(m[0], loc='r', label='test')
    axs[3].colorbar(m[0], loc='l', label='test')
    fig.save('~/Downloads/test1.png', dpi=300)

And here's a new documentation example:

import proplot as pplt

# Stress test of the tight layout algorithm
# Add large labels along the edge of one subplot
equals = [('unequal', False), ('equal', True)]
groups = [('grouped', True), ('ungrouped', False)]
for i in range(2):
    name1, equal = equals[i]
    name2, group = groups[i]
    print(group)
    fig, axs = pplt.subplots(
        nrows=3, ncols=3, refwidth=1.1, share=False, equal=equal, group=group,
    )
    axs[1].format(
        xlabel='xlabel\nxlabel',
        ylabel='ylabel\nylabel\nylabel\nylabel'
    )
    axs[3:6:2].format(
        title='Title\nTitle',
        titlesize='med',
    )
    axs.format(
        grid=False,
        toplabels=('Column 1', 'Column 2', 'Column 3'),
        leftlabels=('Row 1', 'Row 2', 'Row 3'),
        suptitle=f'Tight layout with "{name1}" and "{name2}" row-column spacing',
    )

[zxdawn]

Nice idea of equalspace and groupspace!