From 9380611161ac87b56f05954acc3c40463272e3b3 Mon Sep 17 00:00:00 2001 From: Matthias Schabel Date: Mon, 10 Aug 2026 13:27:46 -0700 Subject: [PATCH] feat: add per-class colors for infinities, NaN, and masked entries cmap colors three exceptional classes: under, over, and bad. Floating point data has more. Negative and positive infinity are indistinguishable from ordinary out-of-range values, and NaN is indistinguishable from a masked entry. Adds neg_inf, pos_inf, nan, and masked. Each falls back to the color its class uses now: neg_inf to under, pos_inf to over, nan and masked to bad. bad is kept as the joint fallback for both of its children, so code that sets it is unaffected and either child may be set alone. Routing appends four fallback-resolved rows to a call-local copy of the over/under LUT, so a class with no color of its own lands on exactly the row it lands on now. Colormap.lut() is unchanged. The infinity masks are taken before the input is scaled by N: that multiply overflows large finite values to infinity (float16 65504 does it), and those are out of range rather than infinite. Co-Authored-By: Claude Opus 5 (1M context) Reviewed-By: Codex (gpt-5.6-sol, reasoning effort xhigh) --- src/cmap/_colormap.py | 145 ++++++++++++++++++++++++++++++++++---- src/cmap/_external.py | 6 +- tests/test_colormap.py | 98 ++++++++++++++++++++++++++ tests/test_third_party.py | 4 ++ 4 files changed, 238 insertions(+), 15 deletions(-) diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 599e18ef9..b576e9613 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -141,6 +141,18 @@ class Colormap: The color to use for NaN and masked values. When no bad color is set, they are transparent. Note that infinities are not bad values here: they use `under` and `over`. + neg_inf : ColorLike | None + The color to use for negative infinity. When unset, negative infinity uses + `under`. + pos_inf : ColorLike | None + The color to use for positive infinity. When unset, positive infinity uses + `over`. + nan : ColorLike | None + The color to use for NaN. When unset, NaN uses `bad`. + masked : ColorLike | None + The color to use for entries masked by a `numpy.ma` masked array. When unset, + masked entries use `bad`. A masked entry takes this color whatever value it + hides, so a masked infinity is masked rather than infinite. Raises ------ @@ -153,6 +165,7 @@ class Colormap: __slots__ = ( "__weakref__", + "_has_exceptional", "_initialized", "_lut_cache", "bad_color", @@ -161,8 +174,12 @@ class Colormap: "identifier", "info", "interpolation", + "masked_color", "name", + "nan_color", + "neg_inf_color", "over_color", + "pos_inf_color", "under_color", ) @@ -224,6 +241,24 @@ class Colormap: If provided, and `Colormap.lut` is called with `with_over_under=True`, `bad_color` will be the last color in the LUT (`lut[-1]`). + + `nan_color` and `masked_color` override it for their own class. It remains the + color both of them fall back to. + """ + + neg_inf_color: Color | None + """A color to use for negative infinity, overriding `under_color`.""" + + pos_inf_color: Color | None + """A color to use for positive infinity, overriding `over_color`.""" + + nan_color: Color | None + """A color to use for NaN, overriding `bad_color`.""" + + masked_color: Color | None + """A color to use for masked entries, overriding `bad_color`. + + Applies to any entry masked by a `numpy.ma` masked array, whatever value it hides. """ _catalog_instance: Catalog | None = None @@ -246,6 +281,10 @@ def __init__( under: ColorLike | None = None, over: ColorLike | None = None, bad: ColorLike | None = None, + neg_inf: ColorLike | None = None, + pos_inf: ColorLike | None = None, + nan: ColorLike | None = None, + masked: ColorLike | None = None, cmap_kwargs: dict[str, Any] | None = None, ) -> None: self.info: CatalogItem | None = None @@ -318,6 +357,20 @@ def __init__( self.under_color = None if under is None else Color(under) self.over_color = None if over is None else Color(over) self.bad_color = None if bad is None else Color(bad) + self.neg_inf_color = None if neg_inf is None else Color(neg_inf) + self.pos_inf_color = None if pos_inf is None else Color(pos_inf) + self.nan_color = None if nan is None else Color(nan) + self.masked_color = None if masked is None else Color(masked) + # a colormap with none of these takes the same path it did before they existed + self._has_exceptional = any( + c is not None + for c in ( + self.neg_inf_color, + self.pos_inf_color, + self.nan_color, + self.masked_color, + ) + ) self._lut_cache: dict[LutCacheKey, np.ndarray] = {} self._initialized = True @@ -361,12 +414,14 @@ def __call__( For float input, values outside the [0, 1] range and values that are not finite do not map into the ramp: - - values below 0, and negative infinity, use `under_color` (when unset, the - first color in the colormap). - - values above 1, and positive infinity, use `over_color` (when unset, the - last color in the colormap). - - NaN, and entries masked by a `numpy.ma` masked array, use `bad_color` - (when unset, transparent). + - values below 0 use `under_color` (when unset, the first color in the + colormap), and values above 1 use `over_color` (when unset, the last). + - negative and positive infinity use `neg_inf_color` and `pos_inf_color` + (when unset, `under_color` and `over_color`). + - NaN uses `nan_color`, and entries masked by a `numpy.ma` masked array use + `masked_color` (when either is unset, `bad_color`, which is itself + transparent when unset). A masked entry takes the masked color whatever + value it hides. For integer input, which indexes the LUT directly, an index at or beyond N uses `over_color`, and a negative index uses `under_color` rather than @@ -410,16 +465,24 @@ def __call__( >>> colored_img = cmap(data) """ lut = self.lut(N=N, gamma=gamma, with_over_under=True) - if bytes: - lut = (lut * 255).astype(np.uint8) # the lut will have three additional colors at the end for under, over, and bad N = len(lut) - 3 + if self._has_exceptional: + lut = self._with_exceptional_colors(lut) + if bytes: + lut = (lut * 255).astype(np.uint8) xa = np.array(x, copy=True) if not xa.dtype.isnative: # Native byteorder is faster. xa = xa.byteswap().view(xa.dtype.newbyteorder()) - if xa.dtype.kind == "f": + is_float = xa.dtype.kind == "f" + if self._has_exceptional and is_float: + # before the scaling below: it overflows large finite values to infinity + # (float16 65504), and those are out of range rather than infinite. + mask_neg_inf = np.isneginf(xa) + mask_pos_inf = np.isposinf(xa) + if is_float: xa *= N # xa == 1 (== N after multiplication) is not out of range. xa[xa == N] = N - 1 @@ -429,11 +492,12 @@ def __call__( # 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) + mask_masked = x.mask # type: ignore + mask_nan = np.isnan(xa) if is_float else False + mask_bad = (mask_masked | mask_nan) if is_float else mask_masked else: - mask_bad = np.isnan(xa) + mask_masked = False + mask_nan = mask_bad = np.isnan(xa) with np.errstate(invalid="ignore"): # We need this cast for unsigned ints as well as floats @@ -442,16 +506,45 @@ def __call__( xa[mask_under] = N xa[mask_over] = N + 1 xa[mask_bad] = N + 2 + if self._has_exceptional: + # last wins: a masked entry is masked whatever value it hides + if is_float: + xa[mask_neg_inf] = N + 3 + xa[mask_pos_inf] = N + 4 + xa[mask_nan] = N + 5 + xa[mask_masked] = N + 6 rgba = lut.take(xa, axis=0, mode="clip") return rgba if np.iterable(x) else Color(rgba) + def _with_exceptional_colors(self, lut: np.ndarray) -> np.ndarray: + """Return `lut` with four rows appended, one per exceptional value class. + + Each appended row falls back to the row its class would otherwise have used, + so routing a class to its own row cannot change any color while that class + has no color of its own. `lut` must be an over/under LUT. + """ + under, over, bad = lut[-3], lut[-2], lut[-1] + return np.vstack( + ( + lut, + under if self.neg_inf_color is None else self.neg_inf_color.rgba, + over if self.pos_inf_color is None else self.pos_inf_color.rgba, + bad if self.nan_color is None else self.nan_color.rgba, + bad if self.masked_color is None else self.masked_color.rgba, + ) + ) + def with_extremes( self, *, bad: ColorLike | None = None, under: ColorLike | None = None, over: ColorLike | None = None, + neg_inf: ColorLike | None = None, + pos_inf: ColorLike | None = None, + nan: ColorLike | None = None, + masked: ColorLike | None = None, ) -> Colormap: """Return a copy of the colormap with new extreme values.""" return type(self)( @@ -462,6 +555,10 @@ def with_extremes( bad=bad, under=under, over=over, + neg_inf=neg_inf, + pos_inf=pos_inf, + nan=nan, + masked=masked, ) def as_dict(self) -> ColormapDict: @@ -609,6 +706,10 @@ def shifted( under=self.under_color, over=self.over_color, bad=self.bad_color, + neg_inf=self.neg_inf_color, + pos_inf=self.pos_inf_color, + nan=self.nan_color, + masked=self.masked_color, ) def to_css( @@ -670,6 +771,10 @@ def __eq__(self, other: object) -> bool: and self.under_color == other.under_color and self.over_color == other.over_color and self.bad_color == other.bad_color + and self.neg_inf_color == other.neg_inf_color + and self.pos_inf_color == other.pos_inf_color + and self.nan_color == other.nan_color + and self.masked_color == other.masked_color and self.interpolation == other.interpolation ) @@ -718,6 +823,20 @@ def _repr_html_(self) -> str: f"over {_html_color_patch(self.over_color)}" "" ) + if self._has_exceptional: + patches = ( + ("neg_inf", self.neg_inf_color), + ("pos_inf", self.pos_inf_color), + ("nan", self.nan_color), + ("masked", self.masked_color), + ) + swatches = " ".join( + f"{name} {_html_color_patch(c)}" for name, c in patches if c is not None + ) + html += ( + '
' + f"{swatches}
" + ) return html diff --git a/src/cmap/_external.py b/src/cmap/_external.py index 68a5bd8ce..23ed001a4 100644 --- a/src/cmap/_external.py +++ b/src/cmap/_external.py @@ -109,8 +109,10 @@ def to_napari(cm: Colormap) -> NapariColormap: kwargs["interpolation"] = ( "zero" if cm.interpolation == "nearest" else "linear" ) - if "nan_color" in param_names and cm.bad_color is not None: - kwargs["nan_color"] = cm.bad_color.rgba + # napari's nan_color covers nan alone, so prefer cmap's nan color over bad + nan_color = cm.nan_color or cm.bad_color + if "nan_color" in param_names and nan_color is not None: + kwargs["nan_color"] = nan_color.rgba if "high_color" in param_names and cm.over_color is not None: kwargs["high_color"] = cm.over_color.rgba if "low_color" in param_names and cm.under_color is not None: diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 41524269e..bad939546 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -187,6 +187,104 @@ def test_colormap_masked_array_with_unmasked_nan() -> None: npt.assert_array_equal(cmap(all_false), cmap(np.array([0.25, np.nan]))) +def test_exceptional_colors() -> None: + cmap = Colormap( + ["red", "blue"], + under="green", + over="yellow", + bad="black", + neg_inf="cyan", + pos_inf="magenta", + nan="white", + masked="orange", + ) + mask = [False] * 6 + [True] * 3 + data = np.ma.masked_array( + [-np.inf, -0.5, 0.5, 1.5, np.inf, np.nan, np.inf, np.nan, 0.25], mask=mask + ) + expect = np.array( + [ + Color("cyan").rgba, # -inf + Color("green").rgba, # under range, finite + Colormap(["red", "blue"])(0.5).rgba, # in range + Color("yellow").rgba, # over range, finite + Color("magenta").rgba, # +inf + Color("white").rgba, # nan, not masked + Color("orange").rgba, # masked wins over +inf + Color("orange").rgba, # masked wins over nan + Color("orange").rgba, # masked + ] + ) + + npt.assert_array_equal(cmap(data), expect) + npt.assert_array_equal(cmap(data, bytes=True), (expect * 255).astype(np.uint8)) + npt.assert_array_equal(data.mask, mask) + + +def test_exceptional_colors_fall_back_to_the_legacy_extremes() -> None: + cmap = Colormap(["red", "blue"], under="green", over="yellow", bad="black") + data = np.ma.masked_array( + [-np.inf, np.inf, np.nan, 0.5], mask=[False, False, False, True] + ) + legacy = [ + Color("green").rgba, # -inf -> under + Color("yellow").rgba, # +inf -> over + Color("black").rgba, # nan -> bad + Color("black").rgba, # masked -> bad + ] + npt.assert_array_equal(cmap(data), legacy) + + # setting one leaves the other three on their legacy destinations + one = cmap.with_extremes(under="green", over="yellow", bad="black", nan="white") + npt.assert_array_equal(one(data), [*legacy[:2], Color("white").rgba, legacy[3]]) + + +def test_finite_values_that_overflow_when_scaled_are_not_infinite() -> None: + # float16 65504 becomes inf once multiplied by N, but it is over-range, not infinite + cmap = Colormap( + ["red", "blue"], under="green", over="yellow", neg_inf="cyan", pos_inf="magenta" + ) + data = np.array([65504, -65504, np.inf, -np.inf], dtype=np.float16) + with np.errstate(over="ignore"): + rgba = cmap(data) + npt.assert_array_equal( + rgba, + [ + Color("yellow").rgba, + Color("green").rgba, + Color("magenta").rgba, + Color("cyan").rgba, + ], + ) + + +def test_masked_dtypes_keep_their_existing_behavior() -> None: + cmap = Colormap(["red", "blue"], bad="black") + bad = Color("black").rgba + + npt.assert_array_equal(cmap(np.ma.masked_array([0, 1], mask=[True, False]))[0], bad) + + obj = np.ma.masked_array(np.array([0.25, 0.5], dtype=object), mask=[True, False]) + npt.assert_array_equal(cmap(obj)[0], bad) + + # object dtype without a mask reaches np.isnan, which has never accepted it + with pytest.raises(TypeError): + cmap(np.ma.masked_array(np.array([0.25], dtype=object), mask=np.ma.nomask)) + + npt.assert_array_equal(cmap(np.ma.masked_array(0.5, mask=True)).rgba, bad) + + +@pytest.mark.parametrize("field", ["neg_inf", "pos_inf", "nan", "masked"]) +def test_exceptional_colors_are_colormap_state(field: str) -> None: + plain = Colormap(["red", "blue"]) + cmap = Colormap(["red", "blue"], **{field: "orange"}) + + assert cmap != plain + assert plain.with_extremes(**{field: "orange"}) == cmap + assert cmap.shifted(1) == cmap + assert field in cmap._repr_html_() + + 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] diff --git a/tests/test_third_party.py b/tests/test_third_party.py index a161dac75..f02cbc4b3 100644 --- a/tests/test_third_party.py +++ b/tests/test_third_party.py @@ -94,6 +94,10 @@ def test_napari_extreme_colors() -> None: np.testing.assert_allclose(ncm.high_color, Color(OVER).rgba) np.testing.assert_allclose(ncm.nan_color, Color(bad).rgba) + # a more specific nan color takes precedence over bad + specific = Colormap(["black", "white"], bad=bad, nan="magenta").to_napari() + np.testing.assert_allclose(specific.nan_color, Color("magenta").rgba) + @pytest.mark.skipif( sys.platform == "darwin" and sys.version_info >= (3, 13),