refactor(render): give CmapParams norm/cmap behavior, drop private _lut poke (#699)

PR1 of the #699 color/cmap/norm unification. Adds behavior to CmapParams so
callers stop re-implementing stateful-object handling:

- fresh_norm(): a safe copy of the norm. Applying a Normalize autoscales
  vmin/vmax in place, so sharing it leaks one element's range into the next.
  Converts the 4 hand-written copy(cmap_params.norm) sites (render_shapes,
  render_points, render_labels-as-points, datashader outline) to use it.
- cmap_with_alpha() / colormap_with_alpha(): a public-API replacement for the
  private cmap._lut[:, -1] = alpha poke in the 1-channel image path. Rebuilds the
  colormap via a lossless N-bin resample (matplotlib quantizes __call__ to N bins,
  so body colors are byte-identical), preserves the cmap name, and reproduces the
  old bad/NaN-pixel alpha behavior. Also stops mutating the shared cmap in place.

Behavior-preserving except the bad-pixel alpha at alpha<1 with NaN pixels, which
is reproduced exactly. Adds unit tests for the new CmapParams methods.

Out of scope (deferred to PR2): unifying _map_color_seg's shared-norm mutation —
the continuous labels colorbar currently depends on that in-place autoscale, so it
must be fixed together with the colorbar's norm derivation.

Author: timtreis

src/spatialdata_plot/pl/_datashader.py

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-from copy import copy
 from typing import Any, Literal
 
 import dask.dataframe as dd
@@ -572,7 +571,7 @@ def _render_ds_outline_by_column(
         )
         # Apply the user-provided norm (vmin/vmax) the same way the fill path does so
         # an explicit Normalize takes effect for the outline cmap.
-        norm = copy(cmap_params.norm)
+        norm = cmap_params.fresh_norm()
         agg_outline, color_span = _apply_ds_norm(agg_outline, norm)
         shaded = ds.tf.shade(
             agg_outline,

src/spatialdata_plot/pl/render.py

@@ -51,6 +51,7 @@
     PointsRenderParams,
     ShapesRenderParams,
     _DsReduction,
+    colormap_with_alpha,
 )
 from spatialdata_plot.pl.utils import (
     _align_outline_vector_to_length,
@@ -741,7 +742,7 @@ def _render_shapes(
 
     color_vector = _maybe_apply_transfunc(color_source_vector, color_vector, render_params.transfunc)
 
-    norm = copy(render_params.cmap_params.norm)
+    norm = render_params.cmap_params.fresh_norm()
 
     if len(color_vector) == 0:
         color_vector = [render_params.cmap_params.na_color.get_hex_with_alpha()]
@@ -1505,7 +1506,7 @@ def _render_points(
 
     trans, trans_data = _prepare_transformation(sdata.points[element], coordinate_system, ax)
 
-    norm = copy(render_params.cmap_params.norm)
+    norm = render_params.cmap_params.fresh_norm()
 
     method = render_params.method
 
@@ -1966,9 +1967,9 @@ def _render_images(
             else render_params.cmap_params.cmap
         )
 
-        # Overwrite alpha in cmap: https://stackoverflow.com/a/10127675
-        cmap._init()
-        cmap._lut[:, -1] = render_params.alpha
+        # Bake a uniform alpha into the cmap via a public-API rebuild (no private `_lut` poke,
+        # no mutation of the shared cmap). The selected cmap may be palette-derived, so apply to it.
+        cmap = colormap_with_alpha(cmap, render_params.alpha, render_params.cmap_params.na_color.get_hex_with_alpha())
 
         # norm needs to be passed directly to ax.imshow(). If we normalize before, that method would always clip.
         _ax_show_and_transform(
@@ -2426,7 +2427,7 @@ def _render_labels(
             y=xy[:, 1],
             color_vector=point_color_vector,
             color_source_vector=point_color_source_vector,
-            norm=copy(render_params.cmap_params.norm),  # ax.scatter autoscales in place; don't mutate the shared norm
+            norm=render_params.cmap_params.fresh_norm(),  # ax.scatter autoscales in place; don't mutate the shared norm
             na_color=na_color,
             adata=table if table_name is not None else None,
             col_for_color=col_for_color,

src/spatialdata_plot/pl/render_params.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Sequence
+from copy import copy
 from dataclasses import dataclass, field
 from typing import Any, Literal
 
 import numpy as np
 from matplotlib.axes import Axes
 from matplotlib.cm import ScalarMappable
-from matplotlib.colors import Colormap, ListedColormap, Normalize, rgb2hex, to_hex
+from matplotlib.colors import Colormap, ListedColormap, Normalize, rgb2hex, to_hex, to_rgba
 from matplotlib.figure import Figure
 
 _FontWeight = Literal["light", "normal", "medium", "semibold", "bold", "heavy", "black"]
@@ -148,6 +149,22 @@ def is_fully_transparent(self) -> bool:
         return self.alpha == "00"
 
 
+def colormap_with_alpha(cmap: Colormap, alpha: float, na_color: str) -> Colormap:
+    """Return ``cmap`` rebuilt with a uniform ``alpha`` and ``na_color`` as the bad/NaN color.
+
+    Body colors are reproduced exactly: matplotlib quantizes ``Colormap.__call__`` into ``N`` bins,
+    so resampling at ``linspace(0, 1, N)`` is lossless. The bad-row alpha is also set to ``alpha`` to
+    match the historical ``cmap._lut[:, -1] = alpha`` poke this replaces.
+    """
+    lut = cmap(np.linspace(0, 1, cmap.N))
+    lut[:, -1] = alpha
+    new = ListedColormap(lut, name=cmap.name)
+    bad = list(to_rgba(na_color))
+    bad[-1] = alpha
+    new.set_bad(bad)
+    return new
+
+
 @dataclass
 class CmapParams:
     """Cmap params."""
@@ -157,6 +174,23 @@ class CmapParams:
     na_color: Color
     cmap_is_default: bool = True
 
+    def fresh_norm(self) -> Normalize:
+        """Return a copy of ``norm`` safe to apply or autoscale without touching the shared one.
+
+        ``Normalize.__call__`` runs ``autoscale_None``, which sets ``vmin``/``vmax`` in place when
+        they are unset. Applying the shared ``norm`` directly therefore leaks the first element's
+        data range into later elements/channels that reuse the same ``CmapParams``.
+        """
+        return copy(self.norm)
+
+    def cmap_with_alpha(self, alpha: float) -> Colormap:
+        """Return ``cmap`` with a uniform ``alpha``, preserving the na/bad color.
+
+        Public-API replacement for poking ``cmap._lut``. Returning a fresh colormap also avoids
+        mutating the shared ``cmap`` in place.
+        """
+        return colormap_with_alpha(self.cmap, alpha, self.na_color.get_hex_with_alpha())
+
 
 @dataclass
 class FigParams:

tests/pl/test_utils.py

@@ -13,7 +13,7 @@
 
 import spatialdata_plot
 from spatialdata_plot.pl import measure_obs
-from spatialdata_plot.pl.render_params import Color, ColorLike
+from spatialdata_plot.pl.render_params import CmapParams, Color, ColorLike
 from spatialdata_plot.pl.utils import (
     _apply_cmap_alpha_to_datashader_result,
     _datashader_map_aggregate_to_color,
@@ -505,9 +505,7 @@ def _add_shapes_table(sdata: SpatialData, element: str = "blobs_polygons", name:
     adata = AnnData(np.zeros((len(gdf), 1), dtype=np.float32))
     adata.obs["instance_id"] = list(gdf.index)
     adata.obs["region"] = element
-    sdata[name] = TableModel.parse(
-        adata, region_key="region", instance_key="instance_id", region=element
-    )
+    sdata[name] = TableModel.parse(adata, region_key="region", instance_key="instance_id", region=element)
     return sdata
 
 
@@ -547,9 +545,7 @@ def test_writes_centroid_area_diameter_for_labels(self, sdata_blobs: SpatialData
         # area is the pixel count (positive integers); diameter = 2*sqrt(area/pi)
         area = table.obs["area"].to_numpy()
         assert (area > 0).all()
-        np.testing.assert_allclose(
-            table.obs["equivalent_diameter"].to_numpy(), 2.0 * np.sqrt(area / np.pi), rtol=1e-12
-        )
+        np.testing.assert_allclose(table.obs["equivalent_diameter"].to_numpy(), 2.0 * np.sqrt(area / np.pi), rtol=1e-12)
 
     def test_writes_for_shapes(self, sdata_blobs: SpatialData) -> None:
         _add_shapes_table(sdata_blobs, "blobs_polygons")
@@ -635,7 +631,7 @@ def test_existing_nonnumeric_column_raises_before_any_write(self, sdata_blobs: S
     def test_sparse_high_label_ids(self, sdata_blobs: SpatialData) -> None:
         # #5: sparse/high label ids (max id >> n_labels) are measured correctly (dense relabelling).
         arr = np.asarray(sdata_blobs["blobs_labels"].data)
-        hi = (arr.astype(np.int64) * 1000)  # ids become 1000, 2000, ... ; max id is huge, few labels
+        hi = arr.astype(np.int64) * 1000  # ids become 1000, 2000, ... ; max id is huge, few labels
         measure_obs(sd_hi := _labels_sdata(hi), "lab", table_name="t")
         measure_obs(sd_lo := _labels_sdata(arr.astype(np.int64)), "lab", table_name="t")
         # relabelling values does not move pixels -> identical centroid set
@@ -814,3 +810,44 @@ def test_show_renders_all_coordinate_systems_for_distributed_elements():
     axes = axes if isinstance(axes, list) else [axes]
     assert {ax.get_title() for ax in axes} == {"cs_a", "cs_b"}
     plt.close("all")
+
+
+class TestCmapParamsMethods:
+    """Unit tests for the behavior added to ``CmapParams`` (#699 PR1)."""
+
+    @staticmethod
+    def _params(cmap_name: str = "viridis", na: Color | None = None) -> CmapParams:
+        from matplotlib import colormaps
+        from matplotlib.colors import Normalize
+
+        return CmapParams(cmap=colormaps[cmap_name], norm=Normalize(vmin=0.0, vmax=1.0), na_color=na or Color())
+
+    def test_fresh_norm_is_independent_copy(self):
+        params = self._params()
+        fresh = params.fresh_norm()
+        assert fresh is not params.norm
+        assert (fresh.vmin, fresh.vmax) == (params.norm.vmin, params.norm.vmax)
+        # autoscaling the copy must not mutate the shared norm (the bug fresh_norm prevents)
+        fresh.autoscale(np.array([5.0, 10.0]))
+        assert (params.norm.vmin, params.norm.vmax) == (0.0, 1.0)
+
+    def test_cmap_with_alpha_preserves_body_and_sets_alpha(self):
+        from matplotlib import colormaps
+
+        params = self._params()
+        out = params.cmap_with_alpha(0.5)
+        # sample at bin centers so quantization is exact for both colormaps
+        xs = (np.arange(params.cmap.N) + 0.5) / params.cmap.N
+        got = out(xs)
+        np.testing.assert_array_equal(got[:, :3], colormaps["viridis"](xs)[:, :3])
+        np.testing.assert_allclose(got[:, 3], 0.5)
+
+    def test_cmap_with_alpha_bad_color_uses_na_with_requested_alpha(self):
+        from matplotlib.colors import to_rgba
+
+        na = Color()  # default lightgray, fully opaque
+        out = self._params(na=na).cmap_with_alpha(0.25)
+        bad = out(np.nan)
+        np.testing.assert_allclose(bad[:3], to_rgba(na.get_hex_with_alpha())[:3], atol=1e-6)
+        # historical `_lut[:, -1] = alpha` overwrote the bad-row alpha too
+        assert bad[3] == 0.25