test_constrainedlayout.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401
  1. import numpy as np
  2. import pytest
  3. from matplotlib.testing.decorators import image_comparison
  4. import matplotlib.pyplot as plt
  5. import matplotlib.gridspec as gridspec
  6. from matplotlib import ticker, rcParams
  7. def example_plot(ax, fontsize=12, nodec=False):
  8. ax.plot([1, 2])
  9. ax.locator_params(nbins=3)
  10. if not nodec:
  11. ax.set_xlabel('x-label', fontsize=fontsize)
  12. ax.set_ylabel('y-label', fontsize=fontsize)
  13. ax.set_title('Title', fontsize=fontsize)
  14. else:
  15. ax.set_xticklabels('')
  16. ax.set_yticklabels('')
  17. def example_pcolor(ax, fontsize=12):
  18. dx, dy = 0.6, 0.6
  19. y, x = np.mgrid[slice(-3, 3 + dy, dy),
  20. slice(-3, 3 + dx, dx)]
  21. z = (1 - x / 2. + x ** 5 + y ** 3) * np.exp(-x ** 2 - y ** 2)
  22. pcm = ax.pcolormesh(x, y, z[:-1, :-1], cmap='RdBu_r', vmin=-1., vmax=1.,
  23. rasterized=True)
  24. ax.set_xlabel('x-label', fontsize=fontsize)
  25. ax.set_ylabel('y-label', fontsize=fontsize)
  26. ax.set_title('Title', fontsize=fontsize)
  27. return pcm
  28. @image_comparison(['constrained_layout1.png'])
  29. def test_constrained_layout1():
  30. """Test constrained_layout for a single subplot"""
  31. fig = plt.figure(constrained_layout=True)
  32. ax = fig.add_subplot(111)
  33. example_plot(ax, fontsize=24)
  34. @image_comparison(['constrained_layout2.png'])
  35. def test_constrained_layout2():
  36. """Test constrained_layout for 2x2 subplots"""
  37. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  38. for ax in axs.flat:
  39. example_plot(ax, fontsize=24)
  40. @image_comparison(['constrained_layout3.png'])
  41. def test_constrained_layout3():
  42. """Test constrained_layout for colorbars with subplots"""
  43. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  44. for nn, ax in enumerate(axs.flat):
  45. pcm = example_pcolor(ax, fontsize=24)
  46. if nn == 3:
  47. pad = 0.08
  48. else:
  49. pad = 0.02 # default
  50. fig.colorbar(pcm, ax=ax, pad=pad)
  51. @image_comparison(['constrained_layout4'])
  52. def test_constrained_layout4():
  53. """Test constrained_layout for a single colorbar with subplots"""
  54. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  55. for ax in axs.flat:
  56. pcm = example_pcolor(ax, fontsize=24)
  57. fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
  58. @image_comparison(['constrained_layout5.png'], tol=5.e-2)
  59. def test_constrained_layout5():
  60. """
  61. Test constrained_layout for a single colorbar with subplots,
  62. colorbar bottom
  63. """
  64. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  65. for ax in axs.flat:
  66. pcm = example_pcolor(ax, fontsize=24)
  67. fig.colorbar(pcm, ax=axs,
  68. use_gridspec=False, pad=0.01, shrink=0.6,
  69. location='bottom')
  70. @image_comparison(['constrained_layout6.png'])
  71. def test_constrained_layout6():
  72. """Test constrained_layout for nested gridspecs"""
  73. fig = plt.figure(constrained_layout=True)
  74. gs = fig.add_gridspec(1, 2, figure=fig)
  75. gsl = gs[0].subgridspec(2, 2)
  76. gsr = gs[1].subgridspec(1, 2)
  77. axsl = []
  78. for gs in gsl:
  79. ax = fig.add_subplot(gs)
  80. axsl += [ax]
  81. example_plot(ax, fontsize=12)
  82. ax.set_xlabel('x-label\nMultiLine')
  83. axsr = []
  84. for gs in gsr:
  85. ax = fig.add_subplot(gs)
  86. axsr += [ax]
  87. pcm = example_pcolor(ax, fontsize=12)
  88. fig.colorbar(pcm, ax=axsr,
  89. pad=0.01, shrink=0.99, location='bottom',
  90. ticks=ticker.MaxNLocator(nbins=5))
  91. def test_constrained_layout7():
  92. """Test for proper warning if fig not set in GridSpec"""
  93. with pytest.warns(
  94. UserWarning, match=('Calling figure.constrained_layout, but figure '
  95. 'not setup to do constrained layout')):
  96. fig = plt.figure(constrained_layout=True)
  97. gs = gridspec.GridSpec(1, 2)
  98. gsl = gridspec.GridSpecFromSubplotSpec(2, 2, gs[0])
  99. gsr = gridspec.GridSpecFromSubplotSpec(1, 2, gs[1])
  100. for gs in gsl:
  101. fig.add_subplot(gs)
  102. # need to trigger a draw to get warning
  103. fig.draw(fig.canvas.get_renderer())
  104. @image_comparison(['constrained_layout8.png'])
  105. def test_constrained_layout8():
  106. """Test for gridspecs that are not completely full"""
  107. fig = plt.figure(figsize=(10, 5), constrained_layout=True)
  108. gs = gridspec.GridSpec(3, 5, figure=fig)
  109. axs = []
  110. for j in [0, 1]:
  111. if j == 0:
  112. ilist = [1]
  113. else:
  114. ilist = [0, 4]
  115. for i in ilist:
  116. ax = fig.add_subplot(gs[j, i])
  117. axs += [ax]
  118. pcm = example_pcolor(ax, fontsize=9)
  119. if i > 0:
  120. ax.set_ylabel('')
  121. if j < 1:
  122. ax.set_xlabel('')
  123. ax.set_title('')
  124. ax = fig.add_subplot(gs[2, :])
  125. axs += [ax]
  126. pcm = example_pcolor(ax, fontsize=9)
  127. fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
  128. @image_comparison(['constrained_layout9.png'])
  129. def test_constrained_layout9():
  130. """Test for handling suptitle and for sharex and sharey"""
  131. fig, axs = plt.subplots(2, 2, constrained_layout=True,
  132. sharex=False, sharey=False)
  133. for ax in axs.flat:
  134. pcm = example_pcolor(ax, fontsize=24)
  135. ax.set_xlabel('')
  136. ax.set_ylabel('')
  137. ax.set_aspect(2.)
  138. fig.colorbar(pcm, ax=axs, pad=0.01, shrink=0.6)
  139. fig.suptitle('Test Suptitle', fontsize=28)
  140. @image_comparison(['constrained_layout10.png'])
  141. def test_constrained_layout10():
  142. """Test for handling legend outside axis"""
  143. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  144. for ax in axs.flat:
  145. ax.plot(np.arange(12), label='This is a label')
  146. ax.legend(loc='center left', bbox_to_anchor=(0.8, 0.5))
  147. @image_comparison(['constrained_layout11.png'])
  148. def test_constrained_layout11():
  149. """Test for multiple nested gridspecs"""
  150. fig = plt.figure(constrained_layout=True, figsize=(13, 3))
  151. gs0 = gridspec.GridSpec(1, 2, figure=fig)
  152. gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
  153. gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1])
  154. ax = fig.add_subplot(gs0[1])
  155. example_plot(ax, fontsize=9)
  156. axs = []
  157. for gs in gsl0:
  158. ax = fig.add_subplot(gs)
  159. axs += [ax]
  160. pcm = example_pcolor(ax, fontsize=9)
  161. fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
  162. ax = fig.add_subplot(gsl[0])
  163. example_plot(ax, fontsize=9)
  164. @image_comparison(['constrained_layout11rat.png'])
  165. def test_constrained_layout11rat():
  166. """Test for multiple nested gridspecs with width_ratios"""
  167. fig = plt.figure(constrained_layout=True, figsize=(10, 3))
  168. gs0 = gridspec.GridSpec(1, 2, figure=fig, width_ratios=[6, 1])
  169. gsl = gridspec.GridSpecFromSubplotSpec(1, 2, gs0[0])
  170. gsl0 = gridspec.GridSpecFromSubplotSpec(2, 2, gsl[1], height_ratios=[2, 1])
  171. ax = fig.add_subplot(gs0[1])
  172. example_plot(ax, fontsize=9)
  173. axs = []
  174. for gs in gsl0:
  175. ax = fig.add_subplot(gs)
  176. axs += [ax]
  177. pcm = example_pcolor(ax, fontsize=9)
  178. fig.colorbar(pcm, ax=axs, shrink=0.6, aspect=70.)
  179. ax = fig.add_subplot(gsl[0])
  180. example_plot(ax, fontsize=9)
  181. @image_comparison(['constrained_layout12.png'])
  182. def test_constrained_layout12():
  183. """Test that very unbalanced labeling still works."""
  184. fig = plt.figure(constrained_layout=True)
  185. gs0 = gridspec.GridSpec(6, 2, figure=fig)
  186. ax1 = fig.add_subplot(gs0[:3, 1])
  187. ax2 = fig.add_subplot(gs0[3:, 1])
  188. example_plot(ax1, fontsize=24)
  189. example_plot(ax2, fontsize=24)
  190. ax = fig.add_subplot(gs0[0:2, 0])
  191. example_plot(ax, nodec=True)
  192. ax = fig.add_subplot(gs0[2:4, 0])
  193. example_plot(ax, nodec=True)
  194. ax = fig.add_subplot(gs0[4:, 0])
  195. example_plot(ax, nodec=True)
  196. ax.set_xlabel('x-label')
  197. @image_comparison(['constrained_layout13.png'], tol=2.e-2)
  198. def test_constrained_layout13():
  199. """Test that padding works."""
  200. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  201. for ax in axs.flat:
  202. pcm = example_pcolor(ax, fontsize=12)
  203. fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
  204. fig.set_constrained_layout_pads(w_pad=24./72., h_pad=24./72.)
  205. @image_comparison(['constrained_layout14.png'])
  206. def test_constrained_layout14():
  207. """Test that padding works."""
  208. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  209. for ax in axs.flat:
  210. pcm = example_pcolor(ax, fontsize=12)
  211. fig.colorbar(pcm, ax=ax, shrink=0.6, aspect=20., pad=0.02)
  212. fig.set_constrained_layout_pads(
  213. w_pad=3./72., h_pad=3./72.,
  214. hspace=0.2, wspace=0.2)
  215. @image_comparison(['constrained_layout15.png'])
  216. def test_constrained_layout15():
  217. """Test that rcparams work."""
  218. rcParams['figure.constrained_layout.use'] = True
  219. fig, axs = plt.subplots(2, 2)
  220. for ax in axs.flat:
  221. example_plot(ax, fontsize=12)
  222. @image_comparison(['constrained_layout16.png'])
  223. def test_constrained_layout16():
  224. """Test ax.set_position."""
  225. fig, ax = plt.subplots(constrained_layout=True)
  226. example_plot(ax, fontsize=12)
  227. ax2 = fig.add_axes([0.2, 0.2, 0.4, 0.4])
  228. @image_comparison(['constrained_layout17.png'])
  229. def test_constrained_layout17():
  230. """Test uneven gridspecs"""
  231. fig = plt.figure(constrained_layout=True)
  232. gs = gridspec.GridSpec(3, 3, figure=fig)
  233. ax1 = fig.add_subplot(gs[0, 0])
  234. ax2 = fig.add_subplot(gs[0, 1:])
  235. ax3 = fig.add_subplot(gs[1:, 0:2])
  236. ax4 = fig.add_subplot(gs[1:, -1])
  237. example_plot(ax1)
  238. example_plot(ax2)
  239. example_plot(ax3)
  240. example_plot(ax4)
  241. def test_constrained_layout18():
  242. """Test twinx"""
  243. fig, ax = plt.subplots(constrained_layout=True)
  244. ax2 = ax.twinx()
  245. example_plot(ax)
  246. example_plot(ax2, fontsize=24)
  247. fig.canvas.draw()
  248. assert all(ax.get_position().extents == ax2.get_position().extents)
  249. def test_constrained_layout19():
  250. """Test twiny"""
  251. fig, ax = plt.subplots(constrained_layout=True)
  252. ax2 = ax.twiny()
  253. example_plot(ax)
  254. example_plot(ax2, fontsize=24)
  255. ax2.set_title('')
  256. ax.set_title('')
  257. fig.canvas.draw()
  258. assert all(ax.get_position().extents == ax2.get_position().extents)
  259. def test_constrained_layout20():
  260. """Smoke test cl does not mess up added axes"""
  261. gx = np.linspace(-5, 5, 4)
  262. img = np.hypot(gx, gx[:, None])
  263. fig = plt.figure()
  264. ax = fig.add_axes([0, 0, 1, 1])
  265. mesh = ax.pcolormesh(gx, gx, img[:-1, :-1])
  266. fig.colorbar(mesh)
  267. def test_constrained_layout21():
  268. """#11035: repeated calls to suptitle should not alter the layout"""
  269. fig, ax = plt.subplots(constrained_layout=True)
  270. fig.suptitle("Suptitle0")
  271. fig.canvas.draw()
  272. extents0 = np.copy(ax.get_position().extents)
  273. fig.suptitle("Suptitle1")
  274. fig.canvas.draw()
  275. extents1 = np.copy(ax.get_position().extents)
  276. np.testing.assert_allclose(extents0, extents1)
  277. def test_constrained_layout22():
  278. """#11035: suptitle should not be include in CL if manually positioned"""
  279. fig, ax = plt.subplots(constrained_layout=True)
  280. fig.canvas.draw()
  281. extents0 = np.copy(ax.get_position().extents)
  282. fig.suptitle("Suptitle", y=0.5)
  283. fig.canvas.draw()
  284. extents1 = np.copy(ax.get_position().extents)
  285. np.testing.assert_allclose(extents0, extents1)
  286. def test_constrained_layout23():
  287. """
  288. Comment in #11035: suptitle used to cause an exception when
  289. reusing a figure w/ CL with ``clear=True``.
  290. """
  291. for i in range(2):
  292. fig, ax = plt.subplots(num="123", constrained_layout=True, clear=True)
  293. fig.suptitle("Suptitle{}".format(i))
  294. # This test occasionally fails the image comparison tests, so we mark as
  295. # flaky. Apparently the constraint solver occasionally doesn't fully
  296. # optimize. Would be nice if this were more deterministic...
  297. @pytest.mark.timeout(30)
  298. @pytest.mark.flaky(reruns=3)
  299. @image_comparison(['test_colorbar_location.png'],
  300. remove_text=True, style='mpl20')
  301. def test_colorbar_location():
  302. """
  303. Test that colorbar handling is as expected for various complicated
  304. cases...
  305. """
  306. fig, axs = plt.subplots(4, 5, constrained_layout=True)
  307. for ax in axs.flat:
  308. pcm = example_pcolor(ax)
  309. ax.set_xlabel('')
  310. ax.set_ylabel('')
  311. fig.colorbar(pcm, ax=axs[:, 1], shrink=0.4)
  312. fig.colorbar(pcm, ax=axs[-1, :2], shrink=0.5, location='bottom')
  313. fig.colorbar(pcm, ax=axs[0, 2:], shrink=0.5, location='bottom')
  314. fig.colorbar(pcm, ax=axs[-2, 3:], shrink=0.5, location='top')
  315. fig.colorbar(pcm, ax=axs[0, 0], shrink=0.5, location='left')
  316. fig.colorbar(pcm, ax=axs[1:3, 2], shrink=0.5, location='right')
  317. def test_hidden_axes():
  318. # test that if we make an axes not visible that constrained_layout
  319. # still works. Note the axes still takes space in the layout
  320. # (as does a gridspec slot that is empty)
  321. fig, axs = plt.subplots(2, 2, constrained_layout=True)
  322. axs[0, 1].set_visible(False)
  323. fig.canvas.draw()
  324. extents1 = np.copy(axs[0, 0].get_position().extents)
  325. np.testing.assert_allclose(
  326. extents1, [0.045552, 0.548288, 0.47319, 0.982638], rtol=1e-5)