diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 4821f2895..916e9c2fa 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -405,8 +405,14 @@ def __call__( mask_under = xa < 0 mask_over = xa >= N - # If input was masked, get the bad mask from it; else mask out nans. - mask_bad = x.mask if np.ma.is_masked(x) else np.isnan(xa) # type: ignore + # If input was masked, start from its mask: a masked array can still carry + # unmasked nans. `|` rather than `|=`, so x's own mask isn't written to. + if np.ma.is_masked(x): + mask_bad = x.mask # type: ignore + if xa.dtype.kind == "f": + mask_bad = mask_bad | np.isnan(xa) + else: + mask_bad = np.isnan(xa) with np.errstate(invalid="ignore"): # We need this cast for unsigned ints as well as floats diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 62911f808..c17ad6f7a 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -144,6 +144,23 @@ def test_colormap_apply() -> None: assert cmap1(swapped).shape == (10, 10, 4) +def test_colormap_masked_array_with_unmasked_nan() -> None: + cmap = Colormap("viridis", bad="red") + mask = [True, False, False, False] + data = np.ma.masked_array([0.0, 0.25, np.nan, 1.0], mask=mask) + + rgba = cmap(data) + + npt.assert_array_equal(rgba[0], Color("red").rgba) # masked + npt.assert_array_equal(rgba[2], Color("red").rgba) # nan, not masked + npt.assert_array_equal(rgba[[1, 3]], cmap(np.array([0.25, 1.0]))) + npt.assert_array_equal(data.mask, mask) + + # an all-false mask takes the same path as a plain array + all_false = np.ma.masked_array([0.25, np.nan], mask=[False, False]) + npt.assert_array_equal(cmap(all_false), cmap(np.array([0.25, np.nan]))) + + def test_fill_stops() -> None: assert _fill_stops([None, None, None]) == [0, 0.5, 1.0] assert _fill_stops([None, 0.8, None]) == [0, 0.8, 1.0]