Coloring labels by a continuous variable fixed (#344)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: timtreis

CHANGELOG.md

@@ -8,6 +8,22 @@ and this project adheres to [Semantic Versioning][].
 [keep a changelog]: https://keepachangelog.com/en/1.0.0/
 [semantic versioning]: https://semver.org/spec/v2.0.0.html
 
+## [0.2.6] - tbd
+
+### Added
+
+-
+
+### Changed
+
+-   Lowered RMSE-threshold for plot-based tests from 45 to 15 (#344)
+-   When subsetting to `groups`, `NA` isn't automatically added to legend (#344)
+
+### Fixed
+
+-   Filtering with `groups` now preserves original cmap (#344)
+-   Non-selected `groups` are now not shown in `na_color` (#344)
+
 ## [0.2.5] - 2024-08-23
 
 ### Added

src/spatialdata_plot/pl/render.py

@@ -855,7 +855,7 @@ def _draw_labels(seg_erosionpx: int | None, seg_boundaries: bool, alpha: float)
         legend_fontweight=legend_params.legend_fontweight,
         legend_loc=legend_params.legend_loc,
         legend_fontoutline=legend_params.legend_fontoutline,
-        na_in_legend=legend_params.na_in_legend,
+        na_in_legend=legend_params.na_in_legend if groups is None else len(groups) == len(set(color_vector)),
         colorbar=legend_params.colorbar,
         scalebar_dx=scalebar_params.scalebar_dx,
         scalebar_units=scalebar_params.scalebar_units,

src/spatialdata_plot/pl/utils.py

@@ -567,9 +567,6 @@ def _get_subplots(num_images: int, ncols: int = 4, width: int = 4, height: int =
     Union[plt.Figure, plt.Axes]
         Matplotlib figure and axes object.
     """
-    # if num_images <= 1:
-    # raise ValueError("Number of images must be greater than 1.")
-
     if num_images < ncols:
         nrows = 1
         ncols = num_images
@@ -733,8 +730,6 @@ def _set_color_source_vec(
         color = np.full(len(element), na_color)
         return color, color, False
 
-    # model = get_model(sdata[element_name])
-
     # Figure out where to get the color from
     origins = _locate_value(value_key=value_to_plot, sdata=sdata, element_name=element_name, table_name=table_name)
 
@@ -778,16 +773,13 @@ def _set_color_source_vec(
             palette=palette,
             na_color=na_color,
         )
+
         color_source_vector = color_source_vector.set_categories(color_mapping.keys())
         if color_mapping is None:
             raise ValueError("Unable to create color palette.")
 
         # do not rename categories, as colors need not be unique
         color_vector = color_source_vector.map(color_mapping)
-        if color_vector.isna().any():
-            if (na_cat_color := to_hex(na_color)) not in color_vector.categories:
-                color_vector = color_vector.add_categories([na_cat_color])
-            color_vector = color_vector.fillna(to_hex(na_color))
 
         return color_source_vector, color_vector, True
 
@@ -808,44 +800,43 @@ def _map_color_seg(
     seg_boundaries: bool = False,
 ) -> ArrayLike:
     cell_id = np.array(cell_id)
-    if color_vector is not None and isinstance(color_vector.dtype, pd.CategoricalDtype):
-        # users wants to plot a categorical column
+
+    if pd.api.types.is_categorical_dtype(color_vector.dtype):
+        # Case A: users wants to plot a categorical column
         if np.any(color_source_vector.isna()):
             cell_id[color_source_vector.isna()] = 0
-        val_im: ArrayLike = map_array(seg, cell_id, color_vector.codes + 1)
+        val_im: ArrayLike = map_array(seg.copy(), cell_id, color_vector.codes + 1)
         cols = colors.to_rgba_array(color_vector.categories)
-
     elif pd.api.types.is_numeric_dtype(color_vector.dtype):
-        # user wants to plot a continous column
+        # Case B: user wants to plot a continous column
         if isinstance(color_vector, pd.Series):
             color_vector = color_vector.to_numpy()
-        val_im = map_array(seg, cell_id, color_vector)
         cols = cmap_params.cmap(cmap_params.norm(color_vector))
-
+        val_im = map_array(seg.copy(), cell_id, cell_id)
     else:
-        val_im = map_array(seg.copy(), cell_id, cell_id)  # replace with same seg id to remove missing segs
-
-        if val_im.shape[0] == 1:
-            val_im = np.squeeze(val_im, axis=0)
-        if "#" in str(color_vector[0]):
-            # we have hex colors
-            assert all(_is_color_like(c) for c in color_vector), "Not all values are color-like."
-            cols = colors.to_rgba_array(color_vector)
+        # Case C: User didn't specify any colors
+        if color_source_vector is not None and (
+            set(color_vector) == set(color_source_vector)
+            and len(set(color_vector)) == 1
+            and set(color_vector) == {na_color}
+            and not na_color_modified_by_user
+        ):
+            val_im = map_array(seg.copy(), cell_id, cell_id)
+            RNG = default_rng(42)
+            cols = RNG.random((len(color_vector), 3))
         else:
-            cols = cmap_params.cmap(cmap_params.norm(color_vector))
+            # Case D: User didn't specify a column to color by, but modified the na_color
+            val_im = map_array(seg.copy(), cell_id, cell_id)
+            if "#" in str(color_vector[0]):
+                # we have hex colors
+                assert all(_is_color_like(c) for c in color_vector), "Not all values are color-like."
+                cols = colors.to_rgba_array(color_vector)
+            else:
+                cols = cmap_params.cmap(cmap_params.norm(color_vector))
 
     if seg_erosionpx is not None:
         val_im[val_im == erosion(val_im, square(seg_erosionpx))] = 0
 
-    if color_source_vector is not None and (
-        set(color_vector) == set(color_source_vector)
-        and len(set(color_vector)) == 1
-        and set(color_vector) == {na_color}
-        and not na_color_modified_by_user
-    ):
-        RNG = default_rng(42)
-        cols = RNG.random((len(cols), 3))
-
     seg_im: ArrayLike = label2rgb(
         label=val_im,
         colors=cols,
@@ -948,7 +939,7 @@ def _get_categorical_color_mapping(
     else:
         base_mapping = _generate_base_categorial_color_mapping(adata, cluster_key, color_source_vector, na_color)
 
-    return _modify_categorical_color_mapping(base_mapping, groups, palette)
+    return _modify_categorical_color_mapping(mapping=base_mapping, groups=groups, palette=palette)
 
 
 def _maybe_set_colors(
@@ -1587,19 +1578,14 @@ def _type_check_params(param_dict: dict[str, Any], element_type: str) -> dict[st
 
     palette = param_dict["palette"]
 
-    if (groups := param_dict.get("groups")) is not None and palette is None:
-        warnings.warn(
-            "Groups is specified but palette is not. Setting palette to default 'lightgray'", UserWarning, stacklevel=2
-        )
-        param_dict["palette"] = ["lightgray" for _ in range(len(groups))]
-
     if isinstance((palette := param_dict["palette"]), list):
         if not all(isinstance(p, str) for p in palette):
             raise ValueError("If specified, parameter 'palette' must contain only strings.")
     elif isinstance(palette, (str, type(None))) and "palette" in param_dict:
         param_dict["palette"] = [palette] if palette is not None else None
 
     if element_type in ["shapes", "points", "labels"] and (palette := param_dict.get("palette")) is not None:
+        groups = param_dict.get("groups")
         if groups is None:
             raise ValueError("When specifying 'palette', 'groups' must also be specified.")
         if len(groups) != len(palette):