1339 lines
46 KiB
Python
1339 lines
46 KiB
Python
|
from datetime import datetime
|
||
|
import io
|
||
|
import itertools
|
||
|
import platform
|
||
|
import re
|
||
|
from types import SimpleNamespace
|
||
|
|
||
|
import numpy as np
|
||
|
from numpy.testing import assert_array_equal, assert_array_almost_equal
|
||
|
import pytest
|
||
|
|
||
|
import matplotlib as mpl
|
||
|
import matplotlib.pyplot as plt
|
||
|
import matplotlib.collections as mcollections
|
||
|
import matplotlib.colors as mcolors
|
||
|
import matplotlib.path as mpath
|
||
|
import matplotlib.transforms as mtransforms
|
||
|
from matplotlib.collections import (Collection, LineCollection,
|
||
|
EventCollection, PolyCollection)
|
||
|
from matplotlib.testing.decorators import check_figures_equal, image_comparison
|
||
|
|
||
|
|
||
|
@pytest.fixture(params=["pcolormesh", "pcolor"])
|
||
|
def pcfunc(request):
|
||
|
return request.param
|
||
|
|
||
|
|
||
|
def generate_EventCollection_plot():
|
||
|
"""Generate the initial collection and plot it."""
|
||
|
positions = np.array([0., 1., 2., 3., 5., 8., 13., 21.])
|
||
|
extra_positions = np.array([34., 55., 89.])
|
||
|
orientation = 'horizontal'
|
||
|
lineoffset = 1
|
||
|
linelength = .5
|
||
|
linewidth = 2
|
||
|
color = [1, 0, 0, 1]
|
||
|
linestyle = 'solid'
|
||
|
antialiased = True
|
||
|
|
||
|
coll = EventCollection(positions,
|
||
|
orientation=orientation,
|
||
|
lineoffset=lineoffset,
|
||
|
linelength=linelength,
|
||
|
linewidth=linewidth,
|
||
|
color=color,
|
||
|
linestyle=linestyle,
|
||
|
antialiased=antialiased
|
||
|
)
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
ax.add_collection(coll)
|
||
|
ax.set_title('EventCollection: default')
|
||
|
props = {'positions': positions,
|
||
|
'extra_positions': extra_positions,
|
||
|
'orientation': orientation,
|
||
|
'lineoffset': lineoffset,
|
||
|
'linelength': linelength,
|
||
|
'linewidth': linewidth,
|
||
|
'color': color,
|
||
|
'linestyle': linestyle,
|
||
|
'antialiased': antialiased
|
||
|
}
|
||
|
ax.set_xlim(-1, 22)
|
||
|
ax.set_ylim(0, 2)
|
||
|
return ax, coll, props
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__default'])
|
||
|
def test__EventCollection__get_props():
|
||
|
_, coll, props = generate_EventCollection_plot()
|
||
|
# check that the default segments have the correct coordinates
|
||
|
check_segments(coll,
|
||
|
props['positions'],
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
# check that the default positions match the input positions
|
||
|
np.testing.assert_array_equal(props['positions'], coll.get_positions())
|
||
|
# check that the default orientation matches the input orientation
|
||
|
assert props['orientation'] == coll.get_orientation()
|
||
|
# check that the default orientation matches the input orientation
|
||
|
assert coll.is_horizontal()
|
||
|
# check that the default linelength matches the input linelength
|
||
|
assert props['linelength'] == coll.get_linelength()
|
||
|
# check that the default lineoffset matches the input lineoffset
|
||
|
assert props['lineoffset'] == coll.get_lineoffset()
|
||
|
# check that the default linestyle matches the input linestyle
|
||
|
assert coll.get_linestyle() == [(0, None)]
|
||
|
# check that the default color matches the input color
|
||
|
for color in [coll.get_color(), *coll.get_colors()]:
|
||
|
np.testing.assert_array_equal(color, props['color'])
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__set_positions'])
|
||
|
def test__EventCollection__set_positions():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_positions = np.hstack([props['positions'], props['extra_positions']])
|
||
|
coll.set_positions(new_positions)
|
||
|
np.testing.assert_array_equal(new_positions, coll.get_positions())
|
||
|
check_segments(coll, new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: set_positions')
|
||
|
splt.set_xlim(-1, 90)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__add_positions'])
|
||
|
def test__EventCollection__add_positions():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_positions = np.hstack([props['positions'],
|
||
|
props['extra_positions'][0]])
|
||
|
coll.switch_orientation() # Test adding in the vertical orientation, too.
|
||
|
coll.add_positions(props['extra_positions'][0])
|
||
|
coll.switch_orientation()
|
||
|
np.testing.assert_array_equal(new_positions, coll.get_positions())
|
||
|
check_segments(coll,
|
||
|
new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: add_positions')
|
||
|
splt.set_xlim(-1, 35)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__append_positions'])
|
||
|
def test__EventCollection__append_positions():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_positions = np.hstack([props['positions'],
|
||
|
props['extra_positions'][2]])
|
||
|
coll.append_positions(props['extra_positions'][2])
|
||
|
np.testing.assert_array_equal(new_positions, coll.get_positions())
|
||
|
check_segments(coll,
|
||
|
new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: append_positions')
|
||
|
splt.set_xlim(-1, 90)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__extend_positions'])
|
||
|
def test__EventCollection__extend_positions():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_positions = np.hstack([props['positions'],
|
||
|
props['extra_positions'][1:]])
|
||
|
coll.extend_positions(props['extra_positions'][1:])
|
||
|
np.testing.assert_array_equal(new_positions, coll.get_positions())
|
||
|
check_segments(coll,
|
||
|
new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: extend_positions')
|
||
|
splt.set_xlim(-1, 90)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__switch_orientation'])
|
||
|
def test__EventCollection__switch_orientation():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_orientation = 'vertical'
|
||
|
coll.switch_orientation()
|
||
|
assert new_orientation == coll.get_orientation()
|
||
|
assert not coll.is_horizontal()
|
||
|
new_positions = coll.get_positions()
|
||
|
check_segments(coll,
|
||
|
new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'], new_orientation)
|
||
|
splt.set_title('EventCollection: switch_orientation')
|
||
|
splt.set_ylim(-1, 22)
|
||
|
splt.set_xlim(0, 2)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__switch_orientation__2x'])
|
||
|
def test__EventCollection__switch_orientation_2x():
|
||
|
"""
|
||
|
Check that calling switch_orientation twice sets the orientation back to
|
||
|
the default.
|
||
|
"""
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
coll.switch_orientation()
|
||
|
coll.switch_orientation()
|
||
|
new_positions = coll.get_positions()
|
||
|
assert props['orientation'] == coll.get_orientation()
|
||
|
assert coll.is_horizontal()
|
||
|
np.testing.assert_array_equal(props['positions'], new_positions)
|
||
|
check_segments(coll,
|
||
|
new_positions,
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: switch_orientation 2x')
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__set_orientation'])
|
||
|
def test__EventCollection__set_orientation():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_orientation = 'vertical'
|
||
|
coll.set_orientation(new_orientation)
|
||
|
assert new_orientation == coll.get_orientation()
|
||
|
assert not coll.is_horizontal()
|
||
|
check_segments(coll,
|
||
|
props['positions'],
|
||
|
props['linelength'],
|
||
|
props['lineoffset'],
|
||
|
new_orientation)
|
||
|
splt.set_title('EventCollection: set_orientation')
|
||
|
splt.set_ylim(-1, 22)
|
||
|
splt.set_xlim(0, 2)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__set_linelength'])
|
||
|
def test__EventCollection__set_linelength():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_linelength = 15
|
||
|
coll.set_linelength(new_linelength)
|
||
|
assert new_linelength == coll.get_linelength()
|
||
|
check_segments(coll,
|
||
|
props['positions'],
|
||
|
new_linelength,
|
||
|
props['lineoffset'],
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: set_linelength')
|
||
|
splt.set_ylim(-20, 20)
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__set_lineoffset'])
|
||
|
def test__EventCollection__set_lineoffset():
|
||
|
splt, coll, props = generate_EventCollection_plot()
|
||
|
new_lineoffset = -5.
|
||
|
coll.set_lineoffset(new_lineoffset)
|
||
|
assert new_lineoffset == coll.get_lineoffset()
|
||
|
check_segments(coll,
|
||
|
props['positions'],
|
||
|
props['linelength'],
|
||
|
new_lineoffset,
|
||
|
props['orientation'])
|
||
|
splt.set_title('EventCollection: set_lineoffset')
|
||
|
splt.set_ylim(-6, -4)
|
||
|
|
||
|
|
||
|
@image_comparison([
|
||
|
'EventCollection_plot__set_linestyle',
|
||
|
'EventCollection_plot__set_linestyle',
|
||
|
'EventCollection_plot__set_linewidth',
|
||
|
])
|
||
|
def test__EventCollection__set_prop():
|
||
|
for prop, value, expected in [
|
||
|
('linestyle', 'dashed', [(0, (6.0, 6.0))]),
|
||
|
('linestyle', (0, (6., 6.)), [(0, (6.0, 6.0))]),
|
||
|
('linewidth', 5, 5),
|
||
|
]:
|
||
|
splt, coll, _ = generate_EventCollection_plot()
|
||
|
coll.set(**{prop: value})
|
||
|
assert plt.getp(coll, prop) == expected
|
||
|
splt.set_title(f'EventCollection: set_{prop}')
|
||
|
|
||
|
|
||
|
@image_comparison(['EventCollection_plot__set_color'])
|
||
|
def test__EventCollection__set_color():
|
||
|
splt, coll, _ = generate_EventCollection_plot()
|
||
|
new_color = np.array([0, 1, 1, 1])
|
||
|
coll.set_color(new_color)
|
||
|
for color in [coll.get_color(), *coll.get_colors()]:
|
||
|
np.testing.assert_array_equal(color, new_color)
|
||
|
splt.set_title('EventCollection: set_color')
|
||
|
|
||
|
|
||
|
def check_segments(coll, positions, linelength, lineoffset, orientation):
|
||
|
"""
|
||
|
Test helper checking that all values in the segment are correct, given a
|
||
|
particular set of inputs.
|
||
|
"""
|
||
|
segments = coll.get_segments()
|
||
|
if (orientation.lower() == 'horizontal'
|
||
|
or orientation.lower() == 'none' or orientation is None):
|
||
|
# if horizontal, the position in is in the y-axis
|
||
|
pos1 = 1
|
||
|
pos2 = 0
|
||
|
elif orientation.lower() == 'vertical':
|
||
|
# if vertical, the position in is in the x-axis
|
||
|
pos1 = 0
|
||
|
pos2 = 1
|
||
|
else:
|
||
|
raise ValueError("orientation must be 'horizontal' or 'vertical'")
|
||
|
|
||
|
# test to make sure each segment is correct
|
||
|
for i, segment in enumerate(segments):
|
||
|
assert segment[0, pos1] == lineoffset + linelength / 2
|
||
|
assert segment[1, pos1] == lineoffset - linelength / 2
|
||
|
assert segment[0, pos2] == positions[i]
|
||
|
assert segment[1, pos2] == positions[i]
|
||
|
|
||
|
|
||
|
def test_collection_norm_autoscale():
|
||
|
# norm should be autoscaled when array is set, not deferred to draw time
|
||
|
lines = np.arange(24).reshape((4, 3, 2))
|
||
|
coll = mcollections.LineCollection(lines, array=np.arange(4))
|
||
|
assert coll.norm(2) == 2 / 3
|
||
|
# setting a new array shouldn't update the already scaled limits
|
||
|
coll.set_array(np.arange(4) + 5)
|
||
|
assert coll.norm(2) == 2 / 3
|
||
|
|
||
|
|
||
|
def test_null_collection_datalim():
|
||
|
col = mcollections.PathCollection([])
|
||
|
col_data_lim = col.get_datalim(mtransforms.IdentityTransform())
|
||
|
assert_array_equal(col_data_lim.get_points(),
|
||
|
mtransforms.Bbox.null().get_points())
|
||
|
|
||
|
|
||
|
def test_no_offsets_datalim():
|
||
|
# A collection with no offsets and a non transData
|
||
|
# transform should return a null bbox
|
||
|
ax = plt.axes()
|
||
|
coll = mcollections.PathCollection([mpath.Path([(0, 0), (1, 0)])])
|
||
|
ax.add_collection(coll)
|
||
|
coll_data_lim = coll.get_datalim(mtransforms.IdentityTransform())
|
||
|
assert_array_equal(coll_data_lim.get_points(),
|
||
|
mtransforms.Bbox.null().get_points())
|
||
|
|
||
|
|
||
|
def test_add_collection():
|
||
|
# Test if data limits are unchanged by adding an empty collection.
|
||
|
# GitHub issue #1490, pull #1497.
|
||
|
plt.figure()
|
||
|
ax = plt.axes()
|
||
|
ax.scatter([0, 1], [0, 1])
|
||
|
bounds = ax.dataLim.bounds
|
||
|
ax.scatter([], [])
|
||
|
assert ax.dataLim.bounds == bounds
|
||
|
|
||
|
|
||
|
@mpl.style.context('mpl20')
|
||
|
@check_figures_equal(extensions=['png'])
|
||
|
def test_collection_log_datalim(fig_test, fig_ref):
|
||
|
# Data limits should respect the minimum x/y when using log scale.
|
||
|
x_vals = [4.38462e-6, 5.54929e-6, 7.02332e-6, 8.88889e-6, 1.12500e-5,
|
||
|
1.42383e-5, 1.80203e-5, 2.28070e-5, 2.88651e-5, 3.65324e-5,
|
||
|
4.62363e-5, 5.85178e-5, 7.40616e-5, 9.37342e-5, 1.18632e-4]
|
||
|
y_vals = [0.0, 0.1, 0.182, 0.332, 0.604, 1.1, 2.0, 3.64, 6.64, 12.1, 22.0,
|
||
|
39.6, 71.3]
|
||
|
|
||
|
x, y = np.meshgrid(x_vals, y_vals)
|
||
|
x = x.flatten()
|
||
|
y = y.flatten()
|
||
|
|
||
|
ax_test = fig_test.subplots()
|
||
|
ax_test.set_xscale('log')
|
||
|
ax_test.set_yscale('log')
|
||
|
ax_test.margins = 0
|
||
|
ax_test.scatter(x, y)
|
||
|
|
||
|
ax_ref = fig_ref.subplots()
|
||
|
ax_ref.set_xscale('log')
|
||
|
ax_ref.set_yscale('log')
|
||
|
ax_ref.plot(x, y, marker="o", ls="")
|
||
|
|
||
|
|
||
|
def test_quiver_limits():
|
||
|
ax = plt.axes()
|
||
|
x, y = np.arange(8), np.arange(10)
|
||
|
u = v = np.linspace(0, 10, 80).reshape(10, 8)
|
||
|
q = plt.quiver(x, y, u, v)
|
||
|
assert q.get_datalim(ax.transData).bounds == (0., 0., 7., 9.)
|
||
|
|
||
|
plt.figure()
|
||
|
ax = plt.axes()
|
||
|
x = np.linspace(-5, 10, 20)
|
||
|
y = np.linspace(-2, 4, 10)
|
||
|
y, x = np.meshgrid(y, x)
|
||
|
trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
|
||
|
plt.quiver(x, y, np.sin(x), np.cos(y), transform=trans)
|
||
|
assert ax.dataLim.bounds == (20.0, 30.0, 15.0, 6.0)
|
||
|
|
||
|
|
||
|
def test_barb_limits():
|
||
|
ax = plt.axes()
|
||
|
x = np.linspace(-5, 10, 20)
|
||
|
y = np.linspace(-2, 4, 10)
|
||
|
y, x = np.meshgrid(y, x)
|
||
|
trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
|
||
|
plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)
|
||
|
# The calculated bounds are approximately the bounds of the original data,
|
||
|
# this is because the entire path is taken into account when updating the
|
||
|
# datalim.
|
||
|
assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
|
||
|
decimal=1)
|
||
|
|
||
|
|
||
|
@image_comparison(['EllipseCollection_test_image.png'], remove_text=True,
|
||
|
tol=0.021 if platform.machine() == 'arm64' else 0)
|
||
|
def test_EllipseCollection():
|
||
|
# Test basic functionality
|
||
|
fig, ax = plt.subplots()
|
||
|
x = np.arange(4)
|
||
|
y = np.arange(3)
|
||
|
X, Y = np.meshgrid(x, y)
|
||
|
XY = np.vstack((X.ravel(), Y.ravel())).T
|
||
|
|
||
|
ww = X / x[-1]
|
||
|
hh = Y / y[-1]
|
||
|
aa = np.ones_like(ww) * 20 # first axis is 20 degrees CCW from x axis
|
||
|
|
||
|
ec = mcollections.EllipseCollection(
|
||
|
ww, hh, aa, units='x', offsets=XY, offset_transform=ax.transData,
|
||
|
facecolors='none')
|
||
|
ax.add_collection(ec)
|
||
|
ax.autoscale_view()
|
||
|
|
||
|
|
||
|
def test_EllipseCollection_setter_getter():
|
||
|
# Test widths, heights and angle setter
|
||
|
rng = np.random.default_rng(0)
|
||
|
|
||
|
widths = (2, )
|
||
|
heights = (3, )
|
||
|
angles = (45, )
|
||
|
offsets = rng.random((10, 2)) * 10
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
|
||
|
ec = mcollections.EllipseCollection(
|
||
|
widths=widths,
|
||
|
heights=heights,
|
||
|
angles=angles,
|
||
|
offsets=offsets,
|
||
|
units='x',
|
||
|
offset_transform=ax.transData,
|
||
|
)
|
||
|
|
||
|
assert_array_almost_equal(ec._widths, np.array(widths).ravel() * 0.5)
|
||
|
assert_array_almost_equal(ec._heights, np.array(heights).ravel() * 0.5)
|
||
|
assert_array_almost_equal(ec._angles, np.deg2rad(angles).ravel())
|
||
|
|
||
|
assert_array_almost_equal(ec.get_widths(), widths)
|
||
|
assert_array_almost_equal(ec.get_heights(), heights)
|
||
|
assert_array_almost_equal(ec.get_angles(), angles)
|
||
|
|
||
|
ax.add_collection(ec)
|
||
|
ax.set_xlim(-2, 12)
|
||
|
ax.set_ylim(-2, 12)
|
||
|
|
||
|
new_widths = rng.random((10, 2)) * 2
|
||
|
new_heights = rng.random((10, 2)) * 3
|
||
|
new_angles = rng.random((10, 2)) * 180
|
||
|
|
||
|
ec.set(widths=new_widths, heights=new_heights, angles=new_angles)
|
||
|
|
||
|
assert_array_almost_equal(ec.get_widths(), new_widths.ravel())
|
||
|
assert_array_almost_equal(ec.get_heights(), new_heights.ravel())
|
||
|
assert_array_almost_equal(ec.get_angles(), new_angles.ravel())
|
||
|
|
||
|
|
||
|
@image_comparison(['polycollection_close.png'], remove_text=True, style='mpl20')
|
||
|
def test_polycollection_close():
|
||
|
from mpl_toolkits.mplot3d import Axes3D # type: ignore
|
||
|
plt.rcParams['axes3d.automargin'] = True
|
||
|
|
||
|
vertsQuad = [
|
||
|
[[0., 0.], [0., 1.], [1., 1.], [1., 0.]],
|
||
|
[[0., 1.], [2., 3.], [2., 2.], [1., 1.]],
|
||
|
[[2., 2.], [2., 3.], [4., 1.], [3., 1.]],
|
||
|
[[3., 0.], [3., 1.], [4., 1.], [4., 0.]]]
|
||
|
|
||
|
fig = plt.figure()
|
||
|
ax = fig.add_axes(Axes3D(fig))
|
||
|
|
||
|
colors = ['r', 'g', 'b', 'y', 'k']
|
||
|
zpos = list(range(5))
|
||
|
|
||
|
poly = mcollections.PolyCollection(
|
||
|
vertsQuad * len(zpos), linewidth=0.25)
|
||
|
poly.set_alpha(0.7)
|
||
|
|
||
|
# need to have a z-value for *each* polygon = element!
|
||
|
zs = []
|
||
|
cs = []
|
||
|
for z, c in zip(zpos, colors):
|
||
|
zs.extend([z] * len(vertsQuad))
|
||
|
cs.extend([c] * len(vertsQuad))
|
||
|
|
||
|
poly.set_color(cs)
|
||
|
|
||
|
ax.add_collection3d(poly, zs=zs, zdir='y')
|
||
|
|
||
|
# axis limit settings:
|
||
|
ax.set_xlim3d(0, 4)
|
||
|
ax.set_zlim3d(0, 3)
|
||
|
ax.set_ylim3d(0, 4)
|
||
|
|
||
|
|
||
|
@image_comparison(['regularpolycollection_rotate.png'], remove_text=True)
|
||
|
def test_regularpolycollection_rotate():
|
||
|
xx, yy = np.mgrid[:10, :10]
|
||
|
xy_points = np.transpose([xx.flatten(), yy.flatten()])
|
||
|
rotations = np.linspace(0, 2*np.pi, len(xy_points))
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
for xy, alpha in zip(xy_points, rotations):
|
||
|
col = mcollections.RegularPolyCollection(
|
||
|
4, sizes=(100,), rotation=alpha,
|
||
|
offsets=[xy], offset_transform=ax.transData)
|
||
|
ax.add_collection(col, autolim=True)
|
||
|
ax.autoscale_view()
|
||
|
|
||
|
|
||
|
@image_comparison(['regularpolycollection_scale.png'], remove_text=True)
|
||
|
def test_regularpolycollection_scale():
|
||
|
# See issue #3860
|
||
|
|
||
|
class SquareCollection(mcollections.RegularPolyCollection):
|
||
|
def __init__(self, **kwargs):
|
||
|
super().__init__(4, rotation=np.pi/4., **kwargs)
|
||
|
|
||
|
def get_transform(self):
|
||
|
"""Return transform scaling circle areas to data space."""
|
||
|
ax = self.axes
|
||
|
|
||
|
pts2pixels = 72.0 / ax.figure.dpi
|
||
|
|
||
|
scale_x = pts2pixels * ax.bbox.width / ax.viewLim.width
|
||
|
scale_y = pts2pixels * ax.bbox.height / ax.viewLim.height
|
||
|
return mtransforms.Affine2D().scale(scale_x, scale_y)
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
|
||
|
xy = [(0, 0)]
|
||
|
# Unit square has a half-diagonal of `1/sqrt(2)`, so `pi * r**2` equals...
|
||
|
circle_areas = [np.pi / 2]
|
||
|
squares = SquareCollection(
|
||
|
sizes=circle_areas, offsets=xy, offset_transform=ax.transData)
|
||
|
ax.add_collection(squares, autolim=True)
|
||
|
ax.axis([-1, 1, -1, 1])
|
||
|
|
||
|
|
||
|
def test_picking():
|
||
|
fig, ax = plt.subplots()
|
||
|
col = ax.scatter([0], [0], [1000], picker=True)
|
||
|
fig.savefig(io.BytesIO(), dpi=fig.dpi)
|
||
|
mouse_event = SimpleNamespace(x=325, y=240)
|
||
|
found, indices = col.contains(mouse_event)
|
||
|
assert found
|
||
|
assert_array_equal(indices['ind'], [0])
|
||
|
|
||
|
|
||
|
def test_quadmesh_contains():
|
||
|
x = np.arange(4)
|
||
|
X = x[:, None] * x[None, :]
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
mesh = ax.pcolormesh(X)
|
||
|
fig.draw_without_rendering()
|
||
|
xdata, ydata = 0.5, 0.5
|
||
|
x, y = mesh.get_transform().transform((xdata, ydata))
|
||
|
mouse_event = SimpleNamespace(xdata=xdata, ydata=ydata, x=x, y=y)
|
||
|
found, indices = mesh.contains(mouse_event)
|
||
|
assert found
|
||
|
assert_array_equal(indices['ind'], [0])
|
||
|
|
||
|
xdata, ydata = 1.5, 1.5
|
||
|
x, y = mesh.get_transform().transform((xdata, ydata))
|
||
|
mouse_event = SimpleNamespace(xdata=xdata, ydata=ydata, x=x, y=y)
|
||
|
found, indices = mesh.contains(mouse_event)
|
||
|
assert found
|
||
|
assert_array_equal(indices['ind'], [5])
|
||
|
|
||
|
|
||
|
def test_quadmesh_contains_concave():
|
||
|
# Test a concave polygon, V-like shape
|
||
|
x = [[0, -1], [1, 0]]
|
||
|
y = [[0, 1], [1, -1]]
|
||
|
fig, ax = plt.subplots()
|
||
|
mesh = ax.pcolormesh(x, y, [[0]])
|
||
|
fig.draw_without_rendering()
|
||
|
# xdata, ydata, expected
|
||
|
points = [(-0.5, 0.25, True), # left wing
|
||
|
(0, 0.25, False), # between the two wings
|
||
|
(0.5, 0.25, True), # right wing
|
||
|
(0, -0.25, True), # main body
|
||
|
]
|
||
|
for point in points:
|
||
|
xdata, ydata, expected = point
|
||
|
x, y = mesh.get_transform().transform((xdata, ydata))
|
||
|
mouse_event = SimpleNamespace(xdata=xdata, ydata=ydata, x=x, y=y)
|
||
|
found, indices = mesh.contains(mouse_event)
|
||
|
assert found is expected
|
||
|
|
||
|
|
||
|
def test_quadmesh_cursor_data():
|
||
|
x = np.arange(4)
|
||
|
X = x[:, None] * x[None, :]
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
mesh = ax.pcolormesh(X)
|
||
|
# Empty array data
|
||
|
mesh._A = None
|
||
|
fig.draw_without_rendering()
|
||
|
xdata, ydata = 0.5, 0.5
|
||
|
x, y = mesh.get_transform().transform((xdata, ydata))
|
||
|
mouse_event = SimpleNamespace(xdata=xdata, ydata=ydata, x=x, y=y)
|
||
|
# Empty collection should return None
|
||
|
assert mesh.get_cursor_data(mouse_event) is None
|
||
|
|
||
|
# Now test adding the array data, to make sure we do get a value
|
||
|
mesh.set_array(np.ones(X.shape))
|
||
|
assert_array_equal(mesh.get_cursor_data(mouse_event), [1])
|
||
|
|
||
|
|
||
|
def test_quadmesh_cursor_data_multiple_points():
|
||
|
x = [1, 2, 1, 2]
|
||
|
fig, ax = plt.subplots()
|
||
|
mesh = ax.pcolormesh(x, x, np.ones((3, 3)))
|
||
|
fig.draw_without_rendering()
|
||
|
xdata, ydata = 1.5, 1.5
|
||
|
x, y = mesh.get_transform().transform((xdata, ydata))
|
||
|
mouse_event = SimpleNamespace(xdata=xdata, ydata=ydata, x=x, y=y)
|
||
|
# All quads are covering the same square
|
||
|
assert_array_equal(mesh.get_cursor_data(mouse_event), np.ones(9))
|
||
|
|
||
|
|
||
|
def test_linestyle_single_dashes():
|
||
|
plt.scatter([0, 1, 2], [0, 1, 2], linestyle=(0., [2., 2.]))
|
||
|
plt.draw()
|
||
|
|
||
|
|
||
|
@image_comparison(['size_in_xy.png'], remove_text=True)
|
||
|
def test_size_in_xy():
|
||
|
fig, ax = plt.subplots()
|
||
|
|
||
|
widths, heights, angles = (10, 10), 10, 0
|
||
|
widths = 10, 10
|
||
|
coords = [(10, 10), (15, 15)]
|
||
|
e = mcollections.EllipseCollection(
|
||
|
widths, heights, angles, units='xy',
|
||
|
offsets=coords, offset_transform=ax.transData)
|
||
|
|
||
|
ax.add_collection(e)
|
||
|
|
||
|
ax.set_xlim(0, 30)
|
||
|
ax.set_ylim(0, 30)
|
||
|
|
||
|
|
||
|
def test_pandas_indexing(pd):
|
||
|
|
||
|
# Should not fail break when faced with a
|
||
|
# non-zero indexed series
|
||
|
index = [11, 12, 13]
|
||
|
ec = fc = pd.Series(['red', 'blue', 'green'], index=index)
|
||
|
lw = pd.Series([1, 2, 3], index=index)
|
||
|
ls = pd.Series(['solid', 'dashed', 'dashdot'], index=index)
|
||
|
aa = pd.Series([True, False, True], index=index)
|
||
|
|
||
|
Collection(edgecolors=ec)
|
||
|
Collection(facecolors=fc)
|
||
|
Collection(linewidths=lw)
|
||
|
Collection(linestyles=ls)
|
||
|
Collection(antialiaseds=aa)
|
||
|
|
||
|
|
||
|
@mpl.style.context('default')
|
||
|
def test_lslw_bcast():
|
||
|
col = mcollections.PathCollection([])
|
||
|
col.set_linestyles(['-', '-'])
|
||
|
col.set_linewidths([1, 2, 3])
|
||
|
|
||
|
assert col.get_linestyles() == [(0, None)] * 6
|
||
|
assert col.get_linewidths() == [1, 2, 3] * 2
|
||
|
|
||
|
col.set_linestyles(['-', '-', '-'])
|
||
|
assert col.get_linestyles() == [(0, None)] * 3
|
||
|
assert (col.get_linewidths() == [1, 2, 3]).all()
|
||
|
|
||
|
|
||
|
def test_set_wrong_linestyle():
|
||
|
c = Collection()
|
||
|
with pytest.raises(ValueError, match="Do not know how to convert 'fuzzy'"):
|
||
|
c.set_linestyle('fuzzy')
|
||
|
|
||
|
|
||
|
@mpl.style.context('default')
|
||
|
def test_capstyle():
|
||
|
col = mcollections.PathCollection([])
|
||
|
assert col.get_capstyle() is None
|
||
|
col = mcollections.PathCollection([], capstyle='round')
|
||
|
assert col.get_capstyle() == 'round'
|
||
|
col.set_capstyle('butt')
|
||
|
assert col.get_capstyle() == 'butt'
|
||
|
|
||
|
|
||
|
@mpl.style.context('default')
|
||
|
def test_joinstyle():
|
||
|
col = mcollections.PathCollection([])
|
||
|
assert col.get_joinstyle() is None
|
||
|
col = mcollections.PathCollection([], joinstyle='round')
|
||
|
assert col.get_joinstyle() == 'round'
|
||
|
col.set_joinstyle('miter')
|
||
|
assert col.get_joinstyle() == 'miter'
|
||
|
|
||
|
|
||
|
@image_comparison(['cap_and_joinstyle.png'])
|
||
|
def test_cap_and_joinstyle_image():
|
||
|
fig, ax = plt.subplots()
|
||
|
ax.set_xlim([-0.5, 1.5])
|
||
|
ax.set_ylim([-0.5, 2.5])
|
||
|
|
||
|
x = np.array([0.0, 1.0, 0.5])
|
||
|
ys = np.array([[0.0], [0.5], [1.0]]) + np.array([[0.0, 0.0, 1.0]])
|
||
|
|
||
|
segs = np.zeros((3, 3, 2))
|
||
|
segs[:, :, 0] = x
|
||
|
segs[:, :, 1] = ys
|
||
|
line_segments = LineCollection(segs, linewidth=[10, 15, 20])
|
||
|
line_segments.set_capstyle("round")
|
||
|
line_segments.set_joinstyle("miter")
|
||
|
|
||
|
ax.add_collection(line_segments)
|
||
|
ax.set_title('Line collection with customized caps and joinstyle')
|
||
|
|
||
|
|
||
|
@image_comparison(['scatter_post_alpha.png'],
|
||
|
remove_text=True, style='default')
|
||
|
def test_scatter_post_alpha():
|
||
|
fig, ax = plt.subplots()
|
||
|
sc = ax.scatter(range(5), range(5), c=range(5))
|
||
|
sc.set_alpha(.1)
|
||
|
|
||
|
|
||
|
def test_scatter_alpha_array():
|
||
|
x = np.arange(5)
|
||
|
alpha = x / 5
|
||
|
# With colormapping.
|
||
|
fig, (ax0, ax1) = plt.subplots(2)
|
||
|
sc0 = ax0.scatter(x, x, c=x, alpha=alpha)
|
||
|
sc1 = ax1.scatter(x, x, c=x)
|
||
|
sc1.set_alpha(alpha)
|
||
|
plt.draw()
|
||
|
assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
|
||
|
assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
|
||
|
# Without colormapping.
|
||
|
fig, (ax0, ax1) = plt.subplots(2)
|
||
|
sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'], alpha=alpha)
|
||
|
sc1 = ax1.scatter(x, x, color='r', alpha=alpha)
|
||
|
plt.draw()
|
||
|
assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
|
||
|
assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
|
||
|
# Without colormapping, and set alpha afterward.
|
||
|
fig, (ax0, ax1) = plt.subplots(2)
|
||
|
sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'])
|
||
|
sc0.set_alpha(alpha)
|
||
|
sc1 = ax1.scatter(x, x, color='r')
|
||
|
sc1.set_alpha(alpha)
|
||
|
plt.draw()
|
||
|
assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
|
||
|
assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
|
||
|
|
||
|
|
||
|
def test_pathcollection_legend_elements():
|
||
|
np.random.seed(19680801)
|
||
|
x, y = np.random.rand(2, 10)
|
||
|
y = np.random.rand(10)
|
||
|
c = np.random.randint(0, 5, size=10)
|
||
|
s = np.random.randint(10, 300, size=10)
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
sc = ax.scatter(x, y, c=c, s=s, cmap="jet", marker="o", linewidths=0)
|
||
|
|
||
|
h, l = sc.legend_elements(fmt="{x:g}")
|
||
|
assert len(h) == 5
|
||
|
assert l == ["0", "1", "2", "3", "4"]
|
||
|
colors = np.array([line.get_color() for line in h])
|
||
|
colors2 = sc.cmap(np.arange(5)/4)
|
||
|
assert_array_equal(colors, colors2)
|
||
|
l1 = ax.legend(h, l, loc=1)
|
||
|
|
||
|
h2, lab2 = sc.legend_elements(num=9)
|
||
|
assert len(h2) == 9
|
||
|
l2 = ax.legend(h2, lab2, loc=2)
|
||
|
|
||
|
h, l = sc.legend_elements(prop="sizes", alpha=0.5, color="red")
|
||
|
assert all(line.get_alpha() == 0.5 for line in h)
|
||
|
assert all(line.get_markerfacecolor() == "red" for line in h)
|
||
|
l3 = ax.legend(h, l, loc=4)
|
||
|
|
||
|
h, l = sc.legend_elements(prop="sizes", num=4, fmt="{x:.2f}",
|
||
|
func=lambda x: 2*x)
|
||
|
actsizes = [line.get_markersize() for line in h]
|
||
|
labeledsizes = np.sqrt(np.array(l, float) / 2)
|
||
|
assert_array_almost_equal(actsizes, labeledsizes)
|
||
|
l4 = ax.legend(h, l, loc=3)
|
||
|
|
||
|
loc = mpl.ticker.MaxNLocator(nbins=9, min_n_ticks=9-1,
|
||
|
steps=[1, 2, 2.5, 3, 5, 6, 8, 10])
|
||
|
h5, lab5 = sc.legend_elements(num=loc)
|
||
|
assert len(h2) == len(h5)
|
||
|
|
||
|
levels = [-1, 0, 55.4, 260]
|
||
|
h6, lab6 = sc.legend_elements(num=levels, prop="sizes", fmt="{x:g}")
|
||
|
assert [float(l) for l in lab6] == levels[2:]
|
||
|
|
||
|
for l in [l1, l2, l3, l4]:
|
||
|
ax.add_artist(l)
|
||
|
|
||
|
fig.canvas.draw()
|
||
|
|
||
|
|
||
|
def test_EventCollection_nosort():
|
||
|
# Check that EventCollection doesn't modify input in place
|
||
|
arr = np.array([3, 2, 1, 10])
|
||
|
coll = EventCollection(arr)
|
||
|
np.testing.assert_array_equal(arr, np.array([3, 2, 1, 10]))
|
||
|
|
||
|
|
||
|
def test_collection_set_verts_array():
|
||
|
verts = np.arange(80, dtype=np.double).reshape(10, 4, 2)
|
||
|
col_arr = PolyCollection(verts)
|
||
|
col_list = PolyCollection(list(verts))
|
||
|
assert len(col_arr._paths) == len(col_list._paths)
|
||
|
for ap, lp in zip(col_arr._paths, col_list._paths):
|
||
|
assert np.array_equal(ap._vertices, lp._vertices)
|
||
|
assert np.array_equal(ap._codes, lp._codes)
|
||
|
|
||
|
verts_tuple = np.empty(10, dtype=object)
|
||
|
verts_tuple[:] = [tuple(tuple(y) for y in x) for x in verts]
|
||
|
col_arr_tuple = PolyCollection(verts_tuple)
|
||
|
assert len(col_arr._paths) == len(col_arr_tuple._paths)
|
||
|
for ap, atp in zip(col_arr._paths, col_arr_tuple._paths):
|
||
|
assert np.array_equal(ap._vertices, atp._vertices)
|
||
|
assert np.array_equal(ap._codes, atp._codes)
|
||
|
|
||
|
|
||
|
def test_collection_set_array():
|
||
|
vals = [*range(10)]
|
||
|
|
||
|
# Test set_array with list
|
||
|
c = Collection()
|
||
|
c.set_array(vals)
|
||
|
|
||
|
# Test set_array with wrong dtype
|
||
|
with pytest.raises(TypeError, match="^Image data of dtype"):
|
||
|
c.set_array("wrong_input")
|
||
|
|
||
|
# Test if array kwarg is copied
|
||
|
vals[5] = 45
|
||
|
assert np.not_equal(vals, c.get_array()).any()
|
||
|
|
||
|
|
||
|
def test_blended_collection_autolim():
|
||
|
a = [1, 2, 4]
|
||
|
height = .2
|
||
|
|
||
|
xy_pairs = np.column_stack([np.repeat(a, 2), np.tile([0, height], len(a))])
|
||
|
line_segs = xy_pairs.reshape([len(a), 2, 2])
|
||
|
|
||
|
f, ax = plt.subplots()
|
||
|
trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
|
||
|
ax.add_collection(LineCollection(line_segs, transform=trans))
|
||
|
ax.autoscale_view(scalex=True, scaley=False)
|
||
|
np.testing.assert_allclose(ax.get_xlim(), [1., 4.])
|
||
|
|
||
|
|
||
|
def test_singleton_autolim():
|
||
|
fig, ax = plt.subplots()
|
||
|
ax.scatter(0, 0)
|
||
|
np.testing.assert_allclose(ax.get_ylim(), [-0.06, 0.06])
|
||
|
np.testing.assert_allclose(ax.get_xlim(), [-0.06, 0.06])
|
||
|
|
||
|
|
||
|
@pytest.mark.parametrize("transform, expected", [
|
||
|
("transData", (-0.5, 3.5)),
|
||
|
("transAxes", (2.8, 3.2)),
|
||
|
])
|
||
|
def test_autolim_with_zeros(transform, expected):
|
||
|
# 1) Test that a scatter at (0, 0) data coordinates contributes to
|
||
|
# autoscaling even though any(offsets) would be False in that situation.
|
||
|
# 2) Test that specifying transAxes for the transform does not contribute
|
||
|
# to the autoscaling.
|
||
|
fig, ax = plt.subplots()
|
||
|
ax.scatter(0, 0, transform=getattr(ax, transform))
|
||
|
ax.scatter(3, 3)
|
||
|
np.testing.assert_allclose(ax.get_ylim(), expected)
|
||
|
np.testing.assert_allclose(ax.get_xlim(), expected)
|
||
|
|
||
|
|
||
|
def test_quadmesh_set_array_validation(pcfunc):
|
||
|
x = np.arange(11)
|
||
|
y = np.arange(8)
|
||
|
z = np.random.random((7, 10))
|
||
|
fig, ax = plt.subplots()
|
||
|
coll = getattr(ax, pcfunc)(x, y, z)
|
||
|
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (10, 7)")):
|
||
|
coll.set_array(z.reshape(10, 7))
|
||
|
|
||
|
z = np.arange(54).reshape((6, 9))
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (6, 9)")):
|
||
|
coll.set_array(z)
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (54,)")):
|
||
|
coll.set_array(z.ravel())
|
||
|
|
||
|
# RGB(A) tests
|
||
|
z = np.ones((9, 6, 3)) # RGB with wrong X/Y dims
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (9, 6, 3)")):
|
||
|
coll.set_array(z)
|
||
|
|
||
|
z = np.ones((9, 6, 4)) # RGBA with wrong X/Y dims
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (9, 6, 4)")):
|
||
|
coll.set_array(z)
|
||
|
|
||
|
z = np.ones((7, 10, 2)) # Right X/Y dims, bad 3rd dim
|
||
|
with pytest.raises(ValueError, match=re.escape(
|
||
|
"For X (11) and Y (8) with flat shading, A should have shape "
|
||
|
"(7, 10, 3) or (7, 10, 4) or (7, 10) or (70,), not (7, 10, 2)")):
|
||
|
coll.set_array(z)
|
||
|
|
||
|
x = np.arange(10)
|
||
|
y = np.arange(7)
|
||
|
z = np.random.random((7, 10))
|
||
|
fig, ax = plt.subplots()
|
||
|
coll = ax.pcolormesh(x, y, z, shading='gouraud')
|
||
|
|
||
|
|
||
|
def test_polyquadmesh_masked_vertices_array():
|
||
|
xx, yy = np.meshgrid([0, 1, 2], [0, 1, 2, 3])
|
||
|
# 2 x 3 mesh data
|
||
|
zz = (xx*yy)[:-1, :-1]
|
||
|
quadmesh = plt.pcolormesh(xx, yy, zz)
|
||
|
quadmesh.update_scalarmappable()
|
||
|
quadmesh_fc = quadmesh.get_facecolor()[1:, :]
|
||
|
# Mask the origin vertex in x
|
||
|
xx = np.ma.masked_where((xx == 0) & (yy == 0), xx)
|
||
|
polymesh = plt.pcolor(xx, yy, zz)
|
||
|
polymesh.update_scalarmappable()
|
||
|
# One cell should be left out
|
||
|
assert len(polymesh.get_paths()) == 5
|
||
|
# Poly version should have the same facecolors as the end of the quadmesh
|
||
|
assert_array_equal(quadmesh_fc, polymesh.get_facecolor())
|
||
|
|
||
|
# Mask the origin vertex in y
|
||
|
yy = np.ma.masked_where((xx == 0) & (yy == 0), yy)
|
||
|
polymesh = plt.pcolor(xx, yy, zz)
|
||
|
polymesh.update_scalarmappable()
|
||
|
# One cell should be left out
|
||
|
assert len(polymesh.get_paths()) == 5
|
||
|
# Poly version should have the same facecolors as the end of the quadmesh
|
||
|
assert_array_equal(quadmesh_fc, polymesh.get_facecolor())
|
||
|
|
||
|
# Mask the origin cell data
|
||
|
zz = np.ma.masked_where((xx[:-1, :-1] == 0) & (yy[:-1, :-1] == 0), zz)
|
||
|
polymesh = plt.pcolor(zz)
|
||
|
polymesh.update_scalarmappable()
|
||
|
# One cell should be left out
|
||
|
assert len(polymesh.get_paths()) == 5
|
||
|
# Poly version should have the same facecolors as the end of the quadmesh
|
||
|
assert_array_equal(quadmesh_fc, polymesh.get_facecolor())
|
||
|
|
||
|
# Setting array with 1D compressed values is deprecated
|
||
|
with pytest.warns(mpl.MatplotlibDeprecationWarning,
|
||
|
match="Setting a PolyQuadMesh"):
|
||
|
polymesh.set_array(np.ones(5))
|
||
|
|
||
|
# We should also be able to call set_array with a new mask and get
|
||
|
# updated polys
|
||
|
# Remove mask, should add all polys back
|
||
|
zz = np.arange(6).reshape((3, 2))
|
||
|
polymesh.set_array(zz)
|
||
|
polymesh.update_scalarmappable()
|
||
|
assert len(polymesh.get_paths()) == 6
|
||
|
# Add mask should remove polys
|
||
|
zz = np.ma.masked_less(zz, 2)
|
||
|
polymesh.set_array(zz)
|
||
|
polymesh.update_scalarmappable()
|
||
|
assert len(polymesh.get_paths()) == 4
|
||
|
|
||
|
|
||
|
def test_quadmesh_get_coordinates(pcfunc):
|
||
|
x = [0, 1, 2]
|
||
|
y = [2, 4, 6]
|
||
|
z = np.ones(shape=(2, 2))
|
||
|
xx, yy = np.meshgrid(x, y)
|
||
|
coll = getattr(plt, pcfunc)(xx, yy, z)
|
||
|
|
||
|
# shape (3, 3, 2)
|
||
|
coords = np.stack([xx.T, yy.T]).T
|
||
|
assert_array_equal(coll.get_coordinates(), coords)
|
||
|
|
||
|
|
||
|
def test_quadmesh_set_array():
|
||
|
x = np.arange(4)
|
||
|
y = np.arange(4)
|
||
|
z = np.arange(9).reshape((3, 3))
|
||
|
fig, ax = plt.subplots()
|
||
|
coll = ax.pcolormesh(x, y, np.ones(z.shape))
|
||
|
# Test that the collection is able to update with a 2d array
|
||
|
coll.set_array(z)
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_array(), z)
|
||
|
|
||
|
# Check that pre-flattened arrays work too
|
||
|
coll.set_array(np.ones(9))
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_array(), np.ones(9))
|
||
|
|
||
|
z = np.arange(16).reshape((4, 4))
|
||
|
fig, ax = plt.subplots()
|
||
|
coll = ax.pcolormesh(x, y, np.ones(z.shape), shading='gouraud')
|
||
|
# Test that the collection is able to update with a 2d array
|
||
|
coll.set_array(z)
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_array(), z)
|
||
|
|
||
|
# Check that pre-flattened arrays work too
|
||
|
coll.set_array(np.ones(16))
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_array(), np.ones(16))
|
||
|
|
||
|
|
||
|
def test_quadmesh_vmin_vmax(pcfunc):
|
||
|
# test when vmin/vmax on the norm changes, the quadmesh gets updated
|
||
|
fig, ax = plt.subplots()
|
||
|
cmap = mpl.colormaps['plasma']
|
||
|
norm = mpl.colors.Normalize(vmin=0, vmax=1)
|
||
|
coll = getattr(ax, pcfunc)([[1]], cmap=cmap, norm=norm)
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
|
||
|
|
||
|
# Change the vmin/vmax of the norm so that the color is from
|
||
|
# the bottom of the colormap now
|
||
|
norm.vmin, norm.vmax = 1, 2
|
||
|
fig.canvas.draw()
|
||
|
assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
|
||
|
|
||
|
|
||
|
def test_quadmesh_alpha_array(pcfunc):
|
||
|
x = np.arange(4)
|
||
|
y = np.arange(4)
|
||
|
z = np.arange(9).reshape((3, 3))
|
||
|
alpha = z / z.max()
|
||
|
alpha_flat = alpha.ravel()
|
||
|
# Provide 2-D alpha:
|
||
|
fig, (ax0, ax1) = plt.subplots(2)
|
||
|
coll1 = getattr(ax0, pcfunc)(x, y, z, alpha=alpha)
|
||
|
coll2 = getattr(ax0, pcfunc)(x, y, z)
|
||
|
coll2.set_alpha(alpha)
|
||
|
plt.draw()
|
||
|
assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
|
||
|
assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
|
||
|
# Or provide 1-D alpha:
|
||
|
fig, (ax0, ax1) = plt.subplots(2)
|
||
|
coll1 = getattr(ax0, pcfunc)(x, y, z, alpha=alpha)
|
||
|
coll2 = getattr(ax1, pcfunc)(x, y, z)
|
||
|
coll2.set_alpha(alpha)
|
||
|
plt.draw()
|
||
|
assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
|
||
|
assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
|
||
|
|
||
|
|
||
|
def test_alpha_validation(pcfunc):
|
||
|
# Most of the relevant testing is in test_artist and test_colors.
|
||
|
fig, ax = plt.subplots()
|
||
|
pc = getattr(ax, pcfunc)(np.arange(12).reshape((3, 4)))
|
||
|
with pytest.raises(ValueError, match="^Data array shape"):
|
||
|
pc.set_alpha([0.5, 0.6])
|
||
|
pc.update_scalarmappable()
|
||
|
|
||
|
|
||
|
def test_legend_inverse_size_label_relationship():
|
||
|
"""
|
||
|
Ensure legend markers scale appropriately when label and size are
|
||
|
inversely related.
|
||
|
Here label = 5 / size
|
||
|
"""
|
||
|
|
||
|
np.random.seed(19680801)
|
||
|
X = np.random.random(50)
|
||
|
Y = np.random.random(50)
|
||
|
C = 1 - np.random.random(50)
|
||
|
S = 5 / C
|
||
|
|
||
|
legend_sizes = [0.2, 0.4, 0.6, 0.8]
|
||
|
fig, ax = plt.subplots()
|
||
|
sc = ax.scatter(X, Y, s=S)
|
||
|
handles, labels = sc.legend_elements(
|
||
|
prop='sizes', num=legend_sizes, func=lambda s: 5 / s
|
||
|
)
|
||
|
|
||
|
# Convert markersize scale to 's' scale
|
||
|
handle_sizes = [x.get_markersize() for x in handles]
|
||
|
handle_sizes = [5 / x**2 for x in handle_sizes]
|
||
|
|
||
|
assert_array_almost_equal(handle_sizes, legend_sizes, decimal=1)
|
||
|
|
||
|
|
||
|
@mpl.style.context('default')
|
||
|
def test_color_logic(pcfunc):
|
||
|
pcfunc = getattr(plt, pcfunc)
|
||
|
z = np.arange(12).reshape(3, 4)
|
||
|
# Explicitly set an edgecolor.
|
||
|
pc = pcfunc(z, edgecolors='red', facecolors='none')
|
||
|
pc.update_scalarmappable() # This is called in draw().
|
||
|
# Define 2 reference "colors" here for multiple use.
|
||
|
face_default = mcolors.to_rgba_array(pc._get_default_facecolor())
|
||
|
mapped = pc.get_cmap()(pc.norm(z.ravel()))
|
||
|
# GitHub issue #1302:
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), 'red')
|
||
|
# Check setting attributes after initialization:
|
||
|
pc = pcfunc(z)
|
||
|
pc.set_facecolor('none')
|
||
|
pc.set_edgecolor('red')
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_facecolor(), 'none')
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
pc.set_alpha(0.5)
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 0.5]])
|
||
|
pc.set_alpha(None) # restore default alpha
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
# Reset edgecolor to default.
|
||
|
pc.set_edgecolor(None)
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_edgecolor(), mapped)
|
||
|
pc.set_facecolor(None) # restore default for facecolor
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_facecolor(), mapped)
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), 'none')
|
||
|
# Turn off colormapping entirely:
|
||
|
pc.set_array(None)
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), 'none')
|
||
|
assert mcolors.same_color(pc.get_facecolor(), face_default) # not mapped
|
||
|
# Turn it back on by restoring the array (must be 1D!):
|
||
|
pc.set_array(z)
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_facecolor(), mapped)
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), 'none')
|
||
|
# Give color via tuple rather than string.
|
||
|
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=(0, 1, 0))
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_facecolor(), mapped)
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
# Provide an RGB array; mapping overrides it.
|
||
|
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 3)))
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_facecolor(), mapped)
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
# Turn off the mapping.
|
||
|
pc.set_array(None)
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_facecolor(), np.ones((12, 3)))
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
# And an RGBA array.
|
||
|
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 4)))
|
||
|
pc.update_scalarmappable()
|
||
|
assert np.array_equal(pc.get_facecolor(), mapped)
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
# Turn off the mapping.
|
||
|
pc.set_array(None)
|
||
|
pc.update_scalarmappable()
|
||
|
assert mcolors.same_color(pc.get_facecolor(), np.ones((12, 4)))
|
||
|
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
|
||
|
|
||
|
|
||
|
def test_LineCollection_args():
|
||
|
lc = LineCollection(None, linewidth=2.2, edgecolor='r',
|
||
|
zorder=3, facecolors=[0, 1, 0, 1])
|
||
|
assert lc.get_linewidth()[0] == 2.2
|
||
|
assert mcolors.same_color(lc.get_edgecolor(), 'r')
|
||
|
assert lc.get_zorder() == 3
|
||
|
assert mcolors.same_color(lc.get_facecolor(), [[0, 1, 0, 1]])
|
||
|
# To avoid breaking mplot3d, LineCollection internally sets the facecolor
|
||
|
# kwarg if it has not been specified. Hence we need the following test
|
||
|
# for LineCollection._set_default().
|
||
|
lc = LineCollection(None, facecolor=None)
|
||
|
assert mcolors.same_color(lc.get_facecolor(), 'none')
|
||
|
|
||
|
|
||
|
def test_array_dimensions(pcfunc):
|
||
|
# Make sure we can set the 1D, 2D, and 3D array shapes
|
||
|
z = np.arange(12).reshape(3, 4)
|
||
|
pc = getattr(plt, pcfunc)(z)
|
||
|
# 1D
|
||
|
pc.set_array(z.ravel())
|
||
|
pc.update_scalarmappable()
|
||
|
# 2D
|
||
|
pc.set_array(z)
|
||
|
pc.update_scalarmappable()
|
||
|
# 3D RGB is OK as well
|
||
|
z = np.arange(36, dtype=np.uint8).reshape(3, 4, 3)
|
||
|
pc.set_array(z)
|
||
|
pc.update_scalarmappable()
|
||
|
|
||
|
|
||
|
def test_get_segments():
|
||
|
segments = np.tile(np.linspace(0, 1, 256), (2, 1)).T
|
||
|
lc = LineCollection([segments])
|
||
|
|
||
|
readback, = lc.get_segments()
|
||
|
# these should comeback un-changed!
|
||
|
assert np.all(segments == readback)
|
||
|
|
||
|
|
||
|
def test_set_offsets_late():
|
||
|
identity = mtransforms.IdentityTransform()
|
||
|
sizes = [2]
|
||
|
|
||
|
null = mcollections.CircleCollection(sizes=sizes)
|
||
|
|
||
|
init = mcollections.CircleCollection(sizes=sizes, offsets=(10, 10))
|
||
|
|
||
|
late = mcollections.CircleCollection(sizes=sizes)
|
||
|
late.set_offsets((10, 10))
|
||
|
|
||
|
# Bbox.__eq__ doesn't compare bounds
|
||
|
null_bounds = null.get_datalim(identity).bounds
|
||
|
init_bounds = init.get_datalim(identity).bounds
|
||
|
late_bounds = late.get_datalim(identity).bounds
|
||
|
|
||
|
# offsets and transform are applied when set after initialization
|
||
|
assert null_bounds != init_bounds
|
||
|
assert init_bounds == late_bounds
|
||
|
|
||
|
|
||
|
def test_set_offset_transform():
|
||
|
skew = mtransforms.Affine2D().skew(2, 2)
|
||
|
init = mcollections.Collection(offset_transform=skew)
|
||
|
|
||
|
late = mcollections.Collection()
|
||
|
late.set_offset_transform(skew)
|
||
|
|
||
|
assert skew == init.get_offset_transform() == late.get_offset_transform()
|
||
|
|
||
|
|
||
|
def test_set_offset_units():
|
||
|
# passing the offsets in initially (i.e. via scatter)
|
||
|
# should yield the same results as `set_offsets`
|
||
|
x = np.linspace(0, 10, 5)
|
||
|
y = np.sin(x)
|
||
|
d = x * np.timedelta64(24, 'h') + np.datetime64('2021-11-29')
|
||
|
|
||
|
sc = plt.scatter(d, y)
|
||
|
off0 = sc.get_offsets()
|
||
|
sc.set_offsets(list(zip(d, y)))
|
||
|
np.testing.assert_allclose(off0, sc.get_offsets())
|
||
|
|
||
|
# try the other way around
|
||
|
fig, ax = plt.subplots()
|
||
|
sc = ax.scatter(y, d)
|
||
|
off0 = sc.get_offsets()
|
||
|
sc.set_offsets(list(zip(y, d)))
|
||
|
np.testing.assert_allclose(off0, sc.get_offsets())
|
||
|
|
||
|
|
||
|
@image_comparison(baseline_images=["test_check_masked_offsets"],
|
||
|
extensions=["png"], remove_text=True, style="mpl20")
|
||
|
def test_check_masked_offsets():
|
||
|
# Check if masked data is respected by scatter
|
||
|
# Ref: Issue #24545
|
||
|
unmasked_x = [
|
||
|
datetime(2022, 12, 15, 4, 49, 52),
|
||
|
datetime(2022, 12, 15, 4, 49, 53),
|
||
|
datetime(2022, 12, 15, 4, 49, 54),
|
||
|
datetime(2022, 12, 15, 4, 49, 55),
|
||
|
datetime(2022, 12, 15, 4, 49, 56),
|
||
|
]
|
||
|
|
||
|
masked_y = np.ma.array([1, 2, 3, 4, 5], mask=[0, 1, 1, 0, 0])
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
ax.scatter(unmasked_x, masked_y)
|
||
|
|
||
|
|
||
|
@check_figures_equal(extensions=["png"])
|
||
|
def test_masked_set_offsets(fig_ref, fig_test):
|
||
|
x = np.ma.array([1, 2, 3, 4, 5], mask=[0, 0, 1, 1, 0])
|
||
|
y = np.arange(1, 6)
|
||
|
|
||
|
ax_test = fig_test.add_subplot()
|
||
|
scat = ax_test.scatter(x, y)
|
||
|
scat.set_offsets(np.ma.column_stack([x, y]))
|
||
|
ax_test.set_xticks([])
|
||
|
ax_test.set_yticks([])
|
||
|
|
||
|
ax_ref = fig_ref.add_subplot()
|
||
|
ax_ref.scatter([1, 2, 5], [1, 2, 5])
|
||
|
ax_ref.set_xticks([])
|
||
|
ax_ref.set_yticks([])
|
||
|
|
||
|
|
||
|
def test_check_offsets_dtype():
|
||
|
# Check that setting offsets doesn't change dtype
|
||
|
x = np.ma.array([1, 2, 3, 4, 5], mask=[0, 0, 1, 1, 0])
|
||
|
y = np.arange(1, 6)
|
||
|
|
||
|
fig, ax = plt.subplots()
|
||
|
scat = ax.scatter(x, y)
|
||
|
masked_offsets = np.ma.column_stack([x, y])
|
||
|
scat.set_offsets(masked_offsets)
|
||
|
assert isinstance(scat.get_offsets(), type(masked_offsets))
|
||
|
|
||
|
unmasked_offsets = np.column_stack([x, y])
|
||
|
scat.set_offsets(unmasked_offsets)
|
||
|
assert isinstance(scat.get_offsets(), type(unmasked_offsets))
|
||
|
|
||
|
|
||
|
@pytest.mark.parametrize('gapcolor', ['orange', ['r', 'k']])
|
||
|
@check_figures_equal(extensions=['png'])
|
||
|
@mpl.rc_context({'lines.linewidth': 20})
|
||
|
def test_striped_lines(fig_test, fig_ref, gapcolor):
|
||
|
ax_test = fig_test.add_subplot(111)
|
||
|
ax_ref = fig_ref.add_subplot(111)
|
||
|
|
||
|
for ax in [ax_test, ax_ref]:
|
||
|
ax.set_xlim(0, 6)
|
||
|
ax.set_ylim(0, 1)
|
||
|
|
||
|
x = range(1, 6)
|
||
|
linestyles = [':', '-', '--']
|
||
|
|
||
|
ax_test.vlines(x, 0, 1, linestyle=linestyles, gapcolor=gapcolor, alpha=0.5)
|
||
|
|
||
|
if isinstance(gapcolor, str):
|
||
|
gapcolor = [gapcolor]
|
||
|
|
||
|
for x, gcol, ls in zip(x, itertools.cycle(gapcolor),
|
||
|
itertools.cycle(linestyles)):
|
||
|
ax_ref.axvline(x, 0, 1, linestyle=ls, gapcolor=gcol, alpha=0.5)
|