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 += ( + '