feat: colors for +inf, -inf, NaN, and masked values

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)}"
                 "</div>"
             )
+        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 += (
+                '<div style="vertical-align: middle; max-width: 514px;">'
+                f"{swatches}</div>"
+            )
 
         return html
 

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:

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]

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),