Replace NaN-masking with RGBA post-processing for cmap alpha

The NaN-masking approach was fundamentally flawed: masking values to NaN
caused datashader to auto-scale the remaining values, shifting other
shapes' colors and potentially making them white.

New approach: let datashader shade normally (preserving all colors and
spans), then post-process the RGBA output to apply the cmap's alpha
channel. This:
- Does not modify aggregate values or spans
- Preserves exact colors for all non-transparent shapes
- Works for any cmap with transparent entries at any position
- Handles span=None correctly (no auto-scaling issues)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timtreis

src/spatialdata_plot/pl/utils.py

@@ -3194,58 +3194,75 @@ def _prepare_transformation(
     return trans, trans_data
 
 
-def _mask_transparent_cmap_entries(
+def _apply_cmap_alpha_to_datashader_result(
+    result: Any,
     agg: DataArray,
     cmap: str | list[str] | Colormap,
     span: list[float] | tuple[float, float] | None,
-) -> tuple[DataArray, list[float] | tuple[float, float] | None]:
-    """Set aggregate values to NaN where the colormap alpha is zero.
+) -> Any:
+    """Apply the colormap's alpha channel to a datashader RGBA result.
 
     Datashader ignores the per-entry alpha channel of matplotlib colormaps,
-    so entries meant to be fully transparent (alpha=0) must be converted to
-    NaN — which datashader already renders as transparent — before shading.
-    See :issue:`376`.
-
-    Returns ``(masked_agg, span)`` where *span* is the data range computed
-    **before** masking.  The caller must pass this span to ``ds.tf.shade``
-    to prevent datashader from auto-scaling the remaining values into the
-    transparent bin.
+    so pixels that the cmap marks as transparent (alpha=0) are rendered
+    opaque.  This function post-processes the shaded RGBA output to restore
+    the cmap's intended transparency.  See :issue:`376`.
     """
     if not isinstance(cmap, Colormap):
-        return agg, span
+        return result
+
+    # Quick check: does this cmap have any transparent entries?
+    test_vals = np.linspace(0, 1, min(cmap.N, 256))
+    cmap_alphas = cmap(test_vals)[:, 3]
+    if np.all(cmap_alphas >= 1.0):
+        return result
+
+    # Get or ensure we have an (H, W, 4) uint8 array
+    if hasattr(result, "values"):
+        # datashader Image — uint32 packed, convert via to_numpy()
+        rgba = result.to_numpy().base
+        if rgba is None:
+            return result
+    else:
+        rgba = result
 
-    # Only the bottom entry is checked — this is what
-    # ``set_zero_in_cmap_to_transparent`` targets.
-    if cmap(0.0)[3] >= 1.0:
-        return agg, span
+    if rgba.ndim != 3 or rgba.shape[2] != 4:
+        return result
 
-    # For a ListedColormap with N colors, index 0 covers [0, 1/N) in
-    # normalised space. Compute the data value at the upper boundary of that bin.
-    frac = 1.0 / cmap.N
+    # Normalise aggregate values to [0, 1] using the same span datashader used
+    agg_vals = agg.values.astype(np.float64)
+    valid = np.isfinite(agg_vals)
+    if not valid.any():
+        return result
 
     if span is not None:
         lo, hi = float(span[0]), float(span[1])
     else:
-        lo = float(agg.min())
-        hi = float(agg.max())
+        lo = float(np.nanmin(agg_vals))
+        hi = float(np.nanmax(agg_vals))
 
     if hi <= lo or not np.isfinite(lo) or not np.isfinite(hi):
-        return agg, span
+        return result
+
+    normed = np.clip((agg_vals - lo) / (hi - lo), 0.0, 1.0)
 
-    # Freeze the span before masking so datashader doesn't auto-scale
-    # the remaining values into the now-transparent first bin.
-    span = [lo, hi]
+    # Look up cmap alpha for each pixel
+    desired_alpha = cmap(normed)[:, :, 3]
 
-    threshold = lo + frac * (hi - lo)
-    return agg.where(agg >= threshold), span
+    # Zero out pixels where the cmap wants transparency
+    transparent = valid & (desired_alpha < 1.0)
+    if transparent.any():
+        # Scale the existing alpha by the cmap's alpha
+        rgba[transparent, 3] = (rgba[transparent, 3].astype(np.float32) * desired_alpha[transparent]).astype(np.uint8)
+
+    return result
 
 
 def _datashader_map_aggregate_to_color(
     agg: DataArray,
     cmap: str | list[str] | ListedColormap,
     color_key: list[str] | dict[str, str] | None = None,
     min_alpha: float = 40,
-    span: list[float] | tuple[float, float] | None = None,
+    span: None | list[float] = None,
     clip: bool = True,
 ) -> ds.tf.Image | np.ndarray[Any, np.dtype[np.uint8]]:
     """ds.tf.shade() part, ensuring correct clipping behavior.
@@ -3257,8 +3274,6 @@ def _datashader_map_aggregate_to_color(
         # in case we use datashader together with a Normalize object where clip=False
         # why we need this is documented in https://github.com/scverse/spatialdata-plot/issues/372
         agg_in = agg.where((agg >= span[0]) & (agg <= span[1]))
-        agg_in, span = _mask_transparent_cmap_entries(agg_in, cmap, span)
-        assert span is not None  # guaranteed: we passed a non-None span above
         img_in = ds.tf.shade(
             agg_in,
             cmap=cmap,
@@ -3293,17 +3308,18 @@ def _datashader_map_aggregate_to_color(
         img_over = img_over.to_numpy().base
         if img_over is not None:
             stack[stack[:, :, 3] == 0] = img_over[stack[:, :, 3] == 0]
-        return stack
 
-    agg, span = _mask_transparent_cmap_entries(agg, cmap, span)
-    return ds.tf.shade(
+        return _apply_cmap_alpha_to_datashader_result(stack, agg, cmap, span)
+
+    result = ds.tf.shade(
         agg,
         cmap=cmap,
         color_key=color_key,
         min_alpha=min_alpha,
         span=span,
         how="linear",
     )
+    return _apply_cmap_alpha_to_datashader_result(result, agg, cmap, span)
 
 
 def _hex_no_alpha(hex: str) -> str:

tests/pl/test_utils.py

@@ -9,9 +9,9 @@
 
 import spatialdata_plot
 from spatialdata_plot.pl.utils import (
+    _apply_cmap_alpha_to_datashader_result,
     _datashader_map_aggregate_to_color,
     _get_subplots,
-    _mask_transparent_cmap_entries,
     set_zero_in_cmap_to_transparent,
 )
 from tests.conftest import DPI, PlotTester, PlotTesterMeta
@@ -117,74 +117,77 @@ def test_is_color_like(color_result: tuple[ColorLike, bool]):
     assert spatialdata_plot.pl.utils._is_color_like(color) == result
 
 
-class TestMaskTransparentCmapEntries:
+class TestCmapAlphaDatashader:
     """Regression tests for #376: set_zero_in_cmap_to_transparent with datashader."""
 
-    def test_masks_zero_values_when_cmap_has_transparent_entry(self):
+    def test_transparent_pixels_get_alpha_zero(self):
+        """Post-processing sets alpha=0 for pixels mapping to transparent cmap entries."""
+        import datashader as ds
+
         cmap = set_zero_in_cmap_to_transparent("viridis")
-        data = np.array([[0.0, 1.0, 5.0], [0.0, 2.0, 10.0]])
+        data = np.array([[0.0, 5.0, 10.0]], dtype=np.float64)
         agg = xr.DataArray(data, dims=["y", "x"])
 
-        masked, returned_span = _mask_transparent_cmap_entries(agg, cmap, span=[0.0, 10.0])
+        shaded = ds.tf.shade(agg, cmap=cmap, min_alpha=254, how="linear")
+        result = _apply_cmap_alpha_to_datashader_result(shaded, agg, cmap, span=[0.0, 10.0])
+        rgba = result.to_numpy().base if hasattr(result, "to_numpy") else result
+
+        assert rgba[0, 0, 3] == 0, f"Expected alpha=0 at value=0.0, got {rgba[0, 0, 3]}"
+        assert rgba[0, 1, 3] > 0, "Expected non-zero alpha at value=5.0"
+        assert rgba[0, 2, 3] > 0, "Expected non-zero alpha at value=10.0"
 
-        assert np.isnan(masked.values[0, 0])
-        assert np.isnan(masked.values[1, 0])
-        assert masked.values[0, 1] == 1.0
-        assert masked.values[0, 2] == 5.0
-        assert returned_span == [0.0, 10.0]
+    def test_opaque_cmap_unchanged(self):
+        """Post-processing is a no-op for fully opaque cmaps."""
+        import datashader as ds
 
-    def test_no_effect_for_opaque_cmap(self):
         cmap = plt.get_cmap("viridis")
-        data = np.array([[0.0, 5.0, 10.0]])
+        data = np.array([[0.0, 5.0, 10.0]], dtype=np.float64)
         agg = xr.DataArray(data, dims=["y", "x"])
 
-        masked, returned_span = _mask_transparent_cmap_entries(agg, cmap, span=[0.0, 10.0])
-        np.testing.assert_array_equal(masked.values, data)
+        shaded = ds.tf.shade(agg, cmap=cmap, min_alpha=254, how="linear")
+        rgba_before = shaded.to_numpy().base.copy()
+        result = _apply_cmap_alpha_to_datashader_result(shaded, agg, cmap, span=[0.0, 10.0])
+        rgba_after = result.to_numpy().base if hasattr(result, "to_numpy") else result
+        np.testing.assert_array_equal(rgba_before, rgba_after)
 
-    def test_no_effect_for_string_cmap(self):
-        data = np.array([[0.0, 5.0, 10.0]])
+    def test_string_cmap_passthrough(self):
+        """Post-processing is a no-op for string cmaps (early return)."""
+        dummy_rgba = np.zeros((2, 3, 4), dtype=np.uint8)
+        dummy_rgba[:, :, 3] = 200
+        data = np.array([[0.0, 5.0, 10.0]], dtype=np.float64)
         agg = xr.DataArray(data, dims=["y", "x"])
 
-        masked, _ = _mask_transparent_cmap_entries(agg, "viridis", span=[0.0, 10.0])
-        np.testing.assert_array_equal(masked.values, data)
+        result = _apply_cmap_alpha_to_datashader_result(dummy_rgba, agg, "viridis", span=[0.0, 10.0])
+        np.testing.assert_array_equal(result, dummy_rgba)
 
-    def test_datashader_shade_respects_transparent_cmap(self):
-        """End-to-end: _datashader_map_aggregate_to_color produces alpha=0 for transparent cmap entries."""
+    def test_end_to_end_datashader_map(self):
+        """_datashader_map_aggregate_to_color produces alpha=0 for transparent cmap entries."""
         cmap = set_zero_in_cmap_to_transparent("viridis")
         data = np.array([[0.0, 5.0, 10.0]], dtype=np.float64)
         agg = xr.DataArray(data, dims=["y", "x"])
 
         result = _datashader_map_aggregate_to_color(agg, cmap=cmap, min_alpha=254, span=[0.0, 10.0])
-        img = result.values if hasattr(result, "values") else result
-
-        alpha_at_zero = (int(img[0, 0]) >> 24) & 0xFF
-        alpha_at_five = (int(img[0, 1]) >> 24) & 0xFF
+        img = result.to_numpy().base if hasattr(result, "to_numpy") else result
 
-        assert alpha_at_zero == 0, f"Expected alpha=0 at value=0.0, got {alpha_at_zero}"
-        assert alpha_at_five > 0, f"Expected non-zero alpha at value=5.0, got {alpha_at_five}"
+        assert img[0, 0, 3] == 0, f"Expected alpha=0 at value=0.0, got {img[0, 0, 3]}"
+        assert img[0, 1, 3] > 0, "Expected non-zero alpha at value=5.0"
 
-    def test_span_none_with_zeros(self):
-        """Masking works when span is inferred from the aggregate (span=None)."""
+    def test_span_none_preserves_colors(self):
+        """With span=None, non-transparent shapes keep their correct colors."""
         cmap = set_zero_in_cmap_to_transparent("viridis")
-        data = np.array([[0.0, 3.0, 10.0]])
+        data = np.array([[0.0, 5.0, 10.0]], dtype=np.float64)
         agg = xr.DataArray(data, dims=["y", "x"])
 
-        masked, returned_span = _mask_transparent_cmap_entries(agg, cmap, span=None)
-
-        assert np.isnan(masked.values[0, 0])
-        assert masked.values[0, 1] == 3.0
-        assert masked.values[0, 2] == 10.0
-        assert returned_span == [0.0, 10.0], "span should be frozen from pre-masking range"
-
-    def test_all_nan_aggregate(self):
-        """All-NaN aggregate is returned unchanged."""
-
-        cmap = set_zero_in_cmap_to_transparent("viridis")
-        data = np.array([[np.nan, np.nan]])
-        agg = xr.DataArray(data, dims=["y", "x"])
+        result = _datashader_map_aggregate_to_color(agg, cmap=cmap, min_alpha=254)
+        img = result.to_numpy().base if hasattr(result, "to_numpy") else result
 
-        masked, _ = _mask_transparent_cmap_entries(agg, cmap, span=None)
-        np.testing.assert_array_equal(np.isnan(masked.values), np.isnan(data))
+        # value=0 should be transparent
+        assert img[0, 0, 3] == 0
+        # value=5 and value=10 should be opaque with correct viridis colors (not white)
+        assert img[0, 1, 3] > 0
+        assert img[0, 2, 3] > 0
+        # The non-transparent pixels should NOT be white (R=255,G=255,B=255)
+        assert not (img[0, 1, 0] == 255 and img[0, 1, 1] == 255 and img[0, 1, 2] == 255)
 
 
 def test_extract_scalar_value():