Merge branch 'main' into fix/issue-449

Author: timtreis

src/spatialdata_plot/pl/_datashader.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from copy import copy
 from typing import Any, Literal
 
 import dask.dataframe as dd
@@ -19,13 +20,15 @@
 from spatialdata_plot._logging import logger
 from spatialdata_plot.pl.render_params import Color, FigParams, ShapesRenderParams, _DsReduction
 from spatialdata_plot.pl.utils import (
+    _DS_REDUCTION_FUNCS,
     _ax_show_and_transform,
     _convert_alpha_to_datashader_range,
     _create_image_from_datashader_result,
     _datashader_aggregate_with_function,
     _datashader_map_aggregate_to_color,
     _datshader_get_how_kw_for_spread,
     _hex_no_alpha,
+    _make_continuous_mappable,
 )
 
 # ---------------------------------------------------------------------------
@@ -36,6 +39,11 @@
 # missing (NaN) values.  Must not collide with realistic user category names.
 _DS_NAN_CATEGORY = "ds_nan"
 
+# Private column name under which the outline color vector is attached to the
+# datashader rasterizer element. Must not collide with a real user column;
+# the leading/trailing dunders are deliberate.
+_OUTLINE_INTERNAL_COL = "__sdp_outline_col__"
+
 # ---------------------------------------------------------------------------
 # Low-level helpers
 # ---------------------------------------------------------------------------
@@ -225,6 +233,7 @@ def _ds_shade_continuous(
     na_color_hex: str,
     spread_px: int | None = None,
     ds_reduction: _DsReduction | None = None,
+    how: str = "linear",
 ) -> tuple[Any, Any | None, tuple[Any, Any] | None]:
     """Shade a continuous datashader aggregate, optionally applying spread and NaN coloring.
 
@@ -253,6 +262,7 @@ def _ds_shade_continuous(
         min_alpha=_convert_alpha_to_datashader_range(alpha),
         span=color_span,
         clip=norm.clip,
+        how=how,
     )
     shaded = _apply_user_alpha(shaded, alpha)
 
@@ -276,6 +286,8 @@ def _ds_shade_categorical(
     color_vector: Any,
     alpha: float,
     spread_px: int | None = None,
+    how: str = "linear",
+    density: bool = False,
 ) -> Any:
     """Shade a categorical or no-color datashader aggregate."""
     ds_cmap = None
@@ -284,12 +296,20 @@ def _ds_shade_categorical(
         if isinstance(ds_cmap, str) and ds_cmap[0] == "#":
             ds_cmap = _hex_no_alpha(ds_cmap)
 
+    # The default min_alpha (~254) is a near-full-opacity floor — right for scatter
+    # plots, but it collapses the count-driven alpha range and makes categorical
+    # density read as a flat hue cloud. Drop the floor under density so per-pixel
+    # alpha can actually encode count. A small non-zero floor (~15%) keeps the
+    # sparse edges visible under density_how="linear" instead of vanishing.
+    min_alpha = 40.0 if density else _convert_alpha_to_datashader_range(alpha)
+
     agg_to_shade = ds.tf.spread(agg, px=spread_px) if spread_px is not None else agg
     shaded = _datashader_map_aggregate_to_color(
         agg_to_shade,
         cmap=ds_cmap,
         color_key=color_key,
-        min_alpha=_convert_alpha_to_datashader_range(alpha),
+        min_alpha=min_alpha,
+        how=how,
     )
     return _apply_user_alpha(shaded, alpha)
 
@@ -330,8 +350,16 @@ def _render_ds_outlines(
     factor: float,
     x_min: float = 0.0,
     y_min: float = 0.0,
+    outline_color_vector: Any | None = None,
+    outline_color_source_vector: pd.Series | None = None,
 ) -> None:
-    """Aggregate, shade, and render shape outlines (outer and inner) with datashader."""
+    """Aggregate, shade, and render shape outlines (outer and inner) with datashader.
+
+    When ``outline_color_vector`` is provided, the outer outline is colored per-shape
+    via ``ds.by`` (categorical) or a numeric reduction (continuous) instead of a
+    single literal color. The two-outline form is rejected at validation, so this
+    only affects the outer outline.
+    """
     ds_lw_factor = fig_params.fig.dpi / 72
     assert len(render_params.outline_alpha) == 2  # noqa: S101
 
@@ -344,6 +372,24 @@ def _render_ds_outlines(
         alpha = render_params.outline_alpha[idx]
         if alpha <= 0:
             continue
+        if idx == 0 and outline_color_vector is not None:
+            _render_ds_outline_by_column(
+                cvs=cvs,
+                transformed_element=transformed_element,
+                outline_color_vector=outline_color_vector,
+                outline_color_source_vector=outline_color_source_vector,
+                cmap_params=render_params.cmap_params,
+                ds_reduction=render_params.ds_reduction,
+                line_width=linewidth * ds_lw_factor,
+                alpha=alpha,
+                fig_params=fig_params,
+                ax=ax,
+                factor=factor,
+                x_min=x_min,
+                y_min=y_min,
+                zorder=render_params.zorder,
+            )
+            continue
         agg_outline = cvs.line(
             transformed_element,
             geometry="geometry",
@@ -361,6 +407,95 @@ def _render_ds_outlines(
             _ax_show_and_transform(rgba, trans, ax, zorder=render_params.zorder)
 
 
+def _render_ds_outline_by_column(
+    cvs: Any,
+    transformed_element: Any,
+    outline_color_vector: Any | None,
+    outline_color_source_vector: pd.Series | None,
+    cmap_params: Any,
+    ds_reduction: _DsReduction | None,
+    line_width: float,
+    alpha: float,
+    fig_params: FigParams,
+    ax: matplotlib.axes.SubplotBase,
+    factor: float,
+    x_min: float,
+    y_min: float,
+    zorder: int,
+) -> None:
+    """Aggregate + shade an outline colored by an obs column via datashader.
+
+    Two-outline form is not supported for column-driven outline coloring,
+    so this only renders the outer outline.
+    """
+    color_by_categorical = outline_color_source_vector is not None
+    na_color_hex = _hex_no_alpha(cmap_params.na_color.get_hex())
+
+    # Attach the outline vector under a private column name so a fill column with the
+    # same key never gets overwritten. Assign positionally (via a Series indexed to the
+    # element) — `.assign(col=series)` aligns by index, which silently inserts NaN when
+    # the element's index is non-contiguous (e.g. after an inner-join). The NaNs would
+    # then be lifted to the `ds_nan` sentinel and one polygon's outline would render as
+    # `na_color` instead of its real category.
+    transformed_element = transformed_element.copy()
+    if color_by_categorical:
+        cat = pd.Categorical(outline_color_source_vector)
+        attach_cat = _inject_ds_nan_sentinel(pd.Series(cat))
+        transformed_element[_OUTLINE_INTERNAL_COL] = pd.Categorical(
+            attach_cat.to_numpy(), categories=attach_cat.cat.categories
+        )
+    else:
+        transformed_element[_OUTLINE_INTERNAL_COL] = np.asarray(outline_color_vector)
+
+    if color_by_categorical:
+        agg_outline = cvs.line(
+            transformed_element,
+            geometry="geometry",
+            agg=ds.by(_OUTLINE_INTERNAL_COL, ds.count()),
+            line_width=line_width,
+        )
+        color_key = _build_datashader_color_key(
+            _coerce_categorical_source(transformed_element[_OUTLINE_INTERNAL_COL]),
+            outline_color_vector,
+            na_color_hex,
+        )
+        shaded = ds.tf.shade(
+            agg_outline,
+            color_key=color_key,
+            min_alpha=_convert_alpha_to_datashader_range(alpha),
+            how="linear",
+        )
+    else:
+        reduction_name = ds_reduction if ds_reduction is not None else "max"
+        try:
+            reduction_function = _DS_REDUCTION_FUNCS[reduction_name](column=_OUTLINE_INTERNAL_COL)
+        except KeyError as e:
+            raise ValueError(
+                f"Reduction '{reduction_name}' is not supported. Use one of: {', '.join(_DS_REDUCTION_FUNCS.keys())}."
+            ) from e
+        agg_outline = cvs.line(
+            transformed_element,
+            geometry="geometry",
+            agg=reduction_function,
+            line_width=line_width,
+        )
+        # 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)
+        agg_outline, color_span = _apply_ds_norm(agg_outline, norm)
+        shaded = ds.tf.shade(
+            agg_outline,
+            cmap=cmap_params.cmap,
+            span=color_span,
+            min_alpha=_convert_alpha_to_datashader_range(alpha),
+            how="linear",
+        )
+
+    shaded = _apply_user_alpha(shaded, alpha)
+    rgba, trans = _create_image_from_datashader_result(shaded, factor, ax, x_min, y_min)
+    _ax_show_and_transform(rgba, trans, ax, zorder=zorder)
+
+
 def _build_ds_colorbar(
     reduction_bounds: tuple[Any, Any] | None,
     norm: Normalize,
@@ -374,12 +509,4 @@ def _build_ds_colorbar(
         return None
     vmin = reduction_bounds[0].values if norm.vmin is None else norm.vmin
     vmax = reduction_bounds[1].values if norm.vmax is None else norm.vmax
-    if (norm.vmin is not None or norm.vmax is not None) and norm.vmin == norm.vmax:
-        assert norm.vmin is not None
-        assert norm.vmax is not None
-        vmin = norm.vmin - 0.5
-        vmax = norm.vmin + 0.5
-    return ScalarMappable(
-        norm=matplotlib.colors.Normalize(vmin=vmin, vmax=vmax),
-        cmap=cmap,
-    )
+    return _make_continuous_mappable(vmin, vmax, cmap)

src/spatialdata_plot/pl/basic.py

@@ -243,13 +243,18 @@ def render_shapes(
             Width of the border. If 2 values are given (tuple), 2 borders are shown with these widths (outer & inner).
             If `outline_color` and/or `outline_alpha` are used to indicate that one/two outlines should be drawn, the
             default outline widths 1.5 and 0.5 are used for outer/only and inner outline respectively.
-        outline_color : ColorLike | tuple[ColorLike], optional
+        outline_color : ColorLike | tuple[ColorLike] | str, optional
             Color of the border. Can either be a named color ("red"), a hex representation ("#000000") or a list of
             floats that represent RGB/RGBA values (1.0, 0.0, 0.0, 1.0). If the hex representation includes alpha, e.g.
             "#000000ff", and `outline_alpha` is not given, this value controls the opacity of the outline. If 2 values
             are given (tuple), 2 borders are shown with these colors (outer & inner). If `outline_width` and/or
             `outline_alpha` are used to indicate that one/two outlines should be drawn, the default outline colors
             "#000000" and "#ffffff are used for outer/only and inner outline respectively.
+            A string that is not a recognized color is interpreted as a column key (in `obs` of the annotating table
+            or in the element's own dataframe), mirroring how ``color`` is parsed. The outline is then colored
+            per-shape using the same ``palette`` / ``cmap`` / ``na_color`` as the fill. When both ``color`` and
+            ``outline_color`` resolve to columns, two stacked legends are drawn. Column-based outline coloring is
+            only supported for a single outline (not the 2-tuple form).
         outline_alpha : float | int | tuple[float | int, float | int] | None, optional
             Alpha value for the outline of shapes. Invisible by default, meaning outline_alpha=0.0 if both outline_color
             and outline_width are not specified. Else, outlines are rendered with the alpha implied by outline_color, or
@@ -346,6 +351,8 @@ def render_shapes(
                 element=element,
                 color=param_values["color"],
                 col_for_color=param_values["col_for_color"],
+                col_for_outline_color=param_values["col_for_outline_color"],
+                outline_table_name=param_values["outline_table_name"],
                 groups=param_values["groups"],
                 scale=param_values["scale"],
                 outline_params=outline_params,
@@ -387,6 +394,8 @@ def render_points(
         colorbar: bool | str | None = "auto",
         colorbar_params: dict[str, object] | None = None,
         datashader_reduction: _DsReduction | None = None,
+        density: bool = False,
+        density_how: Literal["linear", "log", "cbrt", "eq_hist"] = "linear",
         transfunc: Callable[[float], float] | None = None,
     ) -> sd.SpatialData:
         """
@@ -457,13 +466,38 @@ def render_points(
             in another column of ``var``. Mimics scanpy's ``gene_symbols`` parameter.
         datashader_reduction : Literal["sum", "mean", "any", "count", "std", "var", "max", "min"] | None, optional
             Reduction method for datashader when coloring by continuous values. When ``None``, defaults to ``"sum"``.
+        density : bool, default False
+            Render the points as a 2-D count density via datashader instead of plotting individual markers.
+            When ``True``, ``method`` is forced to ``"datashader"`` (passing ``method="matplotlib"`` raises).
+            Density supports ``color=None`` (plain density) or a categorical ``color`` column (per-category
+            density via :func:`datashader.by`). A continuous ``color`` column or a literal color value
+            (e.g. ``"red"``) raises an error. Under ``density=True`` the following parameters are ignored
+            (with a warning if explicitly set): ``size``, ``transfunc``, ``norm.vmin/vmax``, and
+            ``datashader_reduction``.
+        density_how : Literal["linear", "log", "cbrt", "eq_hist"], default "linear"
+            How datashader maps aggregated counts to color intensity. ``"linear"`` (default) keeps the
+            colorbar axis as a count; ``"log"`` and ``"cbrt"`` compress dynamic range; ``"eq_hist"``
+            equalizes the histogram (rank-based, surfaces the most structure but the colorbar axis is
+            no longer a count). Ignored when ``density=False``.
         transfunc : Callable[[float], float] | None, optional
             Optional transformation applied to the continuous color vector before normalization and colormap mapping.
 
         Returns
         -------
         sd.SpatialData
             A copy of the SpatialData object with the rendering parameters stored in its plotting tree.
+
+        Examples
+        --------
+        Plain density of all transcripts:
+
+        >>> sdata.pl.render_points("transcripts", density=True).pl.show()
+
+        Per-gene density with a categorical palette:
+
+        >>> sdata.pl.render_points(
+        ...     "transcripts", color="gene", groups=["Gad1", "Slc17a7"], palette="tab20", density=True
+        ... ).pl.show()
         """
         params_dict = _validate_points_render_params(
             self._sdata,
@@ -482,6 +516,10 @@ def render_points(
             colorbar=colorbar,
             colorbar_params=colorbar_params,
             gene_symbols=gene_symbols,
+            density=density,
+            density_how=density_how,
+            transfunc=transfunc,
+            method=method,
         )
 
         if method is not None:
@@ -490,6 +528,9 @@ def render_points(
             if method not in ["matplotlib", "datashader"]:
                 raise ValueError("Parameter 'method' must be either 'matplotlib' or 'datashader'.")
 
+        if density and method is None:
+            method = "datashader"
+
         sdata = self._copy()
         sdata = _verify_plotting_tree(sdata)
         n_steps = len(sdata.plotting_tree.keys())
@@ -517,6 +558,8 @@ def render_points(
                 ds_reduction=param_values["ds_reduction"],
                 colorbar=param_values["colorbar"],
                 colorbar_params=param_values["colorbar_params"],
+                density=density,
+                density_how=density_how,
             )
             n_steps += 1
 
@@ -811,11 +854,14 @@ def render_labels(
         fill_alpha : float | int | None, optional
             Alpha value for the fill of the labels. By default, it is set to 0.4 or, if a color is given that implies
             an alpha, that value is used for `fill_alpha`.
-        outline_color : ColorLike | None
+        outline_color : ColorLike | str | None
             Color of the outline of the labels. Can either be a named color ("red"), a hex representation
             ("#000000") or a list of floats that represent RGB/RGBA values (1.0, 0.0, 0.0, 1.0). If ``None``,
             the outline inherits from the ``color`` parameter when it is a literal color, or uses data-driven
             per-label colors when ``color`` refers to a column.
+            A string that is not a recognized color is interpreted as a column key (in `obs` of the annotating
+            table), mirroring how ``color`` is parsed. The outline is then colored per-label using the same
+            ``palette`` / ``cmap`` / ``na_color`` as the fill.
         scale :  str | None
             Influences the resolution of the rendering. Possibilities for setting this parameter:
                 1) None (default). The image is rasterized to fit the canvas size. For multiscale images, the best scale
@@ -881,6 +927,8 @@ def render_labels(
                 element=element,
                 color=param_values["color"],
                 col_for_color=param_values["col_for_color"],
+                col_for_outline_color=param_values["col_for_outline_color"],
+                outline_table_name=param_values["outline_table_name"],
                 groups=param_values["groups"],
                 contour_px=param_values["contour_px"],
                 cmap_params=cmap_params,
@@ -1047,6 +1095,8 @@ def show(
         na_in_legend: bool = True,
         colorbar: bool = True,
         colorbar_params: dict[str, object] | None = None,
+        legend_title: str | None = None,
+        outline_legend_title: str | None = None,
         wspace: float | None = None,
         hspace: float = 0.25,
         ncols: int = 4,
@@ -1091,6 +1141,12 @@ def show(
         colorbar_params : dict[str, object] | None
             Global overrides passed to colorbars for all axes. Accepts the same keys as per-layer ``colorbar_params``
             (e.g., ``loc``, ``width``, ``pad``, ``label``).
+        legend_title : str | None
+            Title for the fill categorical legend. When both fill and outline are colored by an obs column, the
+            two legends default to ``"fill"`` / ``"outline"`` to disambiguate; pass an explicit string to override
+            the fill title. Set to ``None`` (default) to keep the auto-title behavior.
+        outline_legend_title : str | None
+            Title for the outline categorical legend. Mirrors ``legend_title`` for the outline channel.
         wspace : float | None
             Horizontal spacing between panels (passed to :class:`matplotlib.gridspec.GridSpec`).
         hspace : float, default 0.25
@@ -1327,6 +1383,8 @@ def show(
             legend_fontoutline=legend_fontoutline,
             na_in_legend=na_in_legend,
             colorbar=colorbar,
+            legend_title=legend_title,
+            outline_legend_title=outline_legend_title,
         )
 
         def _draw_colorbar(