403 lines
13 KiB
Python
403 lines
13 KiB
Python
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from matplotlib.testing.decorators import image_comparison
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import matplotlib.gridspec as gridspec
|
||
|
|
from matplotlib import ticker, rcParams
|
||
|
|
|
||
|
|
|
||
|
|
def example_plot(ax, fontsize=12, nodec=False):
|
||
|
|
ax.plot([1, 2])
|
||
|
|
ax.locator_params(nbins=3)
|
||
|
|
if not nodec:
|
||
|
|
ax.set_xlabel('x-label', fontsize=fontsize)
|
||
|
|
ax.set_ylabel('y-label', fontsize=fontsize)
|
||
|
|
ax.set_title('Title', fontsize=fontsize)
|
||
|
|
else:
|
||
|
|
ax.set_xticklabels('')
|
||
|
|
ax.set_yticklabels('')
|
||
|
|
|
||
|
|
|
||
|
|
def example_pcolor(ax, fontsize=12):
|
||
|
|
dx, dy = 0.6, 0.6
|
||
|
|
y, x = np.mgrid[slice(-3, 3 + dy, dy),
|
||
|
|
slice(-3, 3 + dx, dx)]
|
||
|
|
z = (1 - x / 2. + x ** 5 + y ** 3) * np.exp(-x ** 2 - y ** 2)
|
||
|
|
pcm = ax.pcolormesh(x, y, z, cmap='RdBu_r', vmin=-1., vmax=1.,
|
||
|
|
rasterized=True)
|
||
|
|
# ax.locator_params(nbins=3)
|
||
|
|
ax.set_xlabel('x-label', fontsize=fontsize)
|
||
|
|
ax.set_ylabel('y-label', fontsize=fontsize)
|
||
|
|
ax.set_title('Title', fontsize=fontsize)
|
||
|
|
return pcm
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout1.png'])
|
||
|
|
def test_constrained_layout1():
|
||
|
|
'Test constrained_layout for a single subplot'
|
||
|
|
fig = plt.figure(constrained_layout=True)
|
||
|
|
ax = fig.add_subplot(111)
|
||
|
|
example_plot(ax, fontsize=24)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout2.png'])
|
||
|
|
def test_constrained_layout2():
|
||
|
|
'Test constrained_layout for 2x2 subplots'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
example_plot(ax, fontsize=24)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout3.png'])
|
||
|
|
def test_constrained_layout3():
|
||
|
|
'Test constrained_layout for colorbars with subplots'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for nn, ax in enumerate(axs.flat):
|
||
|
|
pcm = example_pcolor(ax, fontsize=24)
|
||
|
|
if nn == 3:
|
||
|
|
pad = 0.08
|
||
|
|
else:
|
||
|
|
pad = 0.02 # default
|
||
|
|
fig.colorbar(pcm, ax=ax, pad=pad)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout4'])
|
||
|
|
def test_constrained_layout4():
|
||
|
|
'Test constrained_layout for a single colorbar with subplots'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax, fontsize=24)
|
||
|
|
fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout5.png'], tol=5.e-2)
|
||
|
|
def test_constrained_layout5():
|
||
|
|
'''
|
||
|
|
Test constrained_layout for a single colorbar with subplots,
|
||
|
|
colorbar bottom
|
||
|
|
'''
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax, fontsize=24)
|
||
|
|
fig.colorbar(pcm, ax=axs,
|
||
|
|
use_gridspec=False, pad=0.01, shrink=0.6,
|
||
|
|
location='bottom')
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout6.png'])
|
||
|
|
def test_constrained_layout6():
|
||
|
|
'Test constrained_layout for nested gridspecs'
|
||
|
|
fig = plt.figure(constrained_layout=True)
|
||
|
|
gs = fig.add_gridspec(1, 2, figure=fig)
|
||
|
|
gsl = gs[0].subgridspec(2, 2)
|
||
|
|
gsr = gs[1].subgridspec(1, 2)
|
||
|
|
axsl = []
|
||
|
|
for gs in gsl:
|
||
|
|
ax = fig.add_subplot(gs)
|
||
|
|
axsl += [ax]
|
||
|
|
example_plot(ax, fontsize=12)
|
||
|
|
ax.set_xlabel('x-label\nMultiLine')
|
||
|
|
axsr = []
|
||
|
|
for gs in gsr:
|
||
|
|
ax = fig.add_subplot(gs)
|
||
|
|
axsr += [ax]
|
||
|
|
pcm = example_pcolor(ax, fontsize=12)
|
||
|
|
|
||
|
|
fig.colorbar(pcm, ax=axsr,
|
||
|
|
pad=0.01, shrink=0.99, location='bottom',
|
||
|
|
ticks=ticker.MaxNLocator(nbins=5))
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout7():
|
||
|
|
'Test for proper warning if fig not set in GridSpec'
|
||
|
|
with pytest.warns(UserWarning, match='Calling figure.constrained_layout, '
|
||
|
|
'but figure not setup to do constrained layout'):
|
||
|
|
fig = plt.figure(constrained_layout=True)
|
||
|
|
gs = gridspec.GridSpec(1, 2)
|
||
|
|
gsl = gridspec.GridSpecFromSubplotSpec(2, 2, gs[0])
|
||
|
|
gsr = gridspec.GridSpecFromSubplotSpec(1, 2, gs[1])
|
||
|
|
for gs in gsl:
|
||
|
|
fig.add_subplot(gs)
|
||
|
|
# need to trigger a draw to get warning
|
||
|
|
fig.draw(fig.canvas.get_renderer())
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout8.png'])
|
||
|
|
def test_constrained_layout8():
|
||
|
|
'Test for gridspecs that are not completely full'
|
||
|
|
fig = plt.figure(figsize=(10, 5), constrained_layout=True)
|
||
|
|
gs = gridspec.GridSpec(3, 5, figure=fig)
|
||
|
|
axs = []
|
||
|
|
for j in [0, 1]:
|
||
|
|
if j == 0:
|
||
|
|
ilist = [1]
|
||
|
|
else:
|
||
|
|
ilist = [0, 4]
|
||
|
|
for i in ilist:
|
||
|
|
ax = fig.add_subplot(gs[j, i])
|
||
|
|
axs += [ax]
|
||
|
|
pcm = example_pcolor(ax, fontsize=9)
|
||
|
|
if i > 0:
|
||
|
|
ax.set_ylabel('')
|
||
|
|
if j < 1:
|
||
|
|
ax.set_xlabel('')
|
||
|
|
ax.set_title('')
|
||
|
|
ax = fig.add_subplot(gs[2, :])
|
||
|
|
axs += [ax]
|
||
|
|
pcm = example_pcolor(ax, fontsize=9)
|
||
|
|
|
||
|
|
fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout9.png'])
|
||
|
|
def test_constrained_layout9():
|
||
|
|
'Test for handling suptitle and for sharex and sharey'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True,
|
||
|
|
sharex=False, sharey=False)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax, fontsize=24)
|
||
|
|
ax.set_xlabel('')
|
||
|
|
ax.set_ylabel('')
|
||
|
|
ax.set_aspect(2.)
|
||
|
|
fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
|
||
|
|
fig.suptitle('Test Suptitle', fontsize=28)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout10.png'])
|
||
|
|
def test_constrained_layout10():
|
||
|
|
'Test for handling legend outside axis'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
ax.plot(np.arange(12), label='This is a label')
|
||
|
|
ax.legend(loc='center left', bbox_to_anchor=(0.8, 0.5))
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout11.png'])
|
||
|
|
def test_constrained_layout11():
|
||
|
|
'Test for multiple nested gridspecs'
|
||
|
|
fig = plt.figure(constrained_layout=True, figsize=(10, 3))
|
||
|
|
gs0 = gridspec.GridSpec(1, 2, figure=fig)
|
||
|
|
gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
|
||
|
|
gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1])
|
||
|
|
ax = fig.add_subplot(gs0[1])
|
||
|
|
example_plot(ax, fontsize=9)
|
||
|
|
axs = []
|
||
|
|
for gs in gsl0:
|
||
|
|
ax = fig.add_subplot(gs)
|
||
|
|
axs += [ax]
|
||
|
|
pcm = example_pcolor(ax, fontsize=9)
|
||
|
|
fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
|
||
|
|
ax = fig.add_subplot(gsl[0])
|
||
|
|
example_plot(ax, fontsize=9)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout11rat.png'])
|
||
|
|
def test_constrained_layout11rat():
|
||
|
|
'Test for multiple nested gridspecs with width_ratios'
|
||
|
|
fig = plt.figure(constrained_layout=True, figsize=(10, 3))
|
||
|
|
gs0 = gridspec.GridSpec(1, 2, figure=fig, width_ratios=[6., 1.])
|
||
|
|
gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
|
||
|
|
gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1],
|
||
|
|
height_ratios=[2., 1.])
|
||
|
|
ax = fig.add_subplot(gs0[1])
|
||
|
|
example_plot(ax, fontsize=9)
|
||
|
|
axs = []
|
||
|
|
for gs in gsl0:
|
||
|
|
ax = fig.add_subplot(gs)
|
||
|
|
axs += [ax]
|
||
|
|
pcm = example_pcolor(ax, fontsize=9)
|
||
|
|
fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
|
||
|
|
ax = fig.add_subplot(gsl[0])
|
||
|
|
example_plot(ax, fontsize=9)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout12.png'])
|
||
|
|
def test_constrained_layout12():
|
||
|
|
'Test that very unbalanced labeling still works.'
|
||
|
|
fig = plt.figure(constrained_layout=True)
|
||
|
|
|
||
|
|
gs0 = gridspec.GridSpec(6, 2, figure=fig)
|
||
|
|
|
||
|
|
ax1 = fig.add_subplot(gs0[:3, 1])
|
||
|
|
ax2 = fig.add_subplot(gs0[3:, 1])
|
||
|
|
|
||
|
|
example_plot(ax1, fontsize=24)
|
||
|
|
example_plot(ax2, fontsize=24)
|
||
|
|
|
||
|
|
ax = fig.add_subplot(gs0[0:2, 0])
|
||
|
|
example_plot(ax, nodec=True)
|
||
|
|
ax = fig.add_subplot(gs0[2:4, 0])
|
||
|
|
example_plot(ax, nodec=True)
|
||
|
|
ax = fig.add_subplot(gs0[4:, 0])
|
||
|
|
example_plot(ax, nodec=True)
|
||
|
|
ax.set_xlabel('x-label')
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout13.png'], tol=2.e-2)
|
||
|
|
def test_constrained_layout13():
|
||
|
|
'Test that padding works.'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax, fontsize=12)
|
||
|
|
fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
|
||
|
|
fig.set_constrained_layout_pads(w_pad=24./72., h_pad=24./72.)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout14.png'])
|
||
|
|
def test_constrained_layout14():
|
||
|
|
'Test that padding works.'
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax, fontsize=12)
|
||
|
|
fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
|
||
|
|
fig.set_constrained_layout_pads(
|
||
|
|
w_pad=3./72., h_pad=3./72.,
|
||
|
|
hspace=0.2, wspace=0.2)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout15.png'])
|
||
|
|
def test_constrained_layout15():
|
||
|
|
'Test that rcparams work.'
|
||
|
|
rcParams['figure.constrained_layout.use'] = True
|
||
|
|
fig, axs = plt.subplots(2, 2)
|
||
|
|
for ax in axs.flat:
|
||
|
|
example_plot(ax, fontsize=12)
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout16.png'])
|
||
|
|
def test_constrained_layout16():
|
||
|
|
'Test ax.set_position.'
|
||
|
|
fig, ax = plt.subplots(constrained_layout=True)
|
||
|
|
example_plot(ax, fontsize=12)
|
||
|
|
ax2 = fig.add_axes([0.2, 0.2, 0.4, 0.4])
|
||
|
|
|
||
|
|
|
||
|
|
@image_comparison(['constrained_layout17.png'])
|
||
|
|
def test_constrained_layout17():
|
||
|
|
'Test uneven gridspecs'
|
||
|
|
fig = plt.figure(constrained_layout=True)
|
||
|
|
gs = gridspec.GridSpec(3, 3, figure=fig)
|
||
|
|
|
||
|
|
ax1 = fig.add_subplot(gs[0, 0])
|
||
|
|
ax2 = fig.add_subplot(gs[0, 1:])
|
||
|
|
ax3 = fig.add_subplot(gs[1:, 0:2])
|
||
|
|
ax4 = fig.add_subplot(gs[1:, -1])
|
||
|
|
|
||
|
|
example_plot(ax1)
|
||
|
|
example_plot(ax2)
|
||
|
|
example_plot(ax3)
|
||
|
|
example_plot(ax4)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout18():
|
||
|
|
'Test twinx'
|
||
|
|
fig, ax = plt.subplots(constrained_layout=True)
|
||
|
|
ax2 = ax.twinx()
|
||
|
|
example_plot(ax)
|
||
|
|
example_plot(ax2, fontsize=24)
|
||
|
|
fig.canvas.draw()
|
||
|
|
assert all(ax.get_position().extents == ax2.get_position().extents)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout19():
|
||
|
|
'Test twiny'
|
||
|
|
fig, ax = plt.subplots(constrained_layout=True)
|
||
|
|
ax2 = ax.twiny()
|
||
|
|
example_plot(ax)
|
||
|
|
example_plot(ax2, fontsize=24)
|
||
|
|
ax2.set_title('')
|
||
|
|
ax.set_title('')
|
||
|
|
fig.canvas.draw()
|
||
|
|
assert all(ax.get_position().extents == ax2.get_position().extents)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout20():
|
||
|
|
'Smoke test cl does not mess up added axes'
|
||
|
|
gx = np.linspace(-5, 5, 4)
|
||
|
|
img = np.hypot(gx, gx[:, None])
|
||
|
|
|
||
|
|
fig = plt.figure()
|
||
|
|
ax = fig.add_axes([0, 0, 1, 1])
|
||
|
|
mesh = ax.pcolormesh(gx, gx, img)
|
||
|
|
fig.colorbar(mesh)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout21():
|
||
|
|
'#11035: repeated calls to suptitle should not alter the layout'
|
||
|
|
fig, ax = plt.subplots(constrained_layout=True)
|
||
|
|
|
||
|
|
fig.suptitle("Suptitle0")
|
||
|
|
fig.canvas.draw()
|
||
|
|
extents0 = np.copy(ax.get_position().extents)
|
||
|
|
|
||
|
|
fig.suptitle("Suptitle1")
|
||
|
|
fig.canvas.draw()
|
||
|
|
extents1 = np.copy(ax.get_position().extents)
|
||
|
|
|
||
|
|
np.testing.assert_allclose(extents0, extents1)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout22():
|
||
|
|
'#11035: suptitle should not be include in CL if manually positioned'
|
||
|
|
fig, ax = plt.subplots(constrained_layout=True)
|
||
|
|
|
||
|
|
fig.canvas.draw()
|
||
|
|
extents0 = np.copy(ax.get_position().extents)
|
||
|
|
|
||
|
|
fig.suptitle("Suptitle", y=0.5)
|
||
|
|
fig.canvas.draw()
|
||
|
|
extents1 = np.copy(ax.get_position().extents)
|
||
|
|
|
||
|
|
np.testing.assert_allclose(extents0, extents1)
|
||
|
|
|
||
|
|
|
||
|
|
def test_constrained_layout23():
|
||
|
|
'''
|
||
|
|
Comment in #11035: suptitle used to cause an exception when
|
||
|
|
reusing a figure w/ CL with ``clear=True``.
|
||
|
|
'''
|
||
|
|
|
||
|
|
for i in range(2):
|
||
|
|
fig, ax = plt.subplots(num="123", constrained_layout=True, clear=True)
|
||
|
|
fig.suptitle("Suptitle{}".format(i))
|
||
|
|
|
||
|
|
|
||
|
|
# This test occasionally fails the image comparison tests, so we mark as
|
||
|
|
# flaky. Apparently the constraint solver occasionally doesn't fully
|
||
|
|
# optimize. Would be nice if this were more deterministic...
|
||
|
|
@pytest.mark.timeout(30)
|
||
|
|
@pytest.mark.flaky(reruns=3)
|
||
|
|
@image_comparison(['test_colorbar_location.png'],
|
||
|
|
remove_text=True, style='mpl20')
|
||
|
|
def test_colorbar_location():
|
||
|
|
"""
|
||
|
|
Test that colorbar handling is as expected for various complicated
|
||
|
|
cases...
|
||
|
|
"""
|
||
|
|
|
||
|
|
fig, axs = plt.subplots(4, 5, constrained_layout=True)
|
||
|
|
for ax in axs.flat:
|
||
|
|
pcm = example_pcolor(ax)
|
||
|
|
ax.set_xlabel('')
|
||
|
|
ax.set_ylabel('')
|
||
|
|
fig.colorbar(pcm, ax=axs[:, 1], shrink=0.4)
|
||
|
|
fig.colorbar(pcm, ax=axs[-1, :2], shrink=0.5, location='bottom')
|
||
|
|
fig.colorbar(pcm, ax=axs[0, 2:], shrink=0.5, location='bottom')
|
||
|
|
fig.colorbar(pcm, ax=axs[-2, 3:], shrink=0.5, location='top')
|
||
|
|
fig.colorbar(pcm, ax=axs[0, 0], shrink=0.5, location='left')
|
||
|
|
fig.colorbar(pcm, ax=axs[1:3, 2], shrink=0.5, location='right')
|
||
|
|
|
||
|
|
|
||
|
|
def test_hidden_axes():
|
||
|
|
# test that if we make an axes not visible that constrained_layout
|
||
|
|
# still works. Note the axes still takes space in the layout
|
||
|
|
# (as does a gridspec slot that is empty)
|
||
|
|
fig, axs = plt.subplots(2, 2, constrained_layout=True)
|
||
|
|
axs[0, 1].set_visible(False)
|
||
|
|
fig.canvas.draw()
|
||
|
|
extents1 = np.copy(axs[0, 0].get_position().extents)
|
||
|
|
|
||
|
|
np.testing.assert_allclose(extents1,
|
||
|
|
[0.045552, 0.548288, 0.47319, 0.982638], rtol=1e-5)
|