ai-content-maker/.venv/Lib/site-packages/mpl_toolkits/mplot3d/tests/test_legend3d.py

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()