Skip to content

Commit 84c5335

Browse files
jklymakjklymak
authored
Merge pull request #24638 from greglucas/pcolormesh-no-flatten
MNT: Remove auto-flattening of input data to pcolormesh
2 parents a38681c + f56d58e commit 84c5335

File tree

2 files changed

+16
-19
lines changed

2 files changed

+16
-19
lines changed

‎lib/matplotlib/axes/_axes.py

-3
Original file line numberDiff line numberDiff line change
@@ -6220,9 +6220,6 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
62206220
X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
62216221
shading=shading, kwargs=kwargs)
62226222
coords = np.stack([X, Y], axis=-1)
6223-
# convert to one dimensional array, except for 3D RGB(A) arrays
6224-
if C.ndim != 3:
6225-
C = C.ravel()
62266223

62276224
kwargs.setdefault('snap', mpl.rcParams['pcolormesh.snap'])
62286225

‎lib/matplotlib/tests/test_collections.py

+16-16
Original file line numberDiff line numberDiff line change
@@ -914,13 +914,13 @@ def test_quadmesh_vmin_vmax():
914914
norm = mpl.colors.Normalize(vmin=0, vmax=1)
915915
coll = ax.pcolormesh([[1]], cmap=cmap, norm=norm)
916916
fig.canvas.draw()
917-
assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
917+
assert np.array_equal(coll.get_facecolors()[0, 0, :], cmap(norm(1)))
918918

919919
# Change the vmin/vmax of the norm so that the color is from
920920
# the bottom of the colormap now
921921
norm.vmin, norm.vmax = 1, 2
922922
fig.canvas.draw()
923-
assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
923+
assert np.array_equal(coll.get_facecolors()[0, 0, :], cmap(norm(1)))
924924

925925

926926
def test_quadmesh_alpha_array():
@@ -935,16 +935,16 @@ def test_quadmesh_alpha_array():
935935
coll2 = ax1.pcolormesh(x, y, z)
936936
coll2.set_alpha(alpha)
937937
plt.draw()
938-
assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
939-
assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
938+
assert_array_equal(coll1.get_facecolors()[..., -1], alpha)
939+
assert_array_equal(coll2.get_facecolors()[..., -1], alpha)
940940
# Or provide 1-D alpha:
941941
fig, (ax0, ax1) = plt.subplots(2)
942-
coll1 = ax0.pcolormesh(x, y, z, alpha=alpha_flat)
942+
coll1 = ax0.pcolormesh(x, y, z, alpha=alpha)
943943
coll2 = ax1.pcolormesh(x, y, z)
944-
coll2.set_alpha(alpha_flat)
944+
coll2.set_alpha(alpha)
945945
plt.draw()
946-
assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
947-
assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
946+
assert_array_equal(coll1.get_facecolors()[..., -1], alpha)
947+
assert_array_equal(coll2.get_facecolors()[..., -1], alpha)
948948

949949

950950
def test_alpha_validation():
@@ -992,7 +992,7 @@ def test_color_logic(pcfunc):
992992
pc.update_scalarmappable() # This is called in draw().
993993
# Define 2 reference "colors" here for multiple use.
994994
face_default = mcolors.to_rgba_array(pc._get_default_facecolor())
995-
mapped = pc.get_cmap()(pc.norm(z.ravel()))
995+
mapped = pc.get_cmap()(pc.norm(z.ravel() if pcfunc == plt.pcolor else z))
996996
# GitHub issue #1302:
997997
assert mcolors.same_color(pc.get_edgecolor(), 'red')
998998
# Check setting attributes after initialization:
@@ -1011,30 +1011,30 @@ def test_color_logic(pcfunc):
10111011
# Reset edgecolor to default.
10121012
pc.set_edgecolor(None)
10131013
pc.update_scalarmappable()
1014-
assert mcolors.same_color(pc.get_edgecolor(), mapped)
1014+
assert np.array_equal(pc.get_edgecolor(), mapped)
10151015
pc.set_facecolor(None) # restore default for facecolor
10161016
pc.update_scalarmappable()
1017-
assert mcolors.same_color(pc.get_facecolor(), mapped)
1017+
assert np.array_equal(pc.get_facecolor(), mapped)
10181018
assert mcolors.same_color(pc.get_edgecolor(), 'none')
10191019
# Turn off colormapping entirely:
10201020
pc.set_array(None)
10211021
pc.update_scalarmappable()
10221022
assert mcolors.same_color(pc.get_edgecolor(), 'none')
10231023
assert mcolors.same_color(pc.get_facecolor(), face_default) # not mapped
10241024
# Turn it back on by restoring the array (must be 1D!):
1025-
pc.set_array(z.ravel())
1025+
pc.set_array(z.ravel() if pcfunc == plt.pcolor else z)
10261026
pc.update_scalarmappable()
1027-
assert mcolors.same_color(pc.get_facecolor(), mapped)
1027+
assert np.array_equal(pc.get_facecolor(), mapped)
10281028
assert mcolors.same_color(pc.get_edgecolor(), 'none')
10291029
# Give color via tuple rather than string.
10301030
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=(0, 1, 0))
10311031
pc.update_scalarmappable()
1032-
assert mcolors.same_color(pc.get_facecolor(), mapped)
1032+
assert np.array_equal(pc.get_facecolor(), mapped)
10331033
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
10341034
# Provide an RGB array; mapping overrides it.
10351035
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 3)))
10361036
pc.update_scalarmappable()
1037-
assert mcolors.same_color(pc.get_facecolor(), mapped)
1037+
assert np.array_equal(pc.get_facecolor(), mapped)
10381038
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
10391039
# Turn off the mapping.
10401040
pc.set_array(None)
@@ -1044,7 +1044,7 @@ def test_color_logic(pcfunc):
10441044
# And an RGBA array.
10451045
pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 4)))
10461046
pc.update_scalarmappable()
1047-
assert mcolors.same_color(pc.get_facecolor(), mapped)
1047+
assert np.array_equal(pc.get_facecolor(), mapped)
10481048
assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
10491049
# Turn off the mapping.
10501050
pc.set_array(None)

0 commit comments

Comments
 (0)