feat(render_images): opt-in PCA multichannel_strategy for n>=4 channels (#450)

Adds multichannel_strategy={'stack','pca'}|None. Default None->'stack'
(byte-identical to current additive blending); 'pca' is opt-in. PCA reduces the
per-channel-normalized stack to 3 components, percentile-stretches each, blends
with chosen colors (default RGB; palette of 3 or 3 cmap samples), and scales by
total-signal brightness so background fades to black (no segmentation).
svd_solver='full' keeps colors deterministic across runs and sklearn versions.
scikit-learn already a runtime dep.

Author: timtreis

src/spatialdata_plot/pl/_validate.py

@@ -1345,6 +1345,7 @@ def _validate_image_render_params(
     scale: str | None,
     colorbar: bool | str | None,
     colorbar_params: dict[str, object] | None,
+    multichannel_strategy: str | None = None,
 ) -> dict[str, dict[str, Any]]:
     param_dict: dict[str, Any] = {
         "sdata": sdata,
@@ -1397,17 +1398,25 @@ def _validate_image_render_params(
         assert isinstance(palette, list | type(None))  # if present, was converted to list, just to make sure
 
         if isinstance(palette, list):
-            # case A: single palette for all channels
-            if len(palette) == 1:
-                palette_length = len(channel) if channel is not None else len(spatial_element_ch)
-                palette = palette * palette_length
-            # case B: one palette per channel (either given or derived from channel length)
-            channels_to_use = spatial_element_ch if element_params[el]["channel"] is None else channel
-            if channels_to_use is not None and len(palette) != len(channels_to_use):
-                raise ValueError(
-                    f"Palette length ({len(palette)}) does not match channel length "
-                    f"({', '.join(str(c) for c in channels_to_use)})."
-                )
+            if multichannel_strategy == "pca":
+                # palette colors the 3 principal components, not the channels
+                if len(palette) != 3:
+                    raise ValueError(
+                        "multichannel_strategy='pca' maps the 3 principal components to colors; "
+                        f"'palette' must have exactly 3 entries, got {len(palette)}."
+                    )
+            else:
+                # case A: single palette for all channels
+                if len(palette) == 1:
+                    palette_length = len(channel) if channel is not None else len(spatial_element_ch)
+                    palette = palette * palette_length
+                # case B: one palette per channel (either given or derived from channel length)
+                channels_to_use = spatial_element_ch if element_params[el]["channel"] is None else channel
+                if channels_to_use is not None and len(palette) != len(channels_to_use):
+                    raise ValueError(
+                        f"Palette length ({len(palette)}) does not match channel length "
+                        f"({', '.join(str(c) for c in channels_to_use)})."
+                    )
         element_params[el]["palette"] = palette
 
         expected_len = len(channel) if channel is not None else len(spatial_element_ch)

src/spatialdata_plot/pl/basic.py

@@ -740,6 +740,7 @@ def render_images(
         colorbar: bool | str | None = "auto",
         colorbar_params: dict[str, object] | None = None,
         channels_as_legend: bool = False,
+        multichannel_strategy: Literal["stack", "pca"] | None = None,
         method: Literal["matplotlib", "datashader"] | None = None,
         datashader_reduction: _ImageDsReduction | None = None,
     ) -> sd.SpatialData:
@@ -822,6 +823,20 @@ def render_images(
             Ignored for single-channel and RGB(A) images.  When multiple
             ``render_images`` calls use this flag on the same axes, all
             channel entries are combined into a single legend.
+        multichannel_strategy : {"stack", "pca"} | None, default None
+            How to composite an image with more than three channels.
+            ``None`` (default) and ``"stack"`` both use additive blending of
+            per-channel colors (today's behavior, unchanged). ``"pca"`` reduces
+            the per-channel-normalized stack to three principal components mapped
+            to RGB, with per-pixel brightness taken from the total signal so
+            low-signal background fades to black (no segmentation) — useful for
+            highly multiplexed assays (CODEX, IMC, Xenium morphology) where
+            additive blending saturates to a muddy white. The three components
+            map to red/green/blue by default; pass ``palette`` (exactly three
+            colors) or ``cmap`` (three colors are sampled from it) to recolor
+            them. It is an exploratory overview (colors are abstract, not
+            per-channel), so ``channels_as_legend`` is ignored with a warning,
+            and ``"pca"`` requires at least three channels.
         method : str | None, optional
             Whether to use ``'matplotlib'`` (default) or ``'datashader'`` for
             the downsampling step.  When ``'datashader'`` is selected, the
@@ -882,6 +897,7 @@ def render_images(
             scale=scale,
             colorbar=colorbar,
             colorbar_params=colorbar_params,
+            multichannel_strategy=multichannel_strategy,
         )
 
         sdata = self._copy()
@@ -936,6 +952,7 @@ def render_images(
                 transfunc=transfunc,
                 grayscale=grayscale,
                 channels_as_legend=channels_as_legend,
+                multichannel_strategy=multichannel_strategy,
                 method=method,
                 ds_reduction=datashader_reduction,
             )

src/spatialdata_plot/pl/render.py

@@ -1756,6 +1756,69 @@ def _draw_channel_legend(
     )
 
 
+# Explained-variance ratio below which a component is treated as noise (rank-deficient input).
+_PCA_NOISE_FLOOR = 1e-6
+# Percentile clip applied per component (and to the brightness) before rescaling to [0, 1].
+_PCA_STRETCH_PERCENTILES = (1.0, 99.0)
+_PCA_DEFAULT_COLORS = np.eye(3)  # PC1/2/3 -> red/green/blue
+
+
+def _stretch_to_unit(values: np.ndarray) -> np.ndarray:
+    """Clip to ``_PCA_STRETCH_PERCENTILES`` and rescale to ``[0, 1]`` (flat input -> zeros)."""
+    lo, hi = np.percentile(values, _PCA_STRETCH_PERCENTILES)
+    if hi <= lo:
+        return np.zeros_like(values)
+    return np.clip((values - lo) / (hi - lo), 0.0, 1.0)
+
+
+def _pca_compose(
+    layers: dict[Any, np.ndarray], channels: list[Any], colors: np.ndarray | None = None
+) -> tuple[np.ndarray, np.ndarray]:
+    """Reduce a per-channel-normalized multichannel stack to RGB via PCA.
+
+    Channels are treated as features and reduced to up to three principal components; each is
+    percentile-stretched to ``[0, 1]``, additively blended with its ``colors`` row (a ``(3, 3)``
+    RGB array, default red/green/blue), then scaled by per-pixel total signal so background fades
+    to black (no segmentation). Components below ``_PCA_NOISE_FLOOR`` explained variance are
+    dropped (no color) with a warning. ``svd_solver="full"`` plus a per-component max-abs sign
+    convention keep colors deterministic across runs and sklearn versions. Returns
+    ``((y, x, 3) in [0, 1], explained_variance_ratio)``.
+    """
+    from sklearn.decomposition import PCA
+
+    if colors is None:
+        colors = _PCA_DEFAULT_COLORS
+
+    # float32 halves peak memory on the (h*w, n_channels) matrix and everything derived from it;
+    # the percentile stretch and 8-bit display are insensitive to the lost precision.
+    h, w = np.asarray(layers[channels[0]]).shape
+    mat = np.stack([np.asarray(layers[ch], dtype=np.float32).reshape(-1) for ch in channels], axis=1)
+
+    pca = PCA(n_components=min(3, mat.shape[1]), svd_solver="full")
+    comps = pca.fit_transform(mat)
+
+    rgb = np.zeros((h * w, 3), dtype=np.float32)
+    dropped: list[int] = []
+    for k in range(comps.shape[1]):
+        if pca.explained_variance_ratio_[k] < _PCA_NOISE_FLOOR:
+            dropped.append(k)
+            continue
+        col = comps[:, k]
+        if col[np.argmax(np.abs(col))] < 0:  # pin PCA's arbitrary sign
+            col = -col
+        rgb += np.outer(_stretch_to_unit(col), colors[k])
+    np.clip(rgb, 0.0, 1.0, out=rgb)
+
+    if dropped:
+        logger.warning(
+            f"PCA: component(s) {dropped} below the noise floor (explained variance < "
+            f"{_PCA_NOISE_FLOOR:g}); channels appear collinear, so they contribute no color."
+        )
+
+    rgb *= _stretch_to_unit(mat.sum(axis=1))[:, np.newaxis]  # brightness from total signal
+    return rgb.reshape(h, w, 3), pca.explained_variance_ratio_
+
+
 def _render_images(
     sdata: sd.SpatialData,
     render_params: ImageRenderParams,
@@ -2049,6 +2112,40 @@ def _render_images(
         # Colors for the channel legend (set by each branch if applicable)
         legend_colors: list[str] | None = None
 
+        # 2-PCA) Opt-in PCA reduction to RGB for highly multiplexed images, bypassing the
+        # additive seed-color blending that saturates at 4+ channels.
+        if render_params.multichannel_strategy == "pca":
+            if n_channels < 3:
+                raise ValueError(
+                    f"multichannel_strategy='pca' requires at least 3 channels, got {n_channels}. "
+                    "Select 3+ channels via 'channel' or use the default 'stack' strategy."
+                )
+
+            # Color the 3 components by `palette` (exactly 3, validated upstream), 3 cmap
+            # samples, or (default, left as None) red/green/blue inside _pca_compose.
+            component_colors = None
+            if palette is not None:
+                component_colors = np.array([matplotlib.colors.to_rgb(c) for c in palette])
+            elif has_explicit_cmap or user_supplied_multi_cmaps:
+                cmap_obj = (
+                    render_params.cmap_params[0].cmap
+                    if isinstance(render_params.cmap_params, list)
+                    else render_params.cmap_params.cmap
+                )
+                component_colors = np.array([cmap_obj(x)[:3] for x in (0.0, 0.5, 1.0)])
+
+            colored, explained_variance = _pca_compose(layers, channels, component_colors)
+            logger.info(
+                "PCA explained-variance ratio per component: "
+                f"{', '.join(f'{v:.3f}' for v in explained_variance)}."
+            )
+            _ax_show_and_transform(
+                colored, trans_data, ax, render_params.alpha, zorder=render_params.zorder, interpolation=_interp
+            )
+            if render_params.channels_as_legend:
+                logger.warning("channels_as_legend is ignored with multichannel_strategy='pca'.")
+            return
+
         # 2A) Image has 3 channels, no palette info, and no/only one cmap was given
         if palette is None and n_channels == 3 and not isinstance(render_params.cmap_params, list):
             if render_params.cmap_params.cmap_is_default:  # -> use RGB
@@ -2129,8 +2226,9 @@ def _render_images(
                 colored = np.clip(comp_rgb, 0, 1)
                 logger.info(
                     f"Your image has {n_channels} channels. Sampling categorical colors and using "
-                    f"multichannel strategy 'stack' to render."
-                )  # TODO: update when pca is added as strategy
+                    f"multichannel strategy 'stack' to render (pass multichannel_strategy='pca' for a "
+                    f"PCA overview)."
+                )
 
             legend_colors = seed_colors
 

src/spatialdata_plot/pl/render_params.py

@@ -308,6 +308,7 @@ class ImageRenderParams(RenderParams):
     transfunc: Callable[[np.ndarray], np.ndarray] | list[Callable[[np.ndarray], np.ndarray]] | None = None
     grayscale: bool = False
     channels_as_legend: bool = False
+    multichannel_strategy: Literal["stack", "pca"] | None = None
     method: Literal["matplotlib", "datashader"] | None = None
     ds_reduction: _ImageDsReduction | None = None
 

tests/pl/test_render_images.py

@@ -927,3 +927,83 @@ def _render_and_grab(**kwargs):
                 plt.close(fig)
 
         np.testing.assert_array_equal(_render_and_grab(), _render_and_grab(method="matplotlib"))
+
+
+class TestMultichannelPCA(PlotTester, metaclass=PlotTesterMeta):
+    """PCA compositing strategy for n >= 4 channel images (#450)."""
+
+    @staticmethod
+    def _render_array(sdata: SpatialData, **kwargs) -> np.ndarray:
+        fig, ax = plt.subplots(figsize=(2, 2), dpi=50)
+        try:
+            sdata.pl.render_images("blobs_image", **kwargs).pl.show(ax=ax)
+            return np.asarray(ax.get_images()[-1].get_array())
+        finally:
+            plt.close(fig)
+
+    def test_plot_pca_strategy_5_channels(self, sdata_blobs_str: SpatialData):
+        sdata_blobs_str.pl.render_images("blobs_image", multichannel_strategy="pca").pl.show()
+
+    def test_pca_renders_rgb_in_unit_range(self, sdata_blobs_str: SpatialData):
+        arr = self._render_array(sdata_blobs_str, multichannel_strategy="pca")
+        assert arr.ndim == 3 and arr.shape[-1] == 3
+        assert arr.min() >= 0.0 and arr.max() <= 1.0
+
+    def test_pca_is_deterministic(self, sdata_blobs_str: SpatialData):
+        # PCA's SVD sign ambiguity is pinned by the max-abs sign convention -> stable colors.
+        first = self._render_array(sdata_blobs_str, multichannel_strategy="pca")
+        second = self._render_array(sdata_blobs_str, multichannel_strategy="pca")
+        np.testing.assert_array_equal(first, second)
+
+    def test_pca_below_3_channels_raises(self, sdata_blobs_str: SpatialData):
+        with pytest.raises(ValueError, match="at least 3 channels"):
+            sdata_blobs_str.pl.render_images(
+                "blobs_image", channel=["c1", "c2"], multichannel_strategy="pca"
+            ).pl.show()
+        plt.close("all")
+
+    def test_pca_rank_deficient_drops_components_and_warns(self, caplog):
+        # Three identical channels -> rank 1: the surviving component must still render (not
+        # all-black), the two noise-floor components must be dropped (no color) and reported.
+        rng = np.random.default_rng(0)
+        base = rng.random((64, 64), dtype=np.float32)
+        arr = np.stack([base, base, base], axis=0)
+        sdata = SpatialData(images={"blobs_image": Image2DModel.parse(arr, c_coords=["c1", "c2", "c3"])})
+        with logger_warns(caplog, logger, match="below the noise floor"):
+            out = self._render_array(sdata, multichannel_strategy="pca")
+        assert out.shape[-1] == 3
+        assert out[..., 0].max() > 0  # dominant component rendered
+        assert np.allclose(out[..., 1], 0.0) and np.allclose(out[..., 2], 0.0)  # dropped -> no color
+
+    def test_pca_low_signal_background_renders_dark(self):
+        # Signal in the centre, zeros elsewhere: the total-signal brightness gate must pull the
+        # empty border toward black (the non-segmenting background fix).
+        rng = np.random.default_rng(0)
+        arr = np.zeros((4, 64, 64), dtype=np.float32)
+        arr[:, 24:40, 24:40] = rng.random((4, 16, 16), dtype=np.float32) + 0.5
+        sdata = SpatialData(images={"blobs_image": Image2DModel.parse(arr, c_coords=["c1", "c2", "c3", "c4"])})
+        out = self._render_array(sdata, multichannel_strategy="pca")
+        border = np.concatenate([out[:8].ravel(), out[-8:].ravel()])  # the empty margin
+        assert border.mean() < 0.05
+        assert out[24:40, 24:40].mean() > border.mean()
+
+    @pytest.mark.parametrize("recolor", [{"palette": ["cyan", "magenta", "yellow"]}, {"cmap": "viridis"}])
+    def test_pca_recolors_components(self, sdata_blobs_str: SpatialData, recolor):
+        # palette (3 colors) and cmap (3 samples) must change the output vs the default RGB mapping.
+        default = self._render_array(sdata_blobs_str, multichannel_strategy="pca")
+        recolored = self._render_array(sdata_blobs_str, multichannel_strategy="pca", **recolor)
+        assert not np.array_equal(default, recolored)
+
+    def test_pca_palette_wrong_length_raises(self, sdata_blobs_str: SpatialData):
+        with pytest.raises(ValueError, match="must have exactly 3 entries"):
+            sdata_blobs_str.pl.render_images(
+                "blobs_image", multichannel_strategy="pca", palette=["red", "green"]
+            ).pl.show()
+        plt.close("all")
+
+    def test_stack_default_is_backward_compatible(self, sdata_blobs_str: SpatialData):
+        # None (default) and explicit "stack" must be byte-identical to the pre-#450 output.
+        np.testing.assert_array_equal(
+            self._render_array(sdata_blobs_str, multichannel_strategy=None),
+            self._render_array(sdata_blobs_str, multichannel_strategy="stack"),
+        )