64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
|
from contextlib import ExitStack
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import matplotlib as mpl
|
||
|
|
from matplotlib.testing.decorators import image_comparison
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
from mpl_toolkits.axes_grid1 import ImageGrid
|
||
|
|
|
||
|
|
|
||
|
|
# The original version of this test relied on mpl_toolkits's slightly different
|
||
|
|
# colorbar implementation; moving to matplotlib's own colorbar implementation
|
||
|
|
# caused the small image comparison error.
|
||
|
|
@pytest.mark.parametrize("legacy_colorbar", [False, True])
|
||
|
|
@image_comparison(['imagegrid_cbar_mode.png'],
|
||
|
|
remove_text=True, style='mpl20', tol=0.3)
|
||
|
|
def test_imagegrid_cbar_mode_edge(legacy_colorbar):
|
||
|
|
mpl.rcParams["mpl_toolkits.legacy_colorbar"] = legacy_colorbar
|
||
|
|
|
||
|
|
X, Y = np.meshgrid(np.linspace(0, 6, 30), np.linspace(0, 6, 30))
|
||
|
|
arr = np.sin(X) * np.cos(Y) + 1j*(np.sin(3*Y) * np.cos(Y/2.))
|
||
|
|
|
||
|
|
fig = plt.figure(figsize=(18, 9))
|
||
|
|
|
||
|
|
positions = (241, 242, 243, 244, 245, 246, 247, 248)
|
||
|
|
directions = ['row']*4 + ['column']*4
|
||
|
|
cbar_locations = ['left', 'right', 'top', 'bottom']*2
|
||
|
|
|
||
|
|
for position, direction, location in zip(
|
||
|
|
positions, directions, cbar_locations):
|
||
|
|
grid = ImageGrid(fig, position,
|
||
|
|
nrows_ncols=(2, 2),
|
||
|
|
direction=direction,
|
||
|
|
cbar_location=location,
|
||
|
|
cbar_size='20%',
|
||
|
|
cbar_mode='edge')
|
||
|
|
ax1, ax2, ax3, ax4, = grid
|
||
|
|
|
||
|
|
ax1.imshow(arr.real, cmap='nipy_spectral')
|
||
|
|
ax2.imshow(arr.imag, cmap='hot')
|
||
|
|
ax3.imshow(np.abs(arr), cmap='jet')
|
||
|
|
ax4.imshow(np.arctan2(arr.imag, arr.real), cmap='hsv')
|
||
|
|
|
||
|
|
with (pytest.warns(mpl.MatplotlibDeprecationWarning) if legacy_colorbar
|
||
|
|
else ExitStack()):
|
||
|
|
# In each row/column, the "first" colorbars must be overwritten by
|
||
|
|
# the "second" ones. To achieve this, clear out the axes first.
|
||
|
|
for ax in grid:
|
||
|
|
ax.cax.cla()
|
||
|
|
cb = ax.cax.colorbar(
|
||
|
|
ax.images[0],
|
||
|
|
ticks=mpl.ticker.MaxNLocator(5)) # old default locator.
|
||
|
|
|
||
|
|
|
||
|
|
def test_imagegrid():
|
||
|
|
mpl.rcParams["mpl_toolkits.legacy_colorbar"] = False
|
||
|
|
fig = plt.figure()
|
||
|
|
grid = ImageGrid(fig, 111, nrows_ncols=(1, 1))
|
||
|
|
ax = grid[0]
|
||
|
|
im = ax.imshow([[1, 2]], norm=mpl.colors.LogNorm())
|
||
|
|
cb = ax.cax.colorbar(im)
|
||
|
|
assert isinstance(cb.locator, mpl.colorbar._ColorbarLogLocator)
|