Keep User Cluster Colours if Possible

stlearn/_datasets/_datasets.py

@@ -34,11 +34,11 @@ def visium_sge(
 
 
 def xenium_sge(
-    base_url: str="https://cf.10xgenomics.com/samples/xenium/1.0.1",
-    library_id: str="Xenium_FFPE_Human_Breast_Cancer_Rep1",
-    zip_filename: str="outs.zip",
-    image_filename: str="he_image.ome.tif",
-    alignment_filename: str="he_imagealignment.csv",
+    base_url: str = "https://cf.10xgenomics.com/samples/xenium/1.0.1",
+    library_id: str = "Xenium_FFPE_Human_Breast_Cancer_Rep1",
+    zip_filename: str = "outs.zip",
+    image_filename: str = "he_image.ome.tif",
+    alignment_filename: str = "he_imagealignment.csv",
     include_hires_tiff: bool = False,
 ):
     """
@@ -59,11 +59,15 @@ def xenium_sge(
 
     if "xe_outs.zip" in zip_filename:
         files_to_extract = [
-            "cell_feature_matrix.zarr.zip", "cells.zarr.zip", "experiment.xenium"
+            "cell_feature_matrix.zarr.zip",
+            "cells.zarr.zip",
+            "experiment.xenium",
         ]
     else:
         files_to_extract = [
-            "cell_feature_matrix.h5", "cells.csv.gz", "experiment.xenium"
+            "cell_feature_matrix.h5",
+            "cells.csv.gz",
+            "experiment.xenium",
         ]
 
     all_sge_files_exist = all(

stlearn/pl/classes.py

@@ -11,6 +11,8 @@
 )
 
 import matplotlib
+import matplotlib.axes
+import matplotlib.figure
 import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
@@ -20,7 +22,14 @@
 
 from ..classes import Spatial
 from ..utils import Axes, _AxesSubplot, _read_graph
-from .utils import centroidpython, check_sublist, get_cluster, get_cmap, get_node
+from .utils import (
+    centroidpython,
+    check_sublist,
+    get_cluster,
+    get_cmap,
+    get_colors,
+    get_node,
+)
 
 
 class SpatialBasePlot(Spatial):
@@ -656,7 +665,7 @@ def __init__(
 
         self.cmap_ = self._get_cmap(self.cmap)
 
-        self._add_cluster_colors()
+        self.colors = get_colors(self.adata[0], self.use_label, cmap=self.cmap)
 
         self._plot_clusters()
 
@@ -687,31 +696,20 @@ def __init__(
         if fname is not None:
             self._save_output()
 
-    def _add_cluster_colors(self):
-        self.adata[0].uns[self.use_label + "_colors"] = []
-
-        for i, _ in enumerate(self.adata[0].obs.groupby(self.use_label, observed=True)):
-            self.adata[0].uns[self.use_label + "_colors"].append(
-                matplotlib.colors.to_hex(self.cmap_(i / (self.cmap_n - 1)))
-            )
-
     def _plot_clusters(self):
         # Plot scatter plot based on pixel of spots
 
-        for i, cluster in enumerate(
+        for _, cluster in enumerate(
             self.query_adata.obs.groupby(self.use_label, observed=True)
         ):
             # Plot scatter plot based on pixel of spots
             subset_spatial = self.query_adata.obsm["spatial"][
                 check_sublist(list(self.query_adata.obs.index), list(cluster[1].index))
             ]
 
-            if self.use_label + "_colors" in self.adata[0].uns:
-                label_set = self.adata[0].obs[self.use_label].cat.categories.values
-                col_index = np.where(label_set == cluster[0])[0][0]
-                color = self.adata[0].uns[self.use_label + "_colors"][col_index]
-            else:
-                color = self.cmap_(self.query_indexes[i] / (self.cmap_n - 1))
+            label_set = self.adata[0].obs[self.use_label].cat.categories.values
+            col_index = np.where(label_set == cluster[0])[0][0]
+            color = self.colors[col_index]
 
             imgcol_new = subset_spatial[:, 0] * self.scale_factor
             imgrow_new = subset_spatial[:, 1] * self.scale_factor
@@ -767,7 +765,7 @@ def _add_cluster_labels(self):
                 x = 100
                 y = -50
 
-            colors = self.adata[0].uns[self.use_label + "_colors"]
+            colors = self.colors
             index = self.query_indexes[i]
             self.ax.text(
                 centroids[0][0] + x,
@@ -844,7 +842,7 @@ def _add_sub_clusters(self):
                     else:
                         x = 100
                         y = -50
-                    colors = self.adata[0].uns[self.use_label + "_colors"]
+                    colors = self.colors
                     index = self.query_indexes[i]
 
                     self.ax.text(
@@ -863,7 +861,7 @@ def _add_sub_clusters(self):
                     )
 
     def _add_trajectories(self):
-        used_colors = self.adata[0].uns[self.use_label + "_colors"]
+        used_colors = self.colors
         cmaps = matplotlib.colors.LinearSegmentedColormap.from_list("", used_colors)
 
         cmap = plt.get_cmap(cmaps)

stlearn/pl/cluster_plot.py

@@ -2,7 +2,8 @@
     Optional,  # Special
 )
 
-import matplotlib
+import matplotlib.axes
+import matplotlib.figure
 from anndata import AnnData
 from bokeh.io import output_notebook
 from bokeh.plotting import show

stlearn/pl/utils.py

@@ -98,23 +98,23 @@ def check_cmap(cmap):
 
 def get_colors(adata, obs_key, cmap="default", label_set=None):
     """Retrieves colors if present in adata.uns, if not present then will set
-    them as per scanpy & return in order requested.
+    them as per scanpy & return in order requested. If fewer colors are present
+    than there are categories, the existing colors are kept and the remainder
+    are generated.
     """
-    # Checking if colors are already set #
     col_key = f"{obs_key}_colors"
-    if col_key in adata.uns:
-        labels_ordered = adata.obs[obs_key].cat.categories
-        colors_ordered = adata.uns[col_key]
-    else:  # Colors not already present
+
+    if not hasattr(adata.obs[obs_key], "cat"):  # Ensure categorical
+        adata.obs[obs_key] = adata.obs[obs_key].astype("category")
+    labels_ordered = adata.obs[obs_key].cat.categories
+
+    colors_ordered = list(adata.uns.get(col_key, []))
+    if len(colors_ordered) < len(labels_ordered):
         check_cmap(cmap)
         cmap, _ = get_cmap(cmap)
-
-        if not hasattr(adata.obs[obs_key], "cat"):  # Ensure categorical
-            adata.obs[obs_key] = adata.obs[obs_key].astype("category")
-        labels_ordered = adata.obs[obs_key].cat.categories
-        colors_ordered = [
-            matplotlib.colors.rgb2hex(cmap(i / (len(labels_ordered) - 1)))
-            for i in range(len(labels_ordered))
+        colors_ordered += [
+            matplotlib.colors.rgb2hex(cmap(i / max(len(labels_ordered) - 1, 1)))
+            for i in range(len(colors_ordered), len(labels_ordered))
         ]
         adata.uns[col_key] = colors_ordered
 

tests/pl/test_cluster_plot.py

@@ -101,9 +101,34 @@ def test_multiple_calls_same_adata(self):
             self.assertEqual(len(colors1), len(colors2))
             self.assertEqual(colors1, colors2)
 
+    def test_existing_colors_preserved(self):
+        """User-supplied colors must not be overwritten (issue: lr_plot ignores
+        adata.uns[f'{use_label}_colors'])."""
+        label_name = "test_clusters"
+        existing_colors = ["#FF0000", "#00FF00", "#0000FF"]
+        self.adata.uns[f"{label_name}_colors"] = existing_colors.copy()
+
+        with (
+            patch("matplotlib.pyplot.subplots") as mock_subplots,
+            patch.object(ClusterPlot, "_plot_clusters"),
+            patch.object(ClusterPlot, "_add_image"),
+        ):
+            mock_subplots.return_value = (MagicMock(), MagicMock())
+
+            plot = ClusterPlot(
+                adata=self.adata,
+                use_label=label_name,
+                show_image=False,
+                show_color_bar=False,
+            )
+
+            self.assertEqual(
+                list(plot.adata[0].uns[f"{label_name}_colors"]), existing_colors
+            )
+
     def test_insufficient_existing_colors_extended(self):
-        """Test that insufficient existing colors are extended."""
-        # Pre-populate adata with insufficient colors (only 2 colors for 3 clusters)
+        """Test that insufficient existing colours are extended."""
+        # Pre-populate adata with insufficient colors (only 2 colours for 3 clusters)
         existing_colors = ["#FF0000", "#00FF00"]
         label_name = "test_clusters"
         self.adata.uns[f"{label_name}_colors"] = existing_colors
@@ -123,10 +148,10 @@ def test_insufficient_existing_colors_extended(self):
                 show_color_bar=False,
             )
 
-            # Should extend existing colors
+            # Should extend existing colours - keep the user's existing colours set.
             colors = plot.adata[0].uns[f"{label_name}_colors"]
             self.assertEqual(len(colors), 3)
-            self.assertNotEqual(colors[:2], existing_colors)
+            self.assertEqual(colors[:2], existing_colors)
 
     def tearDown(self):
         key = f"{self.__class__._label_name}_colors"