Merge pull request #367 from MannLabs/bump_spatialdata_version

[VERSION] bumped spatialdata-plot version to 0.2.14 to fix incorrect channel scaling in default plotting

Author: sophiamaedler

requirements/requirements.txt

@@ -8,7 +8,7 @@ spatialdata>=0.3.0,<0.6
 pyarrow<22.0.0
 py-lmd>=1.3.1
 
-spatialdata-plot<=0.2.11
+spatialdata-plot>=0.2.14
 matplotlib
 
 tifffile

src/scportrait/plotting/sdata.py

@@ -55,6 +55,8 @@ def _get_shape_element(sdata, element_name) -> tuple[int, int]:
         _, x, y = shape
     elif len(shape) == 2:
         x, y = shape
+    else:
+        raise ValueError(f"Unsupported shape for element '{element_name}': expected 2D or 3D array, got {shape}.")
     return x, y
 
 
@@ -115,6 +117,7 @@ def plot_image(
     if ax is not None:
         if dpi is not None:
             warnings.warn("DPI is ignored when an axis is provided.", stacklevel=2)
+        fig = ax.figure
     else:
         # get size of spatialdata object to plot (required for calculating figure size if DPI is set)
         x, y = _get_shape_element(sdata, image_name)
@@ -189,6 +192,7 @@ def plot_segmentation_mask(
     if ax is not None:
         if dpi is not None:
             warnings.warn("DPI is ignored when an axis is provided.", stacklevel=2)
+        fig = ax.figure
     else:
         # get size of spatialdata object to plot (required for calculating figure size if DPI is set)
         x, y = _get_shape_element(sdata, masks[0])
@@ -217,20 +221,27 @@ def plot_segmentation_mask(
         if selected_channels is not None:
             if not isinstance(selected_channels, Iterable):
                 selected_channels = [selected_channels]
+            if any(i < 0 or i >= len(channel_names) for i in selected_channels):
+                raise ValueError(
+                    f"selected_channels contains out-of-range indices for background image '{background_image}'."
+                )
+            if len(selected_channels) > len(PALETTE):
+                raise ValueError("selected_channels has more entries than the available palette length.")
             channel_names = [channel_names[i] for i in selected_channels]
             c = len(channel_names)
-            palette = [PALETTE[x] for x in selected_channels]
+            palette = PALETTE[:c]
         else:
             if c > max_channels_to_plot:
-                c = 4
+                c = min(c, max_channels_to_plot)
             palette = PALETTE[:c]
             channel_names = list(channel_names[:c])
 
         sdata.pl.render_images(background_image, channel=channel_names, palette=palette).pl.show(ax=ax, colorbar=False)
 
     # plot selected segmentation masks
     for mask in masks:
-        assert mask in sdata, f"Mask {mask} not found in sdata object."
+        if mask not in sdata:
+            raise KeyError(f"Mask {mask} not found in sdata object.")
         if f"{mask}_vectorized" not in sdata:
             sdata[f"{mask}_vectorized"] = spatialdata.to_polygons(sdata[mask])
         sdata.pl.render_shapes(
@@ -298,6 +309,7 @@ def plot_shapes(
     if ax is not None:
         if dpi is not None:
             warnings.warn("DPI is ignored when an axis is provided.", stacklevel=2)
+        fig = ax.figure
     else:
         # get size of spatialdata object to plot (required for calculating figure size if DPI is set)
         x, y = _get_shape_element(sdata, shapes_layer)
@@ -306,7 +318,8 @@ def plot_shapes(
         fig, ax = _create_figure_dpi(x=x, y=y, dpi=dpi)
 
     # plot selected shapes layer
-    assert shapes_layer in sdata, f"Shapes layer {shapes_layer} not found in sdata object."
+    if shapes_layer not in sdata:
+        raise KeyError(f"Shapes layer {shapes_layer} not found in sdata object.")
 
     sdata.pl.render_shapes(
         f"{shapes_layer}",
@@ -374,6 +387,7 @@ def plot_labels(
     if ax is not None:
         if dpi is not None:
             warnings.warn("DPI is ignored when an axis is provided.", stacklevel=2)
+        fig = ax.figure
     else:
         # get size of spatialdata object to plot (required for calculating figure size if DPI is set)
         x, y = _get_shape_element(sdata, label_layer)
@@ -410,18 +424,25 @@ def plot_labels(
                     annotating_table = spatialdata.models.TableModel.parse(annotating_table)
                     break
         if found_annotation is not None:
+            had_annotation = "_annotation" in sdata
+            prev_annotation = sdata["_annotation"] if had_annotation else None
             sdata["_annotation"] = annotating_table
-            sdata.pl.render_shapes(
-                f"{label_layer}_vectorized",
-                color=color,
-                fill_alpha=fill_alpha,
-                outline_alpha=0,
-                cmap=cmap,
-                palette=palette,
-                groups=groups,
-                norm=norm,
-            ).pl.show(ax=ax)
-            del sdata["_annotation"]  # delete element again after plotting
+            try:
+                sdata.pl.render_shapes(
+                    f"{label_layer}_vectorized",
+                    color=color,
+                    fill_alpha=fill_alpha,
+                    outline_alpha=0,
+                    cmap=cmap,
+                    palette=palette,
+                    groups=groups,
+                    norm=norm,
+                ).pl.show(ax=ax)
+            finally:
+                if had_annotation:
+                    sdata["_annotation"] = prev_annotation
+                else:
+                    del sdata["_annotation"]  # delete element again after plotting
         else:
             try:
                 sdata.pl.render_labels(

tests/unit_tests/plotting/test_sdata.py

@@ -35,6 +35,21 @@ def test_plot_image(sdata_with_labels, channel_names, palette, return_fig, show_
         assert fig is None
 
 
+def test_plot_image_with_ax_returns_fig(sdata_with_labels):
+    fig, ax = plt.subplots()
+    returned = plotting.plot_image(
+        sdata=sdata_with_labels,
+        image_name="blobs_image",
+        channel_names=[0],
+        palette=["red"],
+        return_fig=True,
+        show_fig=False,
+        ax=ax,
+    )
+    assert returned is fig
+    plt.close(fig)
+
+
 @pytest.mark.parametrize(
     "selected_channels, background_image",
     [
@@ -56,6 +71,33 @@ def test_plot_segmentation_mask(sdata_with_labels, selected_channels, background
     plt.close(fig)
 
 
+def test_plot_segmentation_mask_selected_channels_out_of_range(sdata_with_labels):
+    with pytest.raises(ValueError):
+        plotting.plot_segmentation_mask(
+            sdata=sdata_with_labels,
+            masks=["blobs_labels"],
+            background_image="blobs_image",
+            selected_channels=[999],
+            return_fig=False,
+            show_fig=False,
+        )
+
+
+def test_plot_segmentation_mask_with_ax_returns_fig(sdata_with_labels):
+    fig, ax = plt.subplots()
+    returned = plotting.plot_segmentation_mask(
+        sdata=sdata_with_labels,
+        masks=["blobs_labels"],
+        background_image="blobs_image",
+        selected_channels=[0],
+        return_fig=True,
+        show_fig=False,
+        ax=ax,
+    )
+    assert returned is fig
+    plt.close(fig)
+
+
 @pytest.mark.parametrize(
     "vectorized, color",
     [
@@ -76,3 +118,31 @@ def test_plot_labels(sdata_with_labels, vectorized, color):
     )
     assert isinstance(fig, plt.Figure)
     plt.close(fig)
+
+
+def test_plot_labels_with_ax_returns_fig(sdata_with_labels):
+    fig, ax = plt.subplots()
+    returned = plotting.plot_labels(
+        sdata=sdata_with_labels,
+        label_layer="blobs_labels",
+        vectorized=False,
+        color="labelling_categorical",
+        return_fig=True,
+        show_fig=False,
+        ax=ax,
+    )
+    assert returned is fig
+    plt.close(fig)
+
+
+def test_plot_shapes_with_ax_returns_fig(sdata_with_labels):
+    fig, ax = plt.subplots()
+    returned = plotting.plot_shapes(
+        sdata=sdata_with_labels,
+        shapes_layer="blobs_polygons",
+        return_fig=True,
+        show_fig=False,
+        ax=ax,
+    )
+    assert returned is fig
+    plt.close(fig)