Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f14469d

Browse files
tacaswelltacaswell
committed
MNT: convert tests and internal usage way from using mpl.cm.get_cmap
1 parent 0e46eee commit f14469d

12 files changed

+71
-46
lines changed

‎lib/matplotlib/cm.py

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -550,8 +550,8 @@ def set_cmap(self, cmap):
550550
cmap : `.Colormap` or str or None
551551
"""
552552
in_init = self.cmap is None
553-
cmap = get_cmap(cmap)
554-
self.cmap = cmap
553+
554+
self.cmap = _ensure_cmap(cmap)
555555
if not in_init:
556556
self.changed() # Things are not set up properly yet.
557557

@@ -663,3 +663,26 @@ def changed(self):
663663
*vmin*/*vmax* when a *norm* instance is given (but using a `str` *norm*
664664
name together with *vmin*/*vmax* is acceptable).""",
665665
)
666+
667+
668+
def _ensure_cmap(cmap):
669+
"""
670+
Ensure that we have a `.Colormap` object.
671+
672+
Parameters
673+
----------
674+
cmap : None, str, Colormap
675+
676+
- if a `Colormap` return it
677+
- in a string look it up in mpl.colormaps
678+
- if None, look up the default color map in mpl.colormaps
679+
680+
Returns
681+
-------
682+
Colormap
683+
"""
684+
if isinstance(cmap, colors.Colormap):
685+
return cmap
686+
return mpl.colormaps[
687+
cmap if cmap is not None else mpl.rcParams['image.cmap']
688+
]

‎lib/matplotlib/pyplot.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2084,7 +2084,7 @@ def set_cmap(cmap):
20842084
matplotlib.cm.register_cmap
20852085
matplotlib.cm.get_cmap
20862086
"""
2087-
cmap = cm.get_cmap(cmap)
2087+
cmap = colormaps[cmap]
20882088

20892089
rc('image', cmap=cmap.name)
20902090
im = gci()

‎lib/matplotlib/streamplot.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
187187
if use_multicolor_lines:
188188
if norm is None:
189189
norm = mcolors.Normalize(color.min(), color.max())
190-
cmap = cm.get_cmap(cmap)
190+
cmap = cm._ensure_cmap(cmap)
191191

192192
streamlines = []
193193
arrows = []

‎lib/matplotlib/tests/test_artist.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
import pytest
77

8-
from matplotlib import cm
98
import matplotlib.colors as mcolors
109
import matplotlib.pyplot as plt
1110
import matplotlib.patches as mpatches
@@ -14,6 +13,7 @@
1413
import matplotlib.transforms as mtransforms
1514
import matplotlib.collections as mcollections
1615
import matplotlib.artist as martist
16+
import matplotlib as mpl
1717
from matplotlib.testing.decorators import check_figures_equal, image_comparison
1818

1919

@@ -415,7 +415,7 @@ def test_format_cursor_data_BoundaryNorm():
415415
# map range -1..1 to 0..256 in 0.01 steps
416416
fig, ax = plt.subplots()
417417
fig.suptitle("-1..1 to 0..256 in 0.01")
418-
cmap = cm.get_cmap('RdBu_r', 200)
418+
cmap = mpl.colormaps['RdBu_r'].resample(200)
419419
norm = mcolors.BoundaryNorm(np.linspace(-1, 1, 200), 200)
420420
img = ax.imshow(X, cmap=cmap, norm=norm)
421421

@@ -439,7 +439,7 @@ def test_format_cursor_data_BoundaryNorm():
439439
# map range -1..1 to 0..256 in 0.01 steps
440440
fig, ax = plt.subplots()
441441
fig.suptitle("-1..1 to 0..256 in 0.001")
442-
cmap = cm.get_cmap('RdBu_r', 2000)
442+
cmap = mpl.colormaps['RdBu_r'].resample(2000)
443443
norm = mcolors.BoundaryNorm(np.linspace(-1, 1, 2000), 2000)
444444
img = ax.imshow(X, cmap=cmap, norm=norm)
445445

‎lib/matplotlib/tests/test_axes.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1220,7 +1220,7 @@ def test_pcolormesh_alpha():
12201220
Qy = Y + np.sin(X)
12211221
Z = np.hypot(X, Y) / 5
12221222
Z = (Z - Z.min()) / Z.ptp()
1223-
vir = plt.get_cmap("viridis", 16)
1223+
vir = mpl.colormaps["viridis"].resample(16)
12241224
# make another colormap with varying alpha
12251225
colors = vir(np.arange(16))
12261226
colors[:, 3] = 0.5 + 0.5*np.sin(np.arange(16))
@@ -2250,7 +2250,7 @@ def test_contour_hatching():
22502250
x, y, z = contour_dat()
22512251
fig, ax = plt.subplots()
22522252
ax.contourf(x, y, z, 7, hatches=['/', '\\', '//', '-'],
2253-
cmap=plt.get_cmap('gray'),
2253+
cmap=mpl.colormaps['gray'],
22542254
extend='both', alpha=0.5)
22552255

22562256

@@ -2260,7 +2260,7 @@ def test_contour_colorbar():
22602260

22612261
fig, ax = plt.subplots()
22622262
cs = ax.contourf(x, y, z, levels=np.arange(-1.8, 1.801, 0.2),
2263-
cmap=plt.get_cmap('RdBu'),
2263+
cmap=mpl.colormaps['RdBu'],
22642264
vmin=-0.6,
22652265
vmax=0.6,
22662266
extend='both')
@@ -2444,7 +2444,7 @@ def test_scatter_edgecolor_RGB(self):
24442444
@check_figures_equal(extensions=["png"])
24452445
def test_scatter_invalid_color(self, fig_test, fig_ref):
24462446
ax = fig_test.subplots()
2447-
cmap = plt.get_cmap("viridis", 16)
2447+
cmap = mpl.colormaps["viridis"].resample(16)
24482448
cmap.set_bad("k", 1)
24492449
# Set a nonuniform size to prevent the last call to `scatter` (plotting
24502450
# the invalid points separately in fig_ref) from using the marker
@@ -2453,15 +2453,15 @@ def test_scatter_invalid_color(self, fig_test, fig_ref):
24532453
c=[1, np.nan, 2, np.nan], s=[1, 2, 3, 4],
24542454
cmap=cmap, plotnonfinite=True)
24552455
ax = fig_ref.subplots()
2456-
cmap = plt.get_cmap("viridis", 16)
2456+
cmap = mpl.colormaps["viridis"].resample(16)
24572457
ax.scatter([0, 2], [0, 2], c=[1, 2], s=[1, 3], cmap=cmap)
24582458
ax.scatter([1, 3], [1, 3], s=[2, 4], color="k")
24592459

24602460
@check_figures_equal(extensions=["png"])
24612461
def test_scatter_no_invalid_color(self, fig_test, fig_ref):
24622462
# With plotnonfinite=False we plot only 2 points.
24632463
ax = fig_test.subplots()
2464-
cmap = plt.get_cmap("viridis", 16)
2464+
cmap = mpl.colormaps["viridis"].resample(16)
24652465
cmap.set_bad("k", 1)
24662466
ax.scatter(range(4), range(4),
24672467
c=[1, np.nan, 2, np.nan], s=[1, 2, 3, 4],

‎lib/matplotlib/tests/test_collections.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -955,7 +955,7 @@ def test_quadmesh_set_array():
955955
def test_quadmesh_vmin_vmax():
956956
# test when vmin/vmax on the norm changes, the quadmesh gets updated
957957
fig, ax = plt.subplots()
958-
cmap = mpl.cm.get_cmap('plasma')
958+
cmap = mpl.colormaps['plasma']
959959
norm = mpl.colors.Normalize(vmin=0, vmax=1)
960960
coll = ax.pcolormesh([[1]], cmap=cmap, norm=norm)
961961
fig.canvas.draw()

‎lib/matplotlib/tests/test_colorbar.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from matplotlib import cm
55
import matplotlib.colors as mcolors
6+
import matplotlib as mpl
67

78
from matplotlib import rc_context
89
from matplotlib.testing.decorators import image_comparison
@@ -24,7 +25,7 @@ def _get_cmap_norms():
2425
colorbar_extension_length.
2526
"""
2627
# Create a colormap and specify the levels it represents.
27-
cmap = cm.get_cmap("RdBu", lut=5)
28+
cmap = mpl.colormaps["RdBu"].resample(5)
2829
clevs = [-5., -2.5, -.5, .5, 1.5, 3.5]
2930
# Define norms for the colormaps.
3031
norms = dict()
@@ -132,8 +133,8 @@ def test_colorbar_extension_inverted_axis(orientation, extend, expected):
132133
"""Test extension color with an inverted axis"""
133134
data = np.arange(12).reshape(3, 4)
134135
fig, ax = plt.subplots()
135-
cmap = plt.get_cmap("viridis").with_extremes(under=(0, 0, 0, 1),
136-
over=(1, 1, 1, 1))
136+
cmap = mpl.colormaps["viridis"].with_extremes(under=(0, 0, 0, 1),
137+
over=(1, 1, 1, 1))
137138
im = ax.imshow(data, cmap=cmap)
138139
cbar = fig.colorbar(im, orientation=orientation, extend=extend)
139140
if orientation == "horizontal":
@@ -268,7 +269,7 @@ def test_colorbar_single_scatter():
268269
plt.figure()
269270
x = y = [0]
270271
z = [50]
271-
cmap = plt.get_cmap('jet', 16)
272+
cmap = mpl.colormaps['jet'].resample(16)
272273
cs = plt.scatter(x, y, z, c=z, cmap=cmap)
273274
plt.colorbar(cs)
274275

@@ -326,7 +327,7 @@ def test_colorbar_closed_patch():
326327
ax4 = fig.add_axes([0.05, 0.25, 0.9, 0.1])
327328
ax5 = fig.add_axes([0.05, 0.05, 0.9, 0.1])
328329

329-
cmap = cm.get_cmap("RdBu", lut=5)
330+
cmap = mpl.colormaps["RdBu"].resample(5)
330331

331332
im = ax1.pcolormesh(np.linspace(0, 10, 16).reshape((4, 4)), cmap=cmap)
332333

@@ -957,7 +958,7 @@ def test_colorbar_extend_drawedges():
957958
fig.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95)
958959

959960
for ax, (extend, coloroffset, res) in zip(axs, params):
960-
cmap = plt.get_cmap("viridis")
961+
cmap = mpl.colormaps["viridis"]
961962
bounds = np.arange(5)
962963
nb_colors = len(bounds) + coloroffset
963964
colors = cmap(np.linspace(100, 255, nb_colors).astype(int))
@@ -980,7 +981,7 @@ def test_colorbar_extend_drawedges():
980981

981982
def test_negative_boundarynorm():
982983
fig, ax = plt.subplots(figsize=(1, 3))
983-
cmap = plt.get_cmap("viridis")
984+
cmap = mpl.colormaps["viridis"]
984985

985986
clevs = np.arange(-94, -85)
986987
norm = BoundaryNorm(clevs, cmap.N)
@@ -1016,7 +1017,7 @@ def test_nonorm():
10161017
fig.subplots_adjust(bottom=0.5)
10171018

10181019
norm = NoNorm(vmin=min(data), vmax=max(data))
1019-
cmap = cm.get_cmap("viridis", len(data))
1020+
cmap = mpl.colormaps["viridis"].resample(len(data))
10201021
mappable = cm.ScalarMappable(norm=norm, cmap=cmap)
10211022
cbar = fig.colorbar(mappable, cax=ax, orientation="horizontal")
10221023

‎lib/matplotlib/tests/test_colors.py

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
from matplotlib import cbook, cm, cycler
1414
import matplotlib
15+
import matplotlib as mpl
1516
import matplotlib.colors as mcolors
1617
import matplotlib.colorbar as mcolorbar
1718
import matplotlib.pyplot as plt
@@ -127,7 +128,7 @@ def test_colormap_copy():
127128

128129

129130
def test_colormap_equals():
130-
cmap = plt.get_cmap("plasma")
131+
cmap = mpl.colormaps["plasma"]
131132
cm_copy = cmap.copy()
132133
# different object id's
133134
assert cm_copy is not cmap
@@ -155,7 +156,7 @@ def test_colormap_endian():
155156
mapping of 1.0 when input from a non-native-byteorder
156157
array.
157158
"""
158-
cmap = cm.get_cmap("jet")
159+
cmap = mpl.colormaps["jet"]
159160
# Test under, over, and invalid along with values 0 and 1.
160161
a = [-0.5, 0, 0.5, 1, 1.5, np.nan]
161162
for dt in ["f2", "f4", "f8"]:
@@ -170,7 +171,7 @@ def test_colormap_invalid():
170171
rather than bad. This tests to make sure all invalid values
171172
(-inf, nan, inf) are mapped respectively to (under, bad, over).
172173
"""
173-
cmap = cm.get_cmap("plasma")
174+
cmap = mpl.colormaps["plasma"]
174175
x = np.array([-np.inf, -1, 0, np.nan, .7, 2, np.inf])
175176

176177
expected = np.array([[0.050383, 0.029803, 0.527975, 1.],
@@ -203,7 +204,7 @@ def test_colormap_return_types():
203204
Make sure that tuples are returned for scalar input and
204205
that the proper shapes are returned for ndarrays.
205206
"""
206-
cmap = cm.get_cmap("plasma")
207+
cmap = mpl.colormaps["plasma"]
207208
# Test return types and shapes
208209
# scalar input needs to return a tuple of length 4
209210
assert isinstance(cmap(0.5), tuple)
@@ -318,7 +319,7 @@ def test_BoundaryNorm():
318319

319320
# Testing extend keyword, with interpolation (large cmap)
320321
bounds = [1, 2, 3]
321-
cmap = cm.get_cmap('viridis')
322+
cmap = mpl.colormaps['viridis']
322323
mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
323324
refnorm = mcolors.BoundaryNorm([0] + bounds + [4], cmap.N)
324325
x = np.random.randn(100) * 10 + 2
@@ -789,7 +790,7 @@ def test_boundarynorm_and_colorbarbase():
789790

790791
# Set the colormap and bounds
791792
bounds = [-1, 2, 5, 7, 12, 15]
792-
cmap = cm.get_cmap('viridis')
793+
cmap = mpl.colormaps['viridis']
793794

794795
# Default behavior
795796
norm = mcolors.BoundaryNorm(bounds, cmap.N)
@@ -1134,13 +1135,13 @@ def test_pandas_iterable(pd):
11341135
assert_array_equal(cm1.colors, cm2.colors)
11351136

11361137

1137-
@pytest.mark.parametrize('name', sorted(plt.colormaps()))
1138+
@pytest.mark.parametrize('name', sorted(mpl.colormaps()))
11381139
def test_colormap_reversing(name):
11391140
"""
11401141
Check the generated _lut data of a colormap and corresponding reversed
11411142
colormap if they are almost the same.
11421143
"""
1143-
cmap = plt.get_cmap(name)
1144+
cmap = mpl.colormaps[name]
11441145
cmap_r = cmap.reversed()
11451146
if not cmap_r._isinit:
11461147
cmap._init()
@@ -1307,7 +1308,7 @@ def test_hex_shorthand_notation():
13071308

13081309

13091310
def test_repr_png():
1310-
cmap = plt.get_cmap('viridis')
1311+
cmap = mpl.colormaps['viridis']
13111312
png = cmap._repr_png_()
13121313
assert len(png) > 0
13131314
img = Image.open(BytesIO(png))
@@ -1320,7 +1321,7 @@ def test_repr_png():
13201321

13211322

13221323
def test_repr_html():
1323-
cmap = plt.get_cmap('viridis')
1324+
cmap = mpl.colormaps['viridis']
13241325
html = cmap._repr_html_()
13251326
assert len(html) > 0
13261327
png = cmap._repr_png_()
@@ -1331,15 +1332,15 @@ def test_repr_html():
13311332

13321333

13331334
def test_get_under_over_bad():
1334-
cmap = plt.get_cmap('viridis')
1335+
cmap = mpl.colormaps['viridis']
13351336
assert_array_equal(cmap.get_under(), cmap(-np.inf))
13361337
assert_array_equal(cmap.get_over(), cmap(np.inf))
13371338
assert_array_equal(cmap.get_bad(), cmap(np.nan))
13381339

13391340

13401341
@pytest.mark.parametrize('kind', ('over', 'under', 'bad'))
13411342
def test_non_mutable_get_values(kind):
1342-
cmap = copy.copy(plt.get_cmap('viridis'))
1343+
cmap = copy.copy(mpl.colormaps['viridis'])
13431344
init_value = getattr(cmap, f'get_{kind}')()
13441345
getattr(cmap, f'set_{kind}')('k')
13451346
black_value = getattr(cmap, f'get_{kind}')()
@@ -1348,7 +1349,7 @@ def test_non_mutable_get_values(kind):
13481349

13491350

13501351
def test_colormap_alpha_array():
1351-
cmap = plt.get_cmap('viridis')
1352+
cmap = mpl.colormaps['viridis']
13521353
vals = [-1, 0.5, 2] # under, valid, over
13531354
with pytest.raises(ValueError, match="alpha is array-like but"):
13541355
cmap(vals, alpha=[1, 1, 1, 1])
@@ -1360,7 +1361,7 @@ def test_colormap_alpha_array():
13601361

13611362

13621363
def test_colormap_bad_data_with_alpha():
1363-
cmap = plt.get_cmap('viridis')
1364+
cmap = mpl.colormaps['viridis']
13641365
c = cmap(np.nan, alpha=0.5)
13651366
assert c == (0, 0, 0, 0)
13661367
c = cmap([0.5, np.nan], alpha=0.5)

0 commit comments

Comments
 (0)