diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 1d7c425e554..8a2666d25fa 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -41,6 +41,9 @@ Bug fixes - :py:func:`concat` can now handle coordinate variables only present in one of the objects to be concatenated when ``coords="different"``. By `Deepak Cherian `_. +- xarray now respects the over, under and bad colors if set on a provided colormap. + (:issue:`3590`, :pull:`3601`) + By `johnomotani `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/plot/utils.py b/xarray/plot/utils.py index 341ff730e01..cb3bef6d409 100644 --- a/xarray/plot/utils.py +++ b/xarray/plot/utils.py @@ -78,6 +78,30 @@ def _build_discrete_cmap(cmap, levels, extend, filled): # copy the old cmap name, for easier testing new_cmap.name = getattr(cmap, "name", cmap) + # copy colors to use for bad, under, and over values in case they have been + # set to non-default values + try: + # matplotlib<3.2 only uses bad color for masked values + bad = cmap(np.ma.masked_invalid([np.nan]))[0] + except TypeError: + # cmap was a str or list rather than a color-map object, so there are + # no bad, under or over values to check or copy + pass + else: + under = cmap(-np.inf) + over = cmap(np.inf) + + new_cmap.set_bad(bad) + + # Only update under and over if they were explicitly changed by the user + # (i.e. are different from the lowest or highest values in cmap). Otherwise + # leave unchanged so new_cmap uses its default values (its own lowest and + # highest values). + if under != cmap(0): + new_cmap.set_under(under) + if over != cmap(cmap.N - 1): + new_cmap.set_over(over) + return new_cmap, cnorm diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index dda9e5de3b2..9ffbcd9c85e 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -1,4 +1,5 @@ import inspect +from copy import deepcopy from datetime import datetime import numpy as np @@ -275,6 +276,71 @@ def test2d_1d_2d_coordinates_contourf(self): a.plot.contourf(x="time", y="depth") a.plot.contourf(x="depth", y="time") + def test_contourf_cmap_set(self): + a = DataArray(easy_array((4, 4)), dims=["z", "time"]) + + cmap = mpl.cm.viridis + + # deepcopy to ensure cmap is not changed by contourf() + # Set vmin and vmax so that _build_discrete_colormap is called with + # extend='both'. extend is passed to + # mpl.colors.from_levels_and_colors(), which returns a result with + # sensible under and over values if extend='both', but not if + # extend='neither' (but if extend='neither' the under and over values + # would not be used because the data would all be within the plotted + # range) + pl = a.plot.contourf(cmap=deepcopy(cmap), vmin=0.1, vmax=0.9) + + # check the set_bad color + assert np.all( + pl.cmap(np.ma.masked_invalid([np.nan]))[0] + == cmap(np.ma.masked_invalid([np.nan]))[0] + ) + + # check the set_under color + assert pl.cmap(-np.inf) == cmap(-np.inf) + + # check the set_over color + assert pl.cmap(np.inf) == cmap(np.inf) + + def test_contourf_cmap_set_with_bad_under_over(self): + a = DataArray(easy_array((4, 4)), dims=["z", "time"]) + + # Make a copy here because we want a local cmap that we will modify. + # Use deepcopy because matplotlib Colormap objects have tuple members + # and we want to ensure we do not change the original. + cmap = deepcopy(mpl.cm.viridis) + + cmap.set_bad("w") + # check we actually changed the set_bad color + assert np.all( + cmap(np.ma.masked_invalid([np.nan]))[0] + != mpl.cm.viridis(np.ma.masked_invalid([np.nan]))[0] + ) + + cmap.set_under("r") + # check we actually changed the set_under color + assert cmap(-np.inf) != mpl.cm.viridis(-np.inf) + + cmap.set_over("g") + # check we actually changed the set_over color + assert cmap(np.inf) != mpl.cm.viridis(-np.inf) + + # deepcopy to ensure cmap is not changed by contourf() + pl = a.plot.contourf(cmap=deepcopy(cmap)) + + # check the set_bad color has been kept + assert np.all( + pl.cmap(np.ma.masked_invalid([np.nan]))[0] + == cmap(np.ma.masked_invalid([np.nan]))[0] + ) + + # check the set_under color has been kept + assert pl.cmap(-np.inf) == cmap(-np.inf) + + # check the set_over color has been kept + assert pl.cmap(np.inf) == cmap(np.inf) + def test3d(self): self.darray.plot()