color bar logic

Author: timtreis

src/spatialdata_plot/pl/basic.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import contextlib
 import sys
+import warnings
 from collections import OrderedDict
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -17,6 +19,7 @@
 from matplotlib.axes import Axes
 from matplotlib.colors import Colormap, Normalize
 from matplotlib.figure import Figure
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 from spatialdata import get_extent
 from spatialdata._utils import _deprecation_alias
 from xarray import DataArray, DataTree
@@ -28,9 +31,11 @@
     _render_labels,
     _render_points,
     _render_shapes,
+    _split_colorbar_params,
 )
 from spatialdata_plot.pl.render_params import (
     CmapParams,
+    ColorbarSpec,
     ImageRenderParams,
     LabelsRenderParams,
     LegendParams,
@@ -172,6 +177,8 @@ def render_shapes(
         table_name: str | None = None,
         table_layer: str | None = None,
         shape: Literal["circle", "hex", "visium_hex", "square"] | None = None,
+        colorbar: bool | str | None = "auto",
+        colorbar_params: dict[str, object] | None = None,
         **kwargs: Any,
     ) -> sd.SpatialData:
         """
@@ -237,6 +244,11 @@ def render_shapes(
         method : str | None, optional
             Whether to use 'matplotlib' and 'datashader'. When None, the method is
             chosen based on the size of the data.
+        colorbar :
+            Whether to request a colorbar for continuous colors. Use "auto" (default) for automatic selection.
+        colorbar_params :
+            Parameters forwarded to Matplotlib's colorbar alongside layout hints such as ``loc``, ``width``, ``pad``,
+            and ``label``.
         table_name: str | None
             Name of the table containing the color(s) columns. If one name is given than the table is used for each
             spatial element to be plotted if the table annotates it. If you want to use different tables for particular
@@ -292,6 +304,8 @@ def render_shapes(
             shape=shape,
             method=method,
             ds_reduction=kwargs.get("datashader_reduction"),
+            colorbar=colorbar,
+            colorbar_params=colorbar_params,
         )
 
         sdata = self._copy()
@@ -326,6 +340,8 @@ def render_shapes(
                 zorder=n_steps,
                 method=param_values["method"],
                 ds_reduction=param_values["ds_reduction"],
+                colorbar=param_values["colorbar"],
+                colorbar_params=param_values["colorbar_params"],
             )
             n_steps += 1
 
@@ -347,6 +363,8 @@ def render_points(
         method: str | None = None,
         table_name: str | None = None,
         table_layer: str | None = None,
+        colorbar: bool | str | None = "auto",
+        colorbar_params: dict[str, object] | None = None,
         **kwargs: Any,
     ) -> sd.SpatialData:
         """
@@ -396,6 +414,11 @@ def render_points(
         method : str | None, optional
             Whether to use 'matplotlib' and 'datashader'. When None, the method is
             chosen based on the size of the data.
+        colorbar :
+            Whether to request a colorbar for continuous colors. Use "auto" (default) for automatic selection.
+        colorbar_params :
+            Parameters forwarded to Matplotlib's colorbar alongside layout hints such as ``loc``, ``width``, ``pad``,
+            and ``label``.
         table_name: str | None
             Name of the table containing the color(s) columns. If one name is given than the table is used for each
             spatial element to be plotted if the table annotates it. If you want to use different tables for particular
@@ -434,6 +457,8 @@ def render_points(
             table_name=table_name,
             table_layer=table_layer,
             ds_reduction=kwargs.get("datashader_reduction"),
+            colorbar=colorbar,
+            colorbar_params=colorbar_params,
         )
 
         if method is not None:
@@ -467,6 +492,8 @@ def render_points(
                 zorder=n_steps,
                 method=method,
                 ds_reduction=param_values["ds_reduction"],
+                colorbar=param_values["colorbar"],
+                colorbar_params=param_values["colorbar_params"],
             )
             n_steps += 1
 
@@ -484,6 +511,8 @@ def render_images(
         palette: list[str] | str | None = None,
         alpha: float | int = 1.0,
         scale: str | None = None,
+        colorbar: bool | str | None = "auto",
+        colorbar_params: dict[str, object] | None = None,
         **kwargs: Any,
     ) -> sd.SpatialData:
         """
@@ -526,6 +555,11 @@ def render_images(
                 3) "full": Renders the full image without rasterization. In the case of
                 multiscale images, the highest resolution scale is selected. Note that
                 this may result in long computing times for large images.
+        colorbar :
+            Whether to request a colorbar for continuous colors. Use "auto" (default) for automatic selection.
+        colorbar_params :
+            Parameters forwarded to Matplotlib's colorbar alongside layout hints such as ``loc``, ``width``, ``pad``,
+            and ``label``.
         kwargs
             Additional arguments to be passed to cmap, norm, and other rendering functions.
 
@@ -547,6 +581,8 @@ def render_images(
             cmap=cmap,
             norm=norm,
             scale=scale,
+            colorbar=colorbar,
+            colorbar_params=colorbar_params,
         )
 
         sdata = self._copy()
@@ -580,6 +616,8 @@ def render_images(
                 alpha=param_values["alpha"],
                 scale=param_values["scale"],
                 zorder=n_steps,
+                colorbar=param_values["colorbar"],
+                colorbar_params=param_values["colorbar_params"],
             )
             n_steps += 1
 
@@ -600,6 +638,8 @@ def render_labels(
         outline_alpha: float | int = 0.0,
         fill_alpha: float | int = 0.4,
         scale: str | None = None,
+        colorbar: bool | str | None = "auto",
+        colorbar_params: dict[str, object] | None = None,
         table_name: str | None = None,
         table_layer: str | None = None,
         **kwargs: Any,
@@ -653,6 +693,11 @@ def render_labels(
                 (exception: a dpi is specified in `show()`. Then the image is rasterized to fit the canvas and dpi).
                 3) "full": render the full image without rasterization. In the case of a multiscale image, the scale
                 with the highest resolution is selected. This can lead to long computing times for large images!
+        colorbar :
+            Whether to request a colorbar for continuous colors. Use "auto" (default) for automatic selection.
+        colorbar_params :
+            Parameters forwarded to Matplotlib's colorbar alongside layout hints such as ``loc``, ``width``, ``pad``,
+            and ``label``.
         table_name: str | None
             Name of the table containing the color columns.
         table_layer: str | None
@@ -681,6 +726,8 @@ def render_labels(
             outline_alpha=outline_alpha,
             palette=palette,
             scale=scale,
+            colorbar=colorbar,
+            colorbar_params=colorbar_params,
             table_name=table_name,
             table_layer=table_layer,
         )
@@ -709,6 +756,8 @@ def render_labels(
                 table_name=param_values["table_name"],
                 table_layer=param_values["table_layer"],
                 zorder=n_steps,
+                colorbar=param_values["colorbar"],
+                colorbar_params=param_values["colorbar_params"],
             )
             n_steps += 1
         return sdata
@@ -722,7 +771,8 @@ def show(
         legend_loc: str | None = "right margin",
         legend_fontoutline: int | None = None,
         na_in_legend: bool = True,
-        colorbar: bool = True,
+        colorbar: bool | None = None,
+        colorbar_params: dict[str, object] | None = None,
         wspace: float | None = None,
         hspace: float = 0.25,
         ncols: int = 4,
@@ -761,7 +811,11 @@ def show(
         return_ax :
             Whether to return the axes object created. False by default.
         colorbar :
-            Whether to plot the colorbar. True by default.
+            DEPRECATED: use per-layer ``colorbar``/``colorbar_params`` instead. If provided, it will be honored
+            but will emit a DeprecationWarning.
+        colorbar_params :
+            Global overrides passed to colorbars for all axes. Accepts the same keys as per-layer ``colorbar_params``
+            (e.g., ``loc``, ``width``, ``pad``, ``label``).
         title :
             The title of the plot. If not provided the plot will have the name of the coordinate system as title.
 
@@ -786,6 +840,7 @@ def show(
             legend_fontoutline,
             na_in_legend,
             colorbar,
+            colorbar_params,
             wspace,
             hspace,
             ncols,
@@ -888,15 +943,89 @@ def show(
             ncols=ncols,
             frameon=frameon,
         )
+        # colorbar is deprecated: warn when explicitly provided
+        legend_colorbar = True if colorbar is None else colorbar
+        if colorbar is not None:
+            warnings.warn(
+                "Parameter 'colorbar' in `.show()` is deprecated. "
+                "Please control colorbars via the per-layer `colorbar` and `colorbar_params` arguments.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         legend_params = LegendParams(
             legend_fontsize=legend_fontsize,
             legend_fontweight=legend_fontweight,
             legend_loc=legend_loc,
             legend_fontoutline=legend_fontoutline,
             na_in_legend=na_in_legend,
-            colorbar=colorbar,
+            colorbar=legend_colorbar,
         )
 
+        def _draw_colorbar(spec: ColorbarSpec, location_offsets: dict[str, float]) -> None:
+            base_layout = {"location": "right", "fraction": 0.08, "pad": 0.01}
+            layer_layout, layer_kwargs, layer_label_override = _split_colorbar_params(spec.params)
+            global_layout, global_kwargs, global_label_override = _split_colorbar_params(colorbar_params)
+            layout = {**base_layout, **layer_layout, **global_layout}
+            cbar_kwargs = {**layer_kwargs, **global_kwargs}
+
+            location = cast(str, layout.get("location", base_layout["location"]))
+            default_orientation = "vertical" if location in {"right", "left"} else "horizontal"
+            cbar_kwargs.setdefault("orientation", default_orientation)
+
+            fraction = float(cast(float | int, layout.get("fraction", base_layout["fraction"])))
+            pad = float(cast(float | int, layout.get("pad", base_layout["pad"])))
+            offset = location_offsets.get(location, 0.0)
+            pad = pad + offset
+            # update offset for the next bar on the same side (space consumed = pad + fraction)
+            location_offsets[location] = offset + pad + fraction
+
+            if location == "right":
+                bbox = [1 + pad, 0, fraction, 1]
+            elif location == "left":
+                bbox = [-(pad + fraction), 0, fraction, 1]
+            elif location == "top":
+                bbox = [0, 1 + pad, 1, fraction]
+            elif location == "bottom":
+                bbox = [0, -(pad + fraction), 1, fraction]
+            else:
+                bbox = [1 + pad, 0, fraction, 1]
+
+            cax = inset_axes(
+                spec.ax,
+                width="100%",
+                height="100%",
+                loc="center",
+                bbox_to_anchor=bbox,
+                bbox_transform=spec.ax.transAxes,
+                borderpad=0.0,
+            )
+
+            cb = fig_params.fig.colorbar(spec.mappable, cax=cax, **cbar_kwargs)
+            if location == "left":
+                cb.ax.yaxis.set_ticks_position("left")
+                cb.ax.yaxis.set_label_position("left")
+                cb.ax.tick_params(labelleft=True, labelright=False)
+            elif location == "top":
+                cb.ax.xaxis.set_ticks_position("top")
+                cb.ax.xaxis.set_label_position("top")
+                cb.ax.tick_params(labeltop=True, labelbottom=False)
+            elif location == "right":
+                cb.ax.yaxis.set_ticks_position("right")
+                cb.ax.yaxis.set_label_position("right")
+                cb.ax.tick_params(labelright=True, labelleft=False)
+            elif location == "bottom":
+                cb.ax.xaxis.set_ticks_position("bottom")
+                cb.ax.xaxis.set_label_position("bottom")
+                cb.ax.tick_params(labelbottom=True, labeltop=False)
+
+            final_label = global_label_override or layer_label_override or spec.label
+            if final_label:
+                cb.set_label(final_label)
+            if spec.alpha is not None:
+                with contextlib.suppress(Exception):
+                    cb.solids.set_alpha(spec.alpha)
+
         cs_contents = _get_cs_contents(sdata)
 
         # go through tree
@@ -908,6 +1037,8 @@ def show(
             )
             ax = fig_params.ax if fig_params.axs is None else fig_params.axs[i]
             assert isinstance(ax, Axes)
+            axis_colorbar_requests: list[ColorbarSpec] | None = [] if legend_params.colorbar else None
+            location_offsets: dict[str, float] = {"left": 0.0, "right": 0.0, "top": 0.0, "bottom": 0.0}
 
             wants_images = False
             wants_labels = False
@@ -937,6 +1068,7 @@ def show(
                             fig_params=fig_params,
                             scalebar_params=scalebar_params,
                             legend_params=legend_params,
+                            colorbar_requests=axis_colorbar_requests,
                             rasterize=rasterize,
                         )
 
@@ -954,6 +1086,7 @@ def show(
                             fig_params=fig_params,
                             scalebar_params=scalebar_params,
                             legend_params=legend_params,
+                            colorbar_requests=axis_colorbar_requests,
                         )
 
                 elif cmd == "render_points" and has_points:
@@ -970,6 +1103,7 @@ def show(
                             fig_params=fig_params,
                             scalebar_params=scalebar_params,
                             legend_params=legend_params,
+                            colorbar_requests=axis_colorbar_requests,
                         )
 
                 elif cmd == "render_labels" and has_labels:
@@ -978,7 +1112,8 @@ def show(
                     )
 
                     if wanted_labels_on_this_cs:
-                        if (table := params_copy.table_name) is not None:
+                        table = params_copy.table_name
+                        if table is not None:
                             assert isinstance(params_copy.color, str)
                             colors = sc.get.obs_df(sdata[table], [params_copy.color])
                             if isinstance(colors[params_copy.color].dtype, pd.CategoricalDtype):
@@ -1002,6 +1137,7 @@ def show(
                             fig_params=fig_params,
                             scalebar_params=scalebar_params,
                             legend_params=legend_params,
+                            colorbar_requests=axis_colorbar_requests,
                             rasterize=rasterize,
                         )
 
@@ -1038,6 +1174,27 @@ def show(
                 ax.set_xlim(x_min, x_max)
                 ax.set_ylim(y_max, y_min)  # (0, 0) is top-left
 
+            if legend_params.colorbar and axis_colorbar_requests:
+                # keep only unique bars per (location, label, layout+kwargs). Last request wins.
+                unique_specs_map: dict[tuple[Any, ...], ColorbarSpec] = {}
+                for spec in axis_colorbar_requests:
+                    layer_layout, layer_kwargs, layer_label_override = _split_colorbar_params(spec.params)
+                    global_layout, global_kwargs, global_label_override = _split_colorbar_params(colorbar_params)
+                    layout = {**{"location": "right"}, **layer_layout, **global_layout}
+                    kwargs_key = tuple(sorted({**layer_kwargs, **global_kwargs}.items()))
+                    label_key = global_label_override or layer_label_override or spec.label
+                    key = (
+                        layout.get("location", "right"),
+                        label_key,
+                        tuple(sorted(layout.items())),
+                        kwargs_key,
+                    )
+                    unique_specs_map[key] = spec
+                unique_specs = list(unique_specs_map.values())
+
+                for spec in unique_specs:
+                    _draw_colorbar(spec, location_offsets)
+
         if fig_params.fig is not None and save is not None:
             save_fig(fig_params.fig, path=save)