118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
|
import numpy as np
|
||
|
|
|
||
|
|
import matplotlib as mpl
|
||
|
|
from matplotlib.colors import same_color
|
||
|
|
from matplotlib.testing.decorators import image_comparison
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
from mpl_toolkits.mplot3d import art3d
|
||
|
|
|
||
|
|
|
||
|
|
# Update style when regenerating the test image
|
||
|
|
@image_comparison(['legend_plot.png'], remove_text=True, style=('mpl20'))
|
||
|
|
def test_legend_plot():
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
x = np.arange(10)
|
||
|
|
ax.plot(x, 5 - x, 'o', zdir='y', label='z=1')
|
||
|
|
ax.plot(x, x - 5, 'o', zdir='y', label='z=-1')
|
||
|
|
ax.legend()
|
||
|
|
|
||
|
|
|
||
|
|
# Update style when regenerating the test image
|
||
|
|
@image_comparison(['legend_bar.png'], remove_text=True, style=('mpl20'))
|
||
|
|
def test_legend_bar():
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
x = np.arange(10)
|
||
|
|
b1 = ax.bar(x, x, zdir='y', align='edge', color='m')
|
||
|
|
b2 = ax.bar(x, x[::-1], zdir='x', align='edge', color='g')
|
||
|
|
ax.legend([b1[0], b2[0]], ['up', 'down'])
|
||
|
|
|
||
|
|
|
||
|
|
# Update style when regenerating the test image
|
||
|
|
@image_comparison(['fancy.png'], remove_text=True, style=('mpl20'))
|
||
|
|
def test_fancy():
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
ax.plot(np.arange(10), np.full(10, 5), np.full(10, 5), 'o--', label='line')
|
||
|
|
ax.scatter(np.arange(10), np.arange(10, 0, -1), label='scatter')
|
||
|
|
ax.errorbar(np.full(10, 5), np.arange(10), np.full(10, 10),
|
||
|
|
xerr=0.5, zerr=0.5, label='errorbar')
|
||
|
|
ax.legend(loc='lower left', ncols=2, title='My legend', numpoints=1)
|
||
|
|
|
||
|
|
|
||
|
|
def test_linecollection_scaled_dashes():
|
||
|
|
lines1 = [[(0, .5), (.5, 1)], [(.3, .6), (.2, .2)]]
|
||
|
|
lines2 = [[[0.7, .2], [.8, .4]], [[.5, .7], [.6, .1]]]
|
||
|
|
lines3 = [[[0.6, .2], [.8, .4]], [[.5, .7], [.1, .1]]]
|
||
|
|
lc1 = art3d.Line3DCollection(lines1, linestyles="--", lw=3)
|
||
|
|
lc2 = art3d.Line3DCollection(lines2, linestyles="-.")
|
||
|
|
lc3 = art3d.Line3DCollection(lines3, linestyles=":", lw=.5)
|
||
|
|
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
ax.add_collection(lc1)
|
||
|
|
ax.add_collection(lc2)
|
||
|
|
ax.add_collection(lc3)
|
||
|
|
|
||
|
|
leg = ax.legend([lc1, lc2, lc3], ['line1', 'line2', 'line 3'])
|
||
|
|
h1, h2, h3 = leg.legend_handles
|
||
|
|
|
||
|
|
for oh, lh in zip((lc1, lc2, lc3), (h1, h2, h3)):
|
||
|
|
assert oh.get_linestyles()[0] == lh._dash_pattern
|
||
|
|
|
||
|
|
|
||
|
|
def test_handlerline3d():
|
||
|
|
# Test marker consistency for monolithic Line3D legend handler.
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
ax.scatter([0, 1], [0, 1], marker="v")
|
||
|
|
handles = [art3d.Line3D([0], [0], [0], marker="v")]
|
||
|
|
leg = ax.legend(handles, ["Aardvark"], numpoints=1)
|
||
|
|
assert handles[0].get_marker() == leg.legend_handles[0].get_marker()
|
||
|
|
|
||
|
|
|
||
|
|
def test_contour_legend_elements():
|
||
|
|
x, y = np.mgrid[1:10, 1:10]
|
||
|
|
h = x * y
|
||
|
|
colors = ['blue', '#00FF00', 'red']
|
||
|
|
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
cs = ax.contour(x, y, h, levels=[10, 30, 50], colors=colors, extend='both')
|
||
|
|
|
||
|
|
artists, labels = cs.legend_elements()
|
||
|
|
assert labels == ['$x = 10.0$', '$x = 30.0$', '$x = 50.0$']
|
||
|
|
assert all(isinstance(a, mpl.lines.Line2D) for a in artists)
|
||
|
|
assert all(same_color(a.get_color(), c)
|
||
|
|
for a, c in zip(artists, colors))
|
||
|
|
|
||
|
|
|
||
|
|
def test_contourf_legend_elements():
|
||
|
|
x, y = np.mgrid[1:10, 1:10]
|
||
|
|
h = x * y
|
||
|
|
|
||
|
|
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
|
||
|
|
cs = ax.contourf(x, y, h, levels=[10, 30, 50],
|
||
|
|
colors=['#FFFF00', '#FF00FF', '#00FFFF'],
|
||
|
|
extend='both')
|
||
|
|
cs.cmap.set_over('red')
|
||
|
|
cs.cmap.set_under('blue')
|
||
|
|
cs.changed()
|
||
|
|
artists, labels = cs.legend_elements()
|
||
|
|
assert labels == ['$x \\leq -1e+250s$',
|
||
|
|
'$10.0 < x \\leq 30.0$',
|
||
|
|
'$30.0 < x \\leq 50.0$',
|
||
|
|
'$x > 1e+250s$']
|
||
|
|
expected_colors = ('blue', '#FFFF00', '#FF00FF', 'red')
|
||
|
|
assert all(isinstance(a, mpl.patches.Rectangle) for a in artists)
|
||
|
|
assert all(same_color(a.get_facecolor(), c)
|
||
|
|
for a, c in zip(artists, expected_colors))
|
||
|
|
|
||
|
|
|
||
|
|
def test_legend_Poly3dCollection():
|
||
|
|
|
||
|
|
verts = np.asarray([[0, 0, 0], [0, 1, 1], [1, 0, 1]])
|
||
|
|
mesh = art3d.Poly3DCollection([verts], label="surface")
|
||
|
|
|
||
|
|
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
|
||
|
|
mesh.set_edgecolor('k')
|
||
|
|
handle = ax.add_collection3d(mesh)
|
||
|
|
leg = ax.legend()
|
||
|
|
assert (leg.legend_handles[0].get_facecolor()
|
||
|
|
== handle.get_facecolor()).all()
|