import platform import numpy as np import pytest from matplotlib import cm import matplotlib.colors as mcolors import matplotlib as mpl from matplotlib import rc_context from matplotlib.testing.decorators import image_comparison import matplotlib.pyplot as plt from matplotlib.colors import ( BoundaryNorm, LogNorm, PowerNorm, Normalize, NoNorm ) from matplotlib.colorbar import Colorbar from matplotlib.ticker import FixedLocator, LogFormatter, StrMethodFormatter from matplotlib.testing.decorators import check_figures_equal def _get_cmap_norms(): """ Define a colormap and appropriate norms for each of the four possible settings of the extend keyword. Helper function for _colorbar_extension_shape and colorbar_extension_length. """ # Create a colormap and specify the levels it represents. cmap = mpl.colormaps["RdBu"].resampled(5) clevs = [-5., -2.5, -.5, .5, 1.5, 3.5] # Define norms for the colormaps. norms = dict() norms['neither'] = BoundaryNorm(clevs, len(clevs) - 1) norms['min'] = BoundaryNorm([-10] + clevs[1:], len(clevs) - 1) norms['max'] = BoundaryNorm(clevs[:-1] + [10], len(clevs) - 1) norms['both'] = BoundaryNorm([-10] + clevs[1:-1] + [10], len(clevs) - 1) return cmap, norms def _colorbar_extension_shape(spacing): """ Produce 4 colorbars with rectangular extensions for either uniform or proportional spacing. Helper function for test_colorbar_extension_shape. """ # Get a colormap and appropriate norms for each extension type. cmap, norms = _get_cmap_norms() # Create a figure and adjust whitespace for subplots. fig = plt.figure() fig.subplots_adjust(hspace=4) for i, extension_type in enumerate(('neither', 'min', 'max', 'both')): # Get the appropriate norm and use it to get colorbar boundaries. norm = norms[extension_type] boundaries = values = norm.boundaries # note that the last value was silently dropped pre 3.3: values = values[:-1] # Create a subplot. cax = fig.add_subplot(4, 1, i + 1) # Generate the colorbar. Colorbar(cax, cmap=cmap, norm=norm, boundaries=boundaries, values=values, extend=extension_type, extendrect=True, orientation='horizontal', spacing=spacing) # Turn off text and ticks. cax.tick_params(left=False, labelleft=False, bottom=False, labelbottom=False) # Return the figure to the caller. return fig def _colorbar_extension_length(spacing): """ Produce 12 colorbars with variable length extensions for either uniform or proportional spacing. Helper function for test_colorbar_extension_length. """ # Get a colormap and appropriate norms for each extension type. cmap, norms = _get_cmap_norms() # Create a figure and adjust whitespace for subplots. fig = plt.figure() fig.subplots_adjust(hspace=.6) for i, extension_type in enumerate(('neither', 'min', 'max', 'both')): # Get the appropriate norm and use it to get colorbar boundaries. norm = norms[extension_type] boundaries = values = norm.boundaries values = values[:-1] for j, extendfrac in enumerate((None, 'auto', 0.1)): # Create a subplot. cax = fig.add_subplot(12, 1, i*3 + j + 1) # Generate the colorbar. Colorbar(cax, cmap=cmap, norm=norm, boundaries=boundaries, values=values, extend=extension_type, extendfrac=extendfrac, orientation='horizontal', spacing=spacing) # Turn off text and ticks. cax.tick_params(left=False, labelleft=False, bottom=False, labelbottom=False) # Return the figure to the caller. return fig @image_comparison(['colorbar_extensions_shape_uniform.png', 'colorbar_extensions_shape_proportional.png']) def test_colorbar_extension_shape(): """Test rectangular colorbar extensions.""" # Remove this line when this test image is regenerated. plt.rcParams['pcolormesh.snap'] = False # Create figures for uniform and proportionally spaced colorbars. _colorbar_extension_shape('uniform') _colorbar_extension_shape('proportional') @image_comparison(['colorbar_extensions_uniform.png', 'colorbar_extensions_proportional.png'], tol=1.0) def test_colorbar_extension_length(): """Test variable length colorbar extensions.""" # Remove this line when this test image is regenerated. plt.rcParams['pcolormesh.snap'] = False # Create figures for uniform and proportionally spaced colorbars. _colorbar_extension_length('uniform') _colorbar_extension_length('proportional') @pytest.mark.parametrize("orientation", ["horizontal", "vertical"]) @pytest.mark.parametrize("extend,expected", [("min", (0, 0, 0, 1)), ("max", (1, 1, 1, 1)), ("both", (1, 1, 1, 1))]) def test_colorbar_extension_inverted_axis(orientation, extend, expected): """Test extension color with an inverted axis""" data = np.arange(12).reshape(3, 4) fig, ax = plt.subplots() cmap = mpl.colormaps["viridis"].with_extremes(under=(0, 0, 0, 1), over=(1, 1, 1, 1)) im = ax.imshow(data, cmap=cmap) cbar = fig.colorbar(im, orientation=orientation, extend=extend) if orientation == "horizontal": cbar.ax.invert_xaxis() else: cbar.ax.invert_yaxis() assert cbar._extend_patches[0].get_facecolor() == expected if extend == "both": assert len(cbar._extend_patches) == 2 assert cbar._extend_patches[1].get_facecolor() == (0, 0, 0, 1) else: assert len(cbar._extend_patches) == 1 @pytest.mark.parametrize('use_gridspec', [True, False]) @image_comparison(['cbar_with_orientation', 'cbar_locationing', 'double_cbar', 'cbar_sharing', ], extensions=['png'], remove_text=True, savefig_kwarg={'dpi': 40}) def test_colorbar_positioning(use_gridspec): # Remove this line when this test image is regenerated. plt.rcParams['pcolormesh.snap'] = False data = np.arange(1200).reshape(30, 40) levels = [0, 200, 400, 600, 800, 1000, 1200] # ------------------- plt.figure() plt.contourf(data, levels=levels) plt.colorbar(orientation='horizontal', use_gridspec=use_gridspec) locations = ['left', 'right', 'top', 'bottom'] plt.figure() for i, location in enumerate(locations): plt.subplot(2, 2, i + 1) plt.contourf(data, levels=levels) plt.colorbar(location=location, use_gridspec=use_gridspec) # ------------------- plt.figure() # make some other data (random integers) data_2nd = np.array([[2, 3, 2, 3], [1.5, 2, 2, 3], [2, 3, 3, 4]]) # make the random data expand to the shape of the main data data_2nd = np.repeat(np.repeat(data_2nd, 10, axis=1), 10, axis=0) color_mappable = plt.contourf(data, levels=levels, extend='both') # test extend frac here hatch_mappable = plt.contourf(data_2nd, levels=[1, 2, 3], colors='none', hatches=['/', 'o', '+'], extend='max') plt.contour(hatch_mappable, colors='black') plt.colorbar(color_mappable, location='left', label='variable 1', use_gridspec=use_gridspec) plt.colorbar(hatch_mappable, location='right', label='variable 2', use_gridspec=use_gridspec) # ------------------- plt.figure() ax1 = plt.subplot(211, anchor='NE', aspect='equal') plt.contourf(data, levels=levels) ax2 = plt.subplot(223) plt.contourf(data, levels=levels) ax3 = plt.subplot(224) plt.contourf(data, levels=levels) plt.colorbar(ax=[ax2, ax3, ax1], location='right', pad=0.0, shrink=0.5, panchor=False, use_gridspec=use_gridspec) plt.colorbar(ax=[ax2, ax3, ax1], location='left', shrink=0.5, panchor=False, use_gridspec=use_gridspec) plt.colorbar(ax=[ax1], location='bottom', panchor=False, anchor=(0.8, 0.5), shrink=0.6, use_gridspec=use_gridspec) def test_colorbar_single_ax_panchor_false(): # Note that this differs from the tests above with panchor=False because # there use_gridspec is actually ineffective: passing *ax* as lists always # disables use_gridspec. ax = plt.subplot(111, anchor='N') plt.imshow([[0, 1]]) plt.colorbar(panchor=False) assert ax.get_anchor() == 'N' @pytest.mark.parametrize('constrained', [False, True], ids=['standard', 'constrained']) def test_colorbar_single_ax_panchor_east(constrained): fig = plt.figure(constrained_layout=constrained) ax = fig.add_subplot(111, anchor='N') plt.imshow([[0, 1]]) plt.colorbar(panchor='E') assert ax.get_anchor() == 'E' @image_comparison(['contour_colorbar.png'], remove_text=True, tol=0 if platform.machine() == 'x86_64' else 0.054) def test_contour_colorbar(): fig, ax = plt.subplots(figsize=(4, 2)) data = np.arange(1200).reshape(30, 40) - 500 levels = np.array([0, 200, 400, 600, 800, 1000, 1200]) - 500 CS = ax.contour(data, levels=levels, extend='both') fig.colorbar(CS, orientation='horizontal', extend='both') fig.colorbar(CS, orientation='vertical') @image_comparison(['cbar_with_subplots_adjust.png'], remove_text=True, savefig_kwarg={'dpi': 40}) def test_gridspec_make_colorbar(): plt.figure() data = np.arange(1200).reshape(30, 40) levels = [0, 200, 400, 600, 800, 1000, 1200] plt.subplot(121) plt.contourf(data, levels=levels) plt.colorbar(use_gridspec=True, orientation='vertical') plt.subplot(122) plt.contourf(data, levels=levels) plt.colorbar(use_gridspec=True, orientation='horizontal') plt.subplots_adjust(top=0.95, right=0.95, bottom=0.2, hspace=0.25) @image_comparison(['colorbar_single_scatter.png'], remove_text=True, savefig_kwarg={'dpi': 40}) def test_colorbar_single_scatter(): # Issue #2642: if a path collection has only one entry, # the norm scaling within the colorbar must ensure a # finite range, otherwise a zero denominator will occur in _locate. plt.figure() x = y = [0] z = [50] cmap = mpl.colormaps['jet'].resampled(16) cs = plt.scatter(x, y, z, c=z, cmap=cmap) plt.colorbar(cs) @pytest.mark.parametrize('use_gridspec', [True, False]) @pytest.mark.parametrize('nested_gridspecs', [True, False]) def test_remove_from_figure(nested_gridspecs, use_gridspec): """Test `remove` with the specified ``use_gridspec`` setting.""" fig = plt.figure() if nested_gridspecs: gs = fig.add_gridspec(2, 2)[1, 1].subgridspec(2, 2) ax = fig.add_subplot(gs[1, 1]) else: ax = fig.add_subplot() sc = ax.scatter([1, 2], [3, 4]) sc.set_array(np.array([5, 6])) pre_position = ax.get_position() cb = fig.colorbar(sc, use_gridspec=use_gridspec) fig.subplots_adjust() cb.remove() fig.subplots_adjust() post_position = ax.get_position() assert (pre_position.get_points() == post_position.get_points()).all() def test_remove_from_figure_cl(): """Test `remove` with constrained_layout.""" fig, ax = plt.subplots(constrained_layout=True) sc = ax.scatter([1, 2], [3, 4]) sc.set_array(np.array([5, 6])) fig.draw_without_rendering() pre_position = ax.get_position() cb = fig.colorbar(sc) cb.remove() fig.draw_without_rendering() post_position = ax.get_position() np.testing.assert_allclose(pre_position.get_points(), post_position.get_points()) def test_colorbarbase(): # smoke test from #3805 ax = plt.gca() Colorbar(ax, cmap=plt.cm.bone) def test_parentless_mappable(): pc = mpl.collections.PatchCollection([], cmap=plt.get_cmap('viridis'), array=[]) with pytest.raises(ValueError, match='Unable to determine Axes to steal'): plt.colorbar(pc) @image_comparison(['colorbar_closed_patch.png'], remove_text=True) def test_colorbar_closed_patch(): # Remove this line when this test image is regenerated. plt.rcParams['pcolormesh.snap'] = False fig = plt.figure(figsize=(8, 6)) ax1 = fig.add_axes([0.05, 0.85, 0.9, 0.1]) ax2 = fig.add_axes([0.1, 0.65, 0.75, 0.1]) ax3 = fig.add_axes([0.05, 0.45, 0.9, 0.1]) ax4 = fig.add_axes([0.05, 0.25, 0.9, 0.1]) ax5 = fig.add_axes([0.05, 0.05, 0.9, 0.1]) cmap = mpl.colormaps["RdBu"].resampled(5) im = ax1.pcolormesh(np.linspace(0, 10, 16).reshape((4, 4)), cmap=cmap) # The use of a "values" kwarg here is unusual. It works only # because it is matched to the data range in the image and to # the number of colors in the LUT. values = np.linspace(0, 10, 5) cbar_kw = dict(orientation='horizontal', values=values, ticks=[]) # The wide line is to show that the closed path is being handled # correctly. See PR #4186. with rc_context({'axes.linewidth': 16}): plt.colorbar(im, cax=ax2, extend='both', extendfrac=0.5, **cbar_kw) plt.colorbar(im, cax=ax3, extend='both', **cbar_kw) plt.colorbar(im, cax=ax4, extend='both', extendrect=True, **cbar_kw) plt.colorbar(im, cax=ax5, extend='neither', **cbar_kw) def test_colorbar_ticks(): # test fix for #5673 fig, ax = plt.subplots() x = np.arange(-3.0, 4.001) y = np.arange(-4.0, 3.001) X, Y = np.meshgrid(x, y) Z = X * Y clevs = np.array([-12, -5, 0, 5, 12], dtype=float) colors = ['r', 'g', 'b', 'c'] cs = ax.contourf(X, Y, Z, clevs, colors=colors, extend='neither') cbar = fig.colorbar(cs, ax=ax, orientation='horizontal', ticks=clevs) assert len(cbar.ax.xaxis.get_ticklocs()) == len(clevs) def test_colorbar_minorticks_on_off(): # test for github issue #11510 and PR #11584 np.random.seed(seed=12345) data = np.random.randn(20, 20) with rc_context({'_internal.classic_mode': False}): fig, ax = plt.subplots() # purposefully setting vmin and vmax to odd fractions # so as to check for the correct locations of the minor ticks im = ax.pcolormesh(data, vmin=-2.3, vmax=3.3) cbar = fig.colorbar(im, extend='both') # testing after minorticks_on() cbar.minorticks_on() np.testing.assert_almost_equal( cbar.ax.yaxis.get_minorticklocs(), [-2.2, -1.8, -1.6, -1.4, -1.2, -0.8, -0.6, -0.4, -0.2, 0.2, 0.4, 0.6, 0.8, 1.2, 1.4, 1.6, 1.8, 2.2, 2.4, 2.6, 2.8, 3.2]) # testing after minorticks_off() cbar.minorticks_off() np.testing.assert_almost_equal(cbar.ax.yaxis.get_minorticklocs(), []) im.set_clim(vmin=-1.2, vmax=1.2) cbar.minorticks_on() np.testing.assert_almost_equal( cbar.ax.yaxis.get_minorticklocs(), [-1.2, -1.1, -0.9, -0.8, -0.7, -0.6, -0.4, -0.3, -0.2, -0.1, 0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.1, 1.2]) # tests for github issue #13257 and PR #13265 data = np.random.uniform(low=1, high=10, size=(20, 20)) fig, ax = plt.subplots() im = ax.pcolormesh(data, norm=LogNorm()) cbar = fig.colorbar(im) fig.canvas.draw() default_minorticklocks = cbar.ax.yaxis.get_minorticklocs() # test that minorticks turn off for LogNorm cbar.minorticks_off() np.testing.assert_equal(cbar.ax.yaxis.get_minorticklocs(), []) # test that minorticks turn back on for LogNorm cbar.minorticks_on() np.testing.assert_equal(cbar.ax.yaxis.get_minorticklocs(), default_minorticklocks) # test issue #13339: minorticks for LogNorm should stay off cbar.minorticks_off() cbar.set_ticks([3, 5, 7, 9]) np.testing.assert_equal(cbar.ax.yaxis.get_minorticklocs(), []) def test_cbar_minorticks_for_rc_xyminortickvisible(): """ issue gh-16468. Making sure that minor ticks on the colorbar are turned on (internally) using the cbar.minorticks_on() method when rcParams['xtick.minor.visible'] = True (for horizontal cbar) rcParams['ytick.minor.visible'] = True (for vertical cbar). Using cbar.minorticks_on() ensures that the minor ticks don't overflow into the extend regions of the colorbar. """ plt.rcParams['ytick.minor.visible'] = True plt.rcParams['xtick.minor.visible'] = True vmin, vmax = 0.4, 2.6 fig, ax = plt.subplots() im = ax.pcolormesh([[1, 2]], vmin=vmin, vmax=vmax) cbar = fig.colorbar(im, extend='both', orientation='vertical') assert cbar.ax.yaxis.get_minorticklocs()[0] >= vmin assert cbar.ax.yaxis.get_minorticklocs()[-1] <= vmax cbar = fig.colorbar(im, extend='both', orientation='horizontal') assert cbar.ax.xaxis.get_minorticklocs()[0] >= vmin assert cbar.ax.xaxis.get_minorticklocs()[-1] <= vmax def test_colorbar_autoticks(): # Test new autotick modes. Needs to be classic because # non-classic doesn't go this route. with rc_context({'_internal.classic_mode': False}): fig, ax = plt.subplots(2, 1) x = np.arange(-3.0, 4.001) y = np.arange(-4.0, 3.001) X, Y = np.meshgrid(x, y) Z = X * Y Z = Z[:-1, :-1] pcm = ax[0].pcolormesh(X, Y, Z) cbar = fig.colorbar(pcm, ax=ax[0], extend='both', orientation='vertical') pcm = ax[1].pcolormesh(X, Y, Z) cbar2 = fig.colorbar(pcm, ax=ax[1], extend='both', orientation='vertical', shrink=0.4) # note only -10 to 10 are visible, np.testing.assert_almost_equal(cbar.ax.yaxis.get_ticklocs(), np.arange(-15, 16, 5)) # note only -10 to 10 are visible np.testing.assert_almost_equal(cbar2.ax.yaxis.get_ticklocs(), np.arange(-20, 21, 10)) def test_colorbar_autotickslog(): # Test new autotick modes... with rc_context({'_internal.classic_mode': False}): fig, ax = plt.subplots(2, 1) x = np.arange(-3.0, 4.001) y = np.arange(-4.0, 3.001) X, Y = np.meshgrid(x, y) Z = X * Y Z = Z[:-1, :-1] pcm = ax[0].pcolormesh(X, Y, 10**Z, norm=LogNorm()) cbar = fig.colorbar(pcm, ax=ax[0], extend='both', orientation='vertical') pcm = ax[1].pcolormesh(X, Y, 10**Z, norm=LogNorm()) cbar2 = fig.colorbar(pcm, ax=ax[1], extend='both', orientation='vertical', shrink=0.4) # note only -12 to +12 are visible np.testing.assert_almost_equal(cbar.ax.yaxis.get_ticklocs(), 10**np.arange(-16., 16.2, 4.)) # note only -24 to +24 are visible np.testing.assert_almost_equal(cbar2.ax.yaxis.get_ticklocs(), 10**np.arange(-24., 25., 12.)) def test_colorbar_get_ticks(): # test feature for #5792 plt.figure() data = np.arange(1200).reshape(30, 40) levels = [0, 200, 400, 600, 800, 1000, 1200] plt.contourf(data, levels=levels) # testing getter for user set ticks userTicks = plt.colorbar(ticks=[0, 600, 1200]) assert userTicks.get_ticks().tolist() == [0, 600, 1200] # testing for getter after calling set_ticks userTicks.set_ticks([600, 700, 800]) assert userTicks.get_ticks().tolist() == [600, 700, 800] # testing for getter after calling set_ticks with some ticks out of bounds # removed #20054: other axes don't trim fixed lists, so colorbars # should not either: # userTicks.set_ticks([600, 1300, 1400, 1500]) # assert userTicks.get_ticks().tolist() == [600] # testing getter when no ticks are assigned defTicks = plt.colorbar(orientation='horizontal') np.testing.assert_allclose(defTicks.get_ticks().tolist(), levels) # test normal ticks and minor ticks fig, ax = plt.subplots() x = np.arange(-3.0, 4.001) y = np.arange(-4.0, 3.001) X, Y = np.meshgrid(x, y) Z = X * Y Z = Z[:-1, :-1] pcm = ax.pcolormesh(X, Y, Z) cbar = fig.colorbar(pcm, ax=ax, extend='both', orientation='vertical') ticks = cbar.get_ticks() np.testing.assert_allclose(ticks, np.arange(-15, 16, 5)) assert len(cbar.get_ticks(minor=True)) == 0 @pytest.mark.parametrize("extend", ['both', 'min', 'max']) def test_colorbar_lognorm_extension(extend): # Test that colorbar with lognorm is extended correctly f, ax = plt.subplots() cb = Colorbar(ax, norm=LogNorm(vmin=0.1, vmax=1000.0), orientation='vertical', extend=extend) assert cb._values[0] >= 0.0 def test_colorbar_powernorm_extension(): # Test that colorbar with powernorm is extended correctly # Just a smoke test that adding the colorbar doesn't raise an error or warning fig, ax = plt.subplots() Colorbar(ax, norm=PowerNorm(gamma=0.5, vmin=0.0, vmax=1.0), orientation='vertical', extend='both') def test_colorbar_axes_kw(): # test fix for #8493: This does only test, that axes-related keywords pass # and do not raise an exception. plt.figure() plt.imshow([[1, 2], [3, 4]]) plt.colorbar(orientation='horizontal', fraction=0.2, pad=0.2, shrink=0.5, aspect=10, anchor=(0., 0.), panchor=(0., 1.)) def test_colorbar_log_minortick_labels(): with rc_context({'_internal.classic_mode': False}): fig, ax = plt.subplots() pcm = ax.imshow([[10000, 50000]], norm=LogNorm()) cb = fig.colorbar(pcm) fig.canvas.draw() lb = [l.get_text() for l in cb.ax.yaxis.get_ticklabels(which='both')] expected = [r'$\mathdefault{10^{4}}$', r'$\mathdefault{2\times10^{4}}$', r'$\mathdefault{3\times10^{4}}$', r'$\mathdefault{4\times10^{4}}$'] for exp in expected: assert exp in lb def test_colorbar_renorm(): x, y = np.ogrid[-4:4:31j, -4:4:31j] z = 120000*np.exp(-x**2 - y**2) fig, ax = plt.subplots() im = ax.imshow(z) cbar = fig.colorbar(im) np.testing.assert_allclose(cbar.ax.yaxis.get_majorticklocs(), np.arange(0, 120000.1, 20000)) cbar.set_ticks([1, 2, 3]) assert isinstance(cbar.locator, FixedLocator) norm = LogNorm(z.min(), z.max()) im.set_norm(norm) np.testing.assert_allclose(cbar.ax.yaxis.get_majorticklocs(), np.logspace(-10, 7, 18)) # note that set_norm removes the FixedLocator... assert np.isclose(cbar.vmin, z.min()) cbar.set_ticks([1, 2, 3]) assert isinstance(cbar.locator, FixedLocator) np.testing.assert_allclose(cbar.ax.yaxis.get_majorticklocs(), [1.0, 2.0, 3.0]) norm = LogNorm(z.min() * 1000, z.max() * 1000) im.set_norm(norm) assert np.isclose(cbar.vmin, z.min() * 1000) assert np.isclose(cbar.vmax, z.max() * 1000) @pytest.mark.parametrize('fmt', ['%4.2e', '{x:.2e}']) def test_colorbar_format(fmt): # make sure that format is passed properly x, y = np.ogrid[-4:4:31j, -4:4:31j] z = 120000*np.exp(-x**2 - y**2) fig, ax = plt.subplots() im = ax.imshow(z) cbar = fig.colorbar(im, format=fmt) fig.canvas.draw() assert cbar.ax.yaxis.get_ticklabels()[4].get_text() == '8.00e+04' # make sure that if we change the clim of the mappable that the # formatting is *not* lost: im.set_clim([4, 200]) fig.canvas.draw() assert cbar.ax.yaxis.get_ticklabels()[4].get_text() == '2.00e+02' # but if we change the norm: im.set_norm(LogNorm(vmin=0.1, vmax=10)) fig.canvas.draw() assert (cbar.ax.yaxis.get_ticklabels()[0].get_text() == '$\\mathdefault{10^{-2}}$') def test_colorbar_scale_reset(): x, y = np.ogrid[-4:4:31j, -4:4:31j] z = 120000*np.exp(-x**2 - y**2) fig, ax = plt.subplots() pcm = ax.pcolormesh(z, cmap='RdBu_r', rasterized=True) cbar = fig.colorbar(pcm, ax=ax) cbar.outline.set_edgecolor('red') assert cbar.ax.yaxis.get_scale() == 'linear' pcm.set_norm(LogNorm(vmin=1, vmax=100)) assert cbar.ax.yaxis.get_scale() == 'log' pcm.set_norm(Normalize(vmin=-20, vmax=20)) assert cbar.ax.yaxis.get_scale() == 'linear' assert cbar.outline.get_edgecolor() == mcolors.to_rgba('red') # log scale with no vmin/vmax set should scale to the data if there # is a mappable already associated with the colorbar, not (0, 1) pcm.norm = LogNorm() assert pcm.norm.vmin == z.min() assert pcm.norm.vmax == z.max() def test_colorbar_get_ticks_2(): plt.rcParams['_internal.classic_mode'] = False fig, ax = plt.subplots() pc = ax.pcolormesh([[.05, .95]]) cb = fig.colorbar(pc) np.testing.assert_allclose(cb.get_ticks(), [0., 0.2, 0.4, 0.6, 0.8, 1.0]) def test_colorbar_inverted_ticks(): fig, axs = plt.subplots(2) ax = axs[0] pc = ax.pcolormesh(10**np.arange(1, 5).reshape(2, 2), norm=LogNorm()) cbar = fig.colorbar(pc, ax=ax, extend='both') ticks = cbar.get_ticks() cbar.ax.invert_yaxis() np.testing.assert_allclose(ticks, cbar.get_ticks()) ax = axs[1] pc = ax.pcolormesh(np.arange(1, 5).reshape(2, 2)) cbar = fig.colorbar(pc, ax=ax, extend='both') cbar.minorticks_on() ticks = cbar.get_ticks() minorticks = cbar.get_ticks(minor=True) assert isinstance(minorticks, np.ndarray) cbar.ax.invert_yaxis() np.testing.assert_allclose(ticks, cbar.get_ticks()) np.testing.assert_allclose(minorticks, cbar.get_ticks(minor=True)) def test_mappable_no_alpha(): fig, ax = plt.subplots() sm = cm.ScalarMappable(norm=mcolors.Normalize(), cmap='viridis') fig.colorbar(sm, ax=ax) sm.set_cmap('plasma') plt.draw() def test_mappable_2d_alpha(): fig, ax = plt.subplots() x = np.arange(1, 5).reshape(2, 2)/4 pc = ax.pcolormesh(x, alpha=x) cb = fig.colorbar(pc, ax=ax) # The colorbar's alpha should be None and the mappable should still have # the original alpha array assert cb.alpha is None assert pc.get_alpha() is x fig.draw_without_rendering() def test_colorbar_label(): """ Test the label parameter. It should just be mapped to the xlabel/ylabel of the axes, depending on the orientation. """ fig, ax = plt.subplots() im = ax.imshow([[1, 2], [3, 4]]) cbar = fig.colorbar(im, label='cbar') assert cbar.ax.get_ylabel() == 'cbar' cbar.set_label(None) assert cbar.ax.get_ylabel() == '' cbar.set_label('cbar 2') assert cbar.ax.get_ylabel() == 'cbar 2' cbar2 = fig.colorbar(im, label=None) assert cbar2.ax.get_ylabel() == '' cbar3 = fig.colorbar(im, orientation='horizontal', label='horizontal cbar') assert cbar3.ax.get_xlabel() == 'horizontal cbar' @image_comparison(['colorbar_keeping_xlabel.png'], style='mpl20') def test_keeping_xlabel(): # github issue #23398 - xlabels being ignored in colorbar axis arr = np.arange(25).reshape((5, 5)) fig, ax = plt.subplots() im = ax.imshow(arr) cbar = plt.colorbar(im) cbar.ax.set_xlabel('Visible Xlabel') cbar.set_label('YLabel') @pytest.mark.parametrize("clim", [(-20000, 20000), (-32768, 0)]) def test_colorbar_int(clim): # Check that we cast to float early enough to not # overflow ``int16(20000) - int16(-20000)`` or # run into ``abs(int16(-32768)) == -32768``. fig, ax = plt.subplots() im = ax.imshow([[*map(np.int16, clim)]]) fig.colorbar(im) assert (im.norm.vmin, im.norm.vmax) == clim def test_anchored_cbar_position_using_specgrid(): data = np.arange(1200).reshape(30, 40) levels = [0, 200, 400, 600, 800, 1000, 1200] shrink = 0.5 anchor_y = 0.3 # right fig, ax = plt.subplots() cs = ax.contourf(data, levels=levels) cbar = plt.colorbar( cs, ax=ax, use_gridspec=True, location='right', anchor=(1, anchor_y), shrink=shrink) # the bottom left corner of one ax is (x0, y0) # the top right corner of one ax is (x1, y1) # p0: the vertical / horizontal position of anchor x0, y0, x1, y1 = ax.get_position().extents cx0, cy0, cx1, cy1 = cbar.ax.get_position().extents p0 = (y1 - y0) * anchor_y + y0 np.testing.assert_allclose( [cy1, cy0], [y1 * shrink + (1 - shrink) * p0, p0 * (1 - shrink) + y0 * shrink]) # left fig, ax = plt.subplots() cs = ax.contourf(data, levels=levels) cbar = plt.colorbar( cs, ax=ax, use_gridspec=True, location='left', anchor=(1, anchor_y), shrink=shrink) # the bottom left corner of one ax is (x0, y0) # the top right corner of one ax is (x1, y1) # p0: the vertical / horizontal position of anchor x0, y0, x1, y1 = ax.get_position().extents cx0, cy0, cx1, cy1 = cbar.ax.get_position().extents p0 = (y1 - y0) * anchor_y + y0 np.testing.assert_allclose( [cy1, cy0], [y1 * shrink + (1 - shrink) * p0, p0 * (1 - shrink) + y0 * shrink]) # top shrink = 0.5 anchor_x = 0.3 fig, ax = plt.subplots() cs = ax.contourf(data, levels=levels) cbar = plt.colorbar( cs, ax=ax, use_gridspec=True, location='top', anchor=(anchor_x, 1), shrink=shrink) # the bottom left corner of one ax is (x0, y0) # the top right corner of one ax is (x1, y1) # p0: the vertical / horizontal position of anchor x0, y0, x1, y1 = ax.get_position().extents cx0, cy0, cx1, cy1 = cbar.ax.get_position().extents p0 = (x1 - x0) * anchor_x + x0 np.testing.assert_allclose( [cx1, cx0], [x1 * shrink + (1 - shrink) * p0, p0 * (1 - shrink) + x0 * shrink]) # bottom shrink = 0.5 anchor_x = 0.3 fig, ax = plt.subplots() cs = ax.contourf(data, levels=levels) cbar = plt.colorbar( cs, ax=ax, use_gridspec=True, location='bottom', anchor=(anchor_x, 1), shrink=shrink) # the bottom left corner of one ax is (x0, y0) # the top right corner of one ax is (x1, y1) # p0: the vertical / horizontal position of anchor x0, y0, x1, y1 = ax.get_position().extents cx0, cy0, cx1, cy1 = cbar.ax.get_position().extents p0 = (x1 - x0) * anchor_x + x0 np.testing.assert_allclose( [cx1, cx0], [x1 * shrink + (1 - shrink) * p0, p0 * (1 - shrink) + x0 * shrink]) @image_comparison(['colorbar_change_lim_scale.png'], remove_text=True, style='mpl20') def test_colorbar_change_lim_scale(): fig, ax = plt.subplots(1, 2, constrained_layout=True) pc = ax[0].pcolormesh(np.arange(100).reshape(10, 10)+1) cb = fig.colorbar(pc, ax=ax[0], extend='both') cb.ax.set_yscale('log') pc = ax[1].pcolormesh(np.arange(100).reshape(10, 10)+1) cb = fig.colorbar(pc, ax=ax[1], extend='both') cb.ax.set_ylim([20, 90]) @check_figures_equal(extensions=["png"]) def test_axes_handles_same_functions(fig_ref, fig_test): # prove that cax and cb.ax are functionally the same for nn, fig in enumerate([fig_ref, fig_test]): ax = fig.add_subplot() pc = ax.pcolormesh(np.ones(300).reshape(10, 30)) cax = fig.add_axes([0.9, 0.1, 0.03, 0.8]) cb = fig.colorbar(pc, cax=cax) if nn == 0: caxx = cax else: caxx = cb.ax caxx.set_yticks(np.arange(0, 20)) caxx.set_yscale('log') caxx.set_position([0.92, 0.1, 0.02, 0.7]) def test_inset_colorbar_layout(): fig, ax = plt.subplots(constrained_layout=True, figsize=(3, 6)) pc = ax.imshow(np.arange(100).reshape(10, 10)) cax = ax.inset_axes([1.02, 0.1, 0.03, 0.8]) cb = fig.colorbar(pc, cax=cax) fig.draw_without_rendering() # make sure this is in the figure. In the colorbar swapping # it was being dropped from the list of children... np.testing.assert_allclose(cb.ax.get_position().bounds, [0.87, 0.342, 0.0237, 0.315], atol=0.01) assert cb.ax in ax.child_axes @image_comparison(['colorbar_twoslope.png'], remove_text=True, style='mpl20') def test_twoslope_colorbar(): # Note that the second tick = 20, and should be in the middle # of the colorbar (white) # There should be no tick right at the bottom, nor at the top. fig, ax = plt.subplots() norm = mcolors.TwoSlopeNorm(20, 5, 95) pc = ax.pcolormesh(np.arange(1, 11), np.arange(1, 11), np.arange(100).reshape(10, 10), norm=norm, cmap='RdBu_r') fig.colorbar(pc) @check_figures_equal(extensions=["png"]) def test_remove_cb_whose_mappable_has_no_figure(fig_ref, fig_test): ax = fig_test.add_subplot() cb = fig_test.colorbar(cm.ScalarMappable(), cax=ax) cb.remove() def test_aspects(): fig, ax = plt.subplots(3, 2, figsize=(8, 8)) aspects = [20, 20, 10] extends = ['neither', 'both', 'both'] cb = [[None, None, None], [None, None, None]] for nn, orient in enumerate(['vertical', 'horizontal']): for mm, (aspect, extend) in enumerate(zip(aspects, extends)): pc = ax[mm, nn].pcolormesh(np.arange(100).reshape(10, 10)) cb[nn][mm] = fig.colorbar(pc, ax=ax[mm, nn], orientation=orient, aspect=aspect, extend=extend) fig.draw_without_rendering() # check the extends are right ratio: np.testing.assert_almost_equal(cb[0][1].ax.get_position().height, cb[0][0].ax.get_position().height * 0.9, decimal=2) # horizontal np.testing.assert_almost_equal(cb[1][1].ax.get_position().width, cb[1][0].ax.get_position().width * 0.9, decimal=2) # check correct aspect: pos = cb[0][0].ax.get_position(original=False) np.testing.assert_almost_equal(pos.height, pos.width * 20, decimal=2) pos = cb[1][0].ax.get_position(original=False) np.testing.assert_almost_equal(pos.height * 20, pos.width, decimal=2) # check twice as wide if aspect is 10 instead of 20 np.testing.assert_almost_equal( cb[0][0].ax.get_position(original=False).width * 2, cb[0][2].ax.get_position(original=False).width, decimal=2) np.testing.assert_almost_equal( cb[1][0].ax.get_position(original=False).height * 2, cb[1][2].ax.get_position(original=False).height, decimal=2) @image_comparison(['proportional_colorbars.png'], remove_text=True, style='mpl20') def test_proportional_colorbars(): x = y = np.arange(-3.0, 3.01, 0.025) X, Y = np.meshgrid(x, y) Z1 = np.exp(-X**2 - Y**2) Z2 = np.exp(-(X - 1)**2 - (Y - 1)**2) Z = (Z1 - Z2) * 2 levels = [-1.25, -0.5, -0.125, 0.125, 0.5, 1.25] cmap = mcolors.ListedColormap( ['0.3', '0.5', 'white', 'lightblue', 'steelblue']) cmap.set_under('darkred') cmap.set_over('crimson') norm = mcolors.BoundaryNorm(levels, cmap.N) extends = ['neither', 'both'] spacings = ['uniform', 'proportional'] fig, axs = plt.subplots(2, 2) for i in range(2): for j in range(2): CS3 = axs[i, j].contourf(X, Y, Z, levels, cmap=cmap, norm=norm, extend=extends[i]) fig.colorbar(CS3, spacing=spacings[j], ax=axs[i, j]) @image_comparison(['extend_drawedges.png'], remove_text=True, style='mpl20') def test_colorbar_extend_drawedges(): params = [ ('both', 1, [[[1.1, 0], [1.1, 1]], [[2, 0], [2, 1]], [[2.9, 0], [2.9, 1]]]), ('min', 0, [[[1.1, 0], [1.1, 1]], [[2, 0], [2, 1]]]), ('max', 0, [[[2, 0], [2, 1]], [[2.9, 0], [2.9, 1]]]), ('neither', -1, [[[2, 0], [2, 1]]]), ] plt.rcParams['axes.linewidth'] = 2 fig = plt.figure(figsize=(10, 4)) subfigs = fig.subfigures(1, 2) for orientation, subfig in zip(['horizontal', 'vertical'], subfigs): if orientation == 'horizontal': axs = subfig.subplots(4, 1) else: axs = subfig.subplots(1, 4) fig.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95) for ax, (extend, coloroffset, res) in zip(axs, params): cmap = mpl.colormaps["viridis"] bounds = np.arange(5) nb_colors = len(bounds) + coloroffset colors = cmap(np.linspace(100, 255, nb_colors).astype(int)) cmap, norm = mcolors.from_levels_and_colors(bounds, colors, extend=extend) cbar = Colorbar(ax, cmap=cmap, norm=norm, orientation=orientation, drawedges=True) # Set limits such that only two colours are visible, and the # dividers would be outside the Axes, to ensure that a) they are # not drawn outside, and b) a divider still appears between the # main colour and the extension. if orientation == 'horizontal': ax.set_xlim(1.1, 2.9) else: ax.set_ylim(1.1, 2.9) res = np.array(res)[:, :, [1, 0]] np.testing.assert_array_equal(cbar.dividers.get_segments(), res) @image_comparison(['contourf_extend_patches.png'], remove_text=True, style='mpl20') def test_colorbar_contourf_extend_patches(): params = [ ('both', 5, ['\\', '//']), ('min', 7, ['+']), ('max', 2, ['|', '-', '/', '\\', '//']), ('neither', 10, ['//', '\\', '||']), ] plt.rcParams['axes.linewidth'] = 2 fig = plt.figure(figsize=(10, 4)) subfigs = fig.subfigures(1, 2) fig.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95) x = np.linspace(-2, 3, 50) y = np.linspace(-2, 3, 30) z = np.cos(x[np.newaxis, :]) + np.sin(y[:, np.newaxis]) cmap = mpl.colormaps["viridis"] for orientation, subfig in zip(['horizontal', 'vertical'], subfigs): axs = subfig.subplots(2, 2).ravel() for ax, (extend, levels, hatches) in zip(axs, params): cs = ax.contourf(x, y, z, levels, hatches=hatches, cmap=cmap, extend=extend) subfig.colorbar(cs, ax=ax, orientation=orientation, fraction=0.4, extendfrac=0.2, aspect=5) def test_negative_boundarynorm(): fig, ax = plt.subplots(figsize=(1, 3)) cmap = mpl.colormaps["viridis"] clevs = np.arange(-94, -85) norm = BoundaryNorm(clevs, cmap.N) cb = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=ax) np.testing.assert_allclose(cb.ax.get_ylim(), [clevs[0], clevs[-1]]) np.testing.assert_allclose(cb.ax.get_yticks(), clevs) clevs = np.arange(85, 94) norm = BoundaryNorm(clevs, cmap.N) cb = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=ax) np.testing.assert_allclose(cb.ax.get_ylim(), [clevs[0], clevs[-1]]) np.testing.assert_allclose(cb.ax.get_yticks(), clevs) clevs = np.arange(-3, 3) norm = BoundaryNorm(clevs, cmap.N) cb = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=ax) np.testing.assert_allclose(cb.ax.get_ylim(), [clevs[0], clevs[-1]]) np.testing.assert_allclose(cb.ax.get_yticks(), clevs) clevs = np.arange(-8, 1) norm = BoundaryNorm(clevs, cmap.N) cb = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=ax) np.testing.assert_allclose(cb.ax.get_ylim(), [clevs[0], clevs[-1]]) np.testing.assert_allclose(cb.ax.get_yticks(), clevs) def test_centerednorm(): # Test default centered norm gets expanded with non-singular limits # when plot data is all equal (autoscale halfrange == 0) fig, ax = plt.subplots(figsize=(1, 3)) norm = mcolors.CenteredNorm() mappable = ax.pcolormesh(np.zeros((3, 3)), norm=norm) fig.colorbar(mappable) assert (norm.vmin, norm.vmax) == (-0.1, 0.1) @image_comparison(['nonorm_colorbars.svg'], style='mpl20') def test_nonorm(): plt.rcParams['svg.fonttype'] = 'none' data = [1, 2, 3, 4, 5] fig, ax = plt.subplots(figsize=(6, 1)) fig.subplots_adjust(bottom=0.5) norm = NoNorm(vmin=min(data), vmax=max(data)) cmap = mpl.colormaps["viridis"].resampled(len(data)) mappable = cm.ScalarMappable(norm=norm, cmap=cmap) cbar = fig.colorbar(mappable, cax=ax, orientation="horizontal") @image_comparison(['test_boundaries.png'], remove_text=True, style='mpl20') def test_boundaries(): np.random.seed(seed=19680808) fig, ax = plt.subplots(figsize=(2, 2)) pc = ax.pcolormesh(np.random.randn(10, 10), cmap='RdBu_r') cb = fig.colorbar(pc, ax=ax, boundaries=np.linspace(-3, 3, 7)) def test_colorbar_no_warning_rcparams_grid_true(): # github issue #21723 - If mpl style has 'axes.grid' = True, # fig.colorbar raises a warning about Auto-removal of grids # by pcolor() and pcolormesh(). This is fixed by PR #22216. plt.rcParams['axes.grid'] = True fig, ax = plt.subplots() ax.grid(False) im = ax.pcolormesh([0, 1], [0, 1], [[1]]) # make sure that no warning is raised by fig.colorbar fig.colorbar(im) def test_colorbar_set_formatter_locator(): # check that the locator properties echo what is on the axis: fig, ax = plt.subplots() pc = ax.pcolormesh(np.random.randn(10, 10)) cb = fig.colorbar(pc) cb.ax.yaxis.set_major_locator(FixedLocator(np.arange(10))) cb.ax.yaxis.set_minor_locator(FixedLocator(np.arange(0, 10, 0.2))) assert cb.locator is cb.ax.yaxis.get_major_locator() assert cb.minorlocator is cb.ax.yaxis.get_minor_locator() cb.ax.yaxis.set_major_formatter(LogFormatter()) cb.ax.yaxis.set_minor_formatter(LogFormatter()) assert cb.formatter is cb.ax.yaxis.get_major_formatter() assert cb.minorformatter is cb.ax.yaxis.get_minor_formatter() # check that the setter works as expected: loc = FixedLocator(np.arange(7)) cb.locator = loc assert cb.ax.yaxis.get_major_locator() is loc loc = FixedLocator(np.arange(0, 7, 0.1)) cb.minorlocator = loc assert cb.ax.yaxis.get_minor_locator() is loc fmt = LogFormatter() cb.formatter = fmt assert cb.ax.yaxis.get_major_formatter() is fmt fmt = LogFormatter() cb.minorformatter = fmt assert cb.ax.yaxis.get_minor_formatter() is fmt assert cb.long_axis is cb.ax.yaxis @image_comparison(['colorbar_extend_alpha.png'], remove_text=True, savefig_kwarg={'dpi': 40}) def test_colorbar_extend_alpha(): fig, ax = plt.subplots() im = ax.imshow([[0, 1], [2, 3]], alpha=0.3, interpolation="none") fig.colorbar(im, extend='both', boundaries=[0.5, 1.5, 2.5]) def test_offset_text_loc(): plt.style.use('mpl20') fig, ax = plt.subplots() np.random.seed(seed=19680808) pc = ax.pcolormesh(np.random.randn(10, 10)*1e6) cb = fig.colorbar(pc, location='right', extend='max') fig.draw_without_rendering() # check that the offsetText is in the proper place above the # colorbar axes. In this case the colorbar axes is the same # height as the parent, so use the parents bbox. assert cb.ax.yaxis.offsetText.get_position()[1] > ax.bbox.y1 def test_title_text_loc(): plt.style.use('mpl20') fig, ax = plt.subplots() np.random.seed(seed=19680808) pc = ax.pcolormesh(np.random.randn(10, 10)) cb = fig.colorbar(pc, location='right', extend='max') cb.ax.set_title('Aardvark') fig.draw_without_rendering() # check that the title is in the proper place above the # colorbar axes, including its extend triangles.... assert (cb.ax.title.get_window_extent(fig.canvas.get_renderer()).ymax > cb.ax.spines['outline'].get_window_extent().ymax) @check_figures_equal(extensions=["png"]) def test_passing_location(fig_ref, fig_test): ax_ref = fig_ref.add_subplot() im = ax_ref.imshow([[0, 1], [2, 3]]) ax_ref.get_figure().colorbar(im, cax=ax_ref.inset_axes([0, 1.05, 1, 0.05]), orientation="horizontal", ticklocation="top") ax_test = fig_test.add_subplot() im = ax_test.imshow([[0, 1], [2, 3]]) ax_test.get_figure().colorbar(im, cax=ax_test.inset_axes([0, 1.05, 1, 0.05]), location="top") @pytest.mark.parametrize("kwargs,error,message", [ ({'location': 'top', 'orientation': 'vertical'}, TypeError, "location and orientation are mutually exclusive"), ({'location': 'top', 'orientation': 'vertical', 'cax': True}, TypeError, "location and orientation are mutually exclusive"), # Different to above ({'ticklocation': 'top', 'orientation': 'vertical', 'cax': True}, ValueError, "'top' is not a valid value for position"), ({'location': 'top', 'extendfrac': (0, None)}, ValueError, "invalid value for extendfrac"), ]) def test_colorbar_errors(kwargs, error, message): fig, ax = plt.subplots() im = ax.imshow([[0, 1], [2, 3]]) if kwargs.get('cax', None) is True: kwargs['cax'] = ax.inset_axes([0, 1.05, 1, 0.05]) with pytest.raises(error, match=message): fig.colorbar(im, **kwargs) def test_colorbar_axes_parmeters(): fig, ax = plt.subplots(2) im = ax[0].imshow([[0, 1], [2, 3]]) # colorbar should accept any form of axes sequence: fig.colorbar(im, ax=ax) fig.colorbar(im, ax=ax[0]) fig.colorbar(im, ax=[_ax for _ax in ax]) fig.colorbar(im, ax=(ax[0], ax[1])) fig.colorbar(im, ax={i: _ax for i, _ax in enumerate(ax)}.values()) fig.draw_without_rendering() def test_colorbar_wrong_figure(): # If we decide in the future to disallow calling colorbar() on the "wrong" figure, # just delete this test. fig_tl = plt.figure(layout="tight") fig_cl = plt.figure(layout="constrained") im = fig_cl.add_subplot().imshow([[0, 1]]) # Make sure this doesn't try to setup a gridspec-controlled colorbar on fig_cl, # which would crash CL. with pytest.warns(UserWarning, match="different Figure"): fig_tl.colorbar(im) fig_tl.draw_without_rendering() fig_cl.draw_without_rendering() def test_colorbar_format_string_and_old(): plt.imshow([[0, 1]]) cb = plt.colorbar(format="{x}%") assert isinstance(cb._formatter, StrMethodFormatter)