another update

Author: giovp

src/spatialdata_plot/pl/_categorical_utils.py

@@ -498,6 +498,8 @@ def _get_colors_for_categorical_obs(categories: Sequence[Union[str, int]]) -> li
 
 
 Palette_t = Optional[Union[str, ListedColormap]]
+from matplotlib.colors import ListedColormap, to_hex, to_rgba
+import matplotlib.pyplot as plt
 
 
 def _get_palette(
@@ -509,15 +511,15 @@ def _get_palette(
 ) -> Mapping[str, str] | None:
     if palette is None:
         try:
-            palette = adata.uns[Key.uns.colors(cluster_key)]  # type: ignore[arg-type]
+            palette = adata.uns[f"{cluster_key}_colors"]  # type: ignore[arg-type]
             if len(palette) != len(categories):
                 raise ValueError(
                     f"Expected palette to be of length `{len(categories)}`, found `{len(palette)}`. "
                     + f"Removing the colors in `adata.uns` with `adata.uns.pop('{cluster_key}_colors')` may help."
                 )
             return {cat: to_hex(to_rgba(col)[:3] + (alpha,), keep_alpha=True) for cat, col in zip(categories, palette)}
         except KeyError as e:
-            logg.error(f"Unable to fetch palette, reason: {e}. Using `None`.")
+            print(e)
             return None
 
     len_cat = len(categories)
@@ -530,3 +532,27 @@ def _get_palette(
         raise TypeError(f"Palette is {type(palette)} but should be string or `ListedColormap`.")
 
     return dict(zip(categories, palette))
+
+
+def _maybe_set_colors(
+    source: AnnData, target: AnnData, key: str, palette: str | ListedColormap | Cycler | Sequence[Any] | None = None
+) -> None:
+    color_key = f"{key}_colors"
+    from scanpy.plotting._utils import add_colors_for_categorical_sample_annotation
+
+    # this is insane, basically the version copied here was from napari
+    # in napari is modified because we have to do some tricks to plot the categorical values
+    # hence it requires the argument vec. But here we don't, so am re importing the original one
+    # from scanpy here.
+    # this is a testament to how broken the categorical color handling is in the scanpy ecosystem and
+    # to the fact that, because I've never fixed it, an embarassing amount of intellectual debt has
+    # been accumulated.
+
+    try:
+        if palette is not None:
+            raise KeyError("Unable to copy the palette when there was other explicitly specified.")
+        target.uns[color_key] = source.uns[color_key]
+    except KeyError:
+        if isinstance(palette, ListedColormap):  # `scanpy` requires it
+            palette = cycler(color=palette.colors)
+        add_colors_for_categorical_sample_annotation(target, key=key, force_update_colors=True, palette=palette)

src/spatialdata_plot/pl/basic.py

@@ -45,6 +45,8 @@
     _get_subplots,
 )
 
+from spatialdata_plot.pl._categorical_utils import _maybe_set_colors
+
 
 @register_spatial_data_accessor("pl")
 class PlotAccessor:
@@ -793,6 +795,9 @@ def show(
                                 vec=_get_color_key_values(sdata, params["color_key"]),
                                 palette=params["palette"],
                             )
+                            _maybe_set_colors(
+                                source=sdata.table, target=sdata.table, key=params["color_key"], palette=None
+                            )
                     else:
                         # If any of the previous conditions are not met, generate random
                         # colors for each cell id

src/spatialdata_plot/pl/render.py

@@ -30,7 +30,7 @@
 from skimage.segmentation import find_boundaries
 from skimage.util import map_array
 from functools import partial
-from ..pl._categorical_utils import _get_colors_for_categorical_obs, _get_palette
+from ..pl._categorical_utils import _get_colors_for_categorical_obs, _get_palette, _maybe_set_colors
 from ..pl.utils import _normalize
 from ..pp.utils import _get_linear_colormap, _get_region_key
 
@@ -221,8 +221,10 @@ def _render_images(
 
     ax.set_title(key)
 
+
 import matplotlib.pyplot as plt
 
+
 def _render_labels(
     sdata: sd.SpatialData,
     params: dict[str, Union[str, int, float]],
@@ -234,7 +236,6 @@ def _render_labels(
 
     # subset table to only the entires specified by 'key'
     table = sdata.table[sdata.table.obs[region_key] == key]
-
     segmentation = sdata.labels[key].values
 
     norm = Normalize(vmin=None, vmax=None)
@@ -271,7 +272,20 @@ def _render_labels(
 
     #     ax.legend(handles=patches, bbox_to_anchor=(0.9, 0.9), loc="upper left", frameon=False)
     # ax.colorbar(pad=0.01, fraction=0.08, aspect=30)
-    plt.colorbar(cax, ax=ax, pad=0.01, fraction=0.08, aspect=30)
+    if is_categorical_dtype(color_source_vector):
+        clusters = color_source_vector.categories
+        palette = _get_palette(table, cluster_key=params["color_key"], categories=clusters)
+        for label in clusters:
+            ax.scatter([], [], c=palette[label], label=label)
+        ax.legend(
+            frameon=False,
+            loc="center left",
+            bbox_to_anchor=(1, 0.5),
+            ncol=(1 if len(clusters) <= 14 else 2 if len(clusters) <= 30 else 3),
+            fontsize=None,
+        )
+    else:
+        plt.colorbar(cax, ax=ax, pad=0.01, fraction=0.08, aspect=30)
     ax.set_title(key)