render_shapes now plots polygons (#41)

Author: timtreis

src/spatialdata_plot/pl/basic.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 from typing import Callable, Optional, Union
 
+import geopandas as gpd
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
@@ -637,15 +638,33 @@ def show(
             # get biggest image after transformations to set ax size
             x_dims = []
             y_dims = []
+
             for cmd, _ in render_cmds.items():
                 if cmd == "render_images":
                     y_dims += [(0, x.shape[1]) for x in sdata.images.values()]
                     x_dims += [(0, x.shape[2]) for x in sdata.images.values()]
 
                 elif cmd == "render_shapes":
-                    for k in sdata.shapes.keys():
-                        x_dims += [(min(sdata.shapes[k].geometry.x), max(sdata.shapes[k].geometry.x))]
-                        y_dims += [(min(sdata.shapes[k].geometry.y), max(sdata.shapes[k].geometry.y))]
+                    for key in sdata.shapes.keys():
+                        points = []
+                        polygons = []
+
+                        for _, row in sdata.shapes[key].iterrows():
+                            if row["geometry"].type == "Point":
+                                points.append(row)
+                            else:
+                                polygons.append(row)
+
+                        if len(points) > 0:
+                            points_df = gpd.GeoDataFrame(data=points)
+                            x_dims += [(min(points_df.geometry.x), max(points_df.geometry.x))]
+                            y_dims += [(min(points_df.geometry.y), max(points_df.geometry.y))]
+
+                        if len(polygons) > 0:
+                            for p in polygons:
+                                minx, miny, maxx, maxy = p.geometry.bounds
+                                x_dims += [(minx, maxx)]
+                                y_dims += [(miny, maxy)]
 
                 elif cmd == "render_labels":
                     y_dims += [(0, x.shape[0]) for x in sdata.labels.values()]
@@ -767,7 +786,6 @@ def show(
                             cell_ids_per_label = {}
                             for key in list(sdata.labels.keys()):
                                 cell_ids_per_label[key] = sdata.labels[key].values.max()
-                            print(cell_ids_per_label)
                             region_key = "tmp_label_id"
                             instance_key = "tmp_cell_id"
                             params["instance_key"] = instance_key
@@ -786,7 +804,6 @@ def show(
                             distinct_cells = max(list(cell_ids_per_label.values()))
 
                         if sdata.table is not None:
-                            # print("Plotting a lot of cells with random colors, might take a while...")
                             sdata.table.uns[f"{instance_key}_colors"] = _get_random_hex_colors(distinct_cells)
 
                         elif sdata.table is None:

src/spatialdata_plot/pl/render.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterable
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import matplotlib
 import matplotlib.patches as mpatches
@@ -53,31 +53,53 @@ def _render_channels(
 
 def _render_shapes(
     sdata: sd.SpatialData,
-    params: dict[str, Union[str, int, float, Iterable[str]]],
+    params: dict[str, Optional[Union[str, int, float, Iterable[str]]]],
     key: str,
     ax: matplotlib.axes.SubplotBase,
     extent: dict[str, list[int]],
 ) -> None:
+    colors: Optional[Union[str, int, float, Iterable[str]]] = None  # to shut up mypy
     if sdata.table is not None and isinstance(params["instance_key"], str) and isinstance(params["color_key"], str):
         colors = [to_rgb(c) for c in sdata.table.uns[f"{params['color_key']}_colors"]]
     elif isinstance(params["palette"], str):
         colors = [params["palette"]]
     elif isinstance(params["palette"], Iterable):
         colors = [to_rgb(c) for c in list(params["palette"])]
     else:
-        colors = [params["palette"]]
+        colors = params["palette"]
 
     ax.set_xlim(extent["x"][0], extent["x"][1])
     ax.set_ylim(extent["y"][0], extent["y"][1])
 
-    shape = sdata.shapes[key]
+    points = []
+    polygons = []
 
-    ax.scatter(
-        x=shape.geometry.x,
-        y=shape.geometry.y,
-        s=shape.radius,
-        color=colors,
-    )
+    for _, row in sdata.shapes[key].iterrows():
+        if row["geometry"].geom_type == "Point":
+            points.append((row[0], row[1]))  # (point, radius)
+        elif row["geometry"].geom_type == "Polygon":
+            polygons.append(row[0])  # just polygon
+        else:
+            raise NotImplementedError(f"Geometry type {row['geometry'].type} not supported.")
+
+    if len(polygons) > 0:
+        for polygon in polygons:
+            ax.add_patch(
+                mpatches.Polygon(
+                    polygon.exterior.coords,
+                    color=colors,
+                )
+            )
+
+    if len(points) > 0:
+        for point, radius in points:
+            ax.add_patch(
+                mpatches.Circle(
+                    (point.x, point.y),
+                    radius=radius,
+                    color=colors,
+                )
+            )
 
     ax.set_title(key)
 
@@ -146,11 +168,6 @@ def _render_images(
         elif n_channels == 2:
             colors = ListedColormap(["#d30cb8", "#6df1d8"])
         elif n_channels == 3:
-            # bg = [(1, 1, 1, 1)]
-            # cmap_red = ListedColormap([(1, 0, 0, i) for i in reversed(range(0, 256, 1))] + bg)
-            # cmap_green = ListedColormap([(0, 1, 0, i) for i in reversed(range(0, 256, 1))] + bg)
-            # cmap_blue = ListedColormap([(0, 0, 1, i) for i in reversed(range(0, 256, 1))] + bg)
-            # colors = [cmap_red, cmap_green, cmap_blue]
             colors = ListedColormap(["red", "blue", "green"])
         else:
             # we do PCA to reduce to 3 channels
@@ -171,10 +188,6 @@ def _render_images(
     )
 
     ax.set_title(key)
-    # ax.set_xlabel("spatial1")
-    # ax.set_ylabel("spatial2")
-    # ax.set_xticks([])
-    # ax.set_yticks([])
 
 
 def _render_labels(
@@ -249,7 +262,3 @@ def _render_labels(
         ax.legend(handles=patches, bbox_to_anchor=(0.9, 0.9), loc="upper left", frameon=False)
 
     ax.set_title(key)
-    # ax.set_xlabel("spatial1")
-    # ax.set_ylabel("spatial2")
-    # ax.set_xticks([])
-    # ax.set_yticks([])

test_sandbox_data.py

@@ -6,11 +6,12 @@
 # - You can use `jupyter nbconvert --to python test_sandbox_data.ipynb` to convert it back
 # - Otherwise run it from the CLI and verify that the plots are okay
 
-# In[4]:
+# In[1]:
 
 
 import matplotlib.pyplot as plt
 import spatialdata as sd
+from spatialdata.datasets import blobs
 
 import spatialdata_plot
 
@@ -19,26 +20,31 @@
 DATA_DIR = "/Users/tim.treis/Documents/GitHub/spatialdata-sandbox/"
 
 
-# ## Load data
-# Adjust paths as neccecary
+# In[2]:
+
 
-# In[15]:
+(blobs().pl.render_images().pl.render_labels().pl.render_shapes().pl.render_points(color_key="genes").pl.show())
+
+
+# In[4]:
 
 
 # Mibi
 
-sdata = sd.read_zarr(DATA_DIR + "mibitof/data.zarr")
-sdata.pl.render_images().pl.render_labels().pl.show()
+mibitof = sd.read_zarr(DATA_DIR + "mibitof/data.zarr")
+
+(mibitof.pl.render_images().pl.render_labels().pl.show())
 
 plt.savefig("mibi.png")
 
 
-# In[16]:
+# In[5]:
 
 
 # Visium
 
-sdata = sd.read_zarr(DATA_DIR + "visium/data.zarr")
-sdata.pl.render_images().pl.render_shapes().pl.show(width=12, height=12)
+visium = sd.read_zarr(DATA_DIR + "visium/data.zarr")
+
+(visium.pl.render_images().pl.render_shapes().pl.show())
 
 plt.savefig("visium.png")