handle render shape colors; add parameter for render method

Author: ckmah

src/spatialdata_plot/pl/basic.py

@@ -160,6 +160,7 @@ def render_shapes(
         cmap: Colormap | str | None = None,
         norm: bool | Normalize = False,
         scale: float | int = 1.0,
+        method: str | None = None,
         **kwargs: Any,
     ) -> sd.SpatialData:
         """
@@ -204,6 +205,9 @@ def render_shapes(
             Colormap normalization for continuous annotations.
         scale : float | int, default 1.0
             Value to scale circles, if present.
+        method : str | None, optional
+            Whether to use 'matplotlib' and 'datashader'. When None, the method is
+            chosen based on the size of the data.
         **kwargs : Any
             Additional arguments to be passed to cmap and norm.
 
@@ -317,6 +321,12 @@ def render_shapes(
         if scale < 0:
             raise ValueError("Parameter 'scale' must be a positive number.")
 
+        if method is not None:
+            if not isinstance(method, str):
+                raise TypeError("Parameter 'method' must be a string.")
+            if method not in ["matplotlib", "datashader"]:
+                raise ValueError("Parameter 'method' must be either 'matplotlib' or 'datashader'.")
+
         sdata = self._copy()
         sdata = _verify_plotting_tree(sdata)
         n_steps = len(sdata.plotting_tree.keys())
@@ -343,6 +353,7 @@ def render_shapes(
             fill_alpha=fill_alpha,
             transfunc=kwargs.get("transfunc", None),
             zorder=n_steps,
+            method=method,
         )
 
         return sdata
@@ -358,6 +369,7 @@ def render_points(
         cmap: Colormap | str | None = None,
         norm: None | Normalize = None,
         size: float | int = 1.0,
+        method: str | None = None,
         **kwargs: Any,
     ) -> sd.SpatialData:
         """
@@ -392,6 +404,9 @@ def render_points(
             Colormap normalization for continuous annotations.
         size : float | int, default 1.0
             Size of the points
+        method : str | None, optional
+            Whether to use 'matplotlib' and 'datashader'. When None, the method is
+            chosen based on the size of the data.
         kwargs
             Additional arguments to be passed to cmap and norm.
 
@@ -479,6 +494,12 @@ def render_points(
         if size < 0:
             raise ValueError("Parameter 'size' must be a positive number.")
 
+        if method is not None:
+            if not isinstance(method, str):
+                raise TypeError("Parameter 'method' must be a string.")
+            if method not in ["matplotlib", "datashader"]:
+                raise ValueError("Parameter 'method' must be either 'matplotlib' or 'datashader'.")
+
         sdata = self._copy()
         sdata = _verify_plotting_tree(sdata)
         n_steps = len(sdata.plotting_tree.keys())
@@ -501,6 +522,7 @@ def render_points(
             transfunc=kwargs.get("transfunc", None),
             size=size,
             zorder=n_steps,
+            method=method,
         )
 
         return sdata

src/spatialdata_plot/pl/render.py

@@ -134,12 +134,21 @@ def _render_shapes(
         # Apply the transformation to the PatchCollection's paths
         trans = get_transformation(sdata_filt.shapes[e], get_all=True)[coordinate_system]
         affine_trans = trans.to_affine_matrix(input_axes=("x", "y"), output_axes=("x", "y"))
-        trans = mtransforms.Affine2D(matrix=affine_trans) + ax.transData
+        trans = mtransforms.Affine2D(matrix=affine_trans)
 
         shapes = gpd.GeoDataFrame(shapes, geometry="geometry")
 
-        if len(shapes) < 1:
-            logger.info("Using matplotlib")
+        # Determine which method to use for rendering. Default is matplotlib for under 100 shapes and datashader for more
+        # User can also specify the method to use
+        method = render_params.method
+
+        if method is None:
+            method = "datashader" if len(shapes) > 100 else "matplotlib"
+        elif method not in ["matplotlib", "datashader"]:
+            raise ValueError("Method must be either 'matplotlib' or 'datashader'.")
+
+        if method == "matplotlib":
+            logger.info(f"Using {method}")
             _cax = _get_collection_shape(
                 shapes=shapes,
                 s=render_params.scale,
@@ -159,8 +168,12 @@ def _render_shapes(
             for path in _cax.get_paths():
                 path.vertices = trans.transform(path.vertices)
                 cax = ax.add_collection(_cax)
-        else:
-            logger.info("Using datashader")
+        elif method == "datashader":
+            logger.info(f"Using {method}")
+
+            # Where to put this
+            trans = mtransforms.Affine2D(matrix=affine_trans) + ax.transData
+
             extent = get_extent(sdata.shapes[e])
             x_ext = extent["x"][1]
             y_ext = extent["y"][1]
@@ -179,10 +192,28 @@ def _render_shapes(
 
             # Handle circles encoded as points with radius
             if is_point.any():  # TODO
-                shapes.loc[is_point, "geometry"] = _geometry[is_point].buffer(shapes[is_point]["radius"])
+                scale = shapes[is_point]["radius"] * render_params.scale
+                shapes.loc[is_point, "geometry"] = _geometry[is_point].buffer(scale)
 
             agg = cvs.polygons(shapes, geometry="geometry", agg=ds.count())
-            ds_result = ds.tf.shade(agg)
+
+            # Render shapes with datashader
+            if render_params.col_for_color is not None and (
+                render_params.groups is None or len(render_params.groups) > 1
+            ):
+                agg = cvs.polygons(shapes, geometry="geometry", agg=ds.by(render_params.col_for_color, ds.count()))
+            else:
+                agg = cvs.polygons(shapes, geometry="geometry", agg=ds.count())
+
+            color_key = (
+                [x[:-2] for x in color_vector.categories.values]
+                if (type(color_vector) == pd.core.arrays.categorical.Categorical)
+                and (len(color_vector.categories.values) > 1)
+                else None
+            )
+            ds_result = ds.tf.shade(
+                agg, cmap=color_vector[0][:-2], alpha=render_params.fill_alpha * 255, color_key=color_key, min_alpha=200
+            )
 
             # Render image
             rgba_image = np.transpose(ds_result.to_numpy().base, (0, 1, 2))
@@ -312,9 +343,13 @@ def _render_points(
 
         norm = copy(render_params.cmap_params.norm)
 
-        # optionally render points using datashader
-        # TODO: maybe move this, add heuristic
-        if len(points) > 50:
+        method = render_params.method
+        if method is None:
+            method = "datashader" if len(points.shape[0]) > 10000 else "matplotlib"
+        elif method not in ["matplotlib", "datashader"]:
+            raise ValueError("Method must be either 'matplotlib' or 'datashader'.")
+
+        if method == "datashader":
             extent = get_extent(sdata_filt.points[e], coordinate_system=coordinate_system)
             x_ext = extent["x"][1]
             y_ext = extent["y"][1]
@@ -368,7 +403,7 @@ def _render_points(
             rbga_image = np.transpose(ds_result.to_numpy().base, (0, 1, 2))
             ax.imshow(rbga_image, zorder=render_params.zorder)
             cax = None
-        else:
+        elif method == "matplotlib":
             # original way of plotting points
             _cax = ax.scatter(
                 adata[:, 0].X.flatten(),

src/spatialdata_plot/pl/render_params.py

@@ -81,6 +81,7 @@ class ShapesRenderParams:
     fill_alpha: float = 0.3
     scale: float = 1.0
     transfunc: Callable[[float], float] | None = None
+    method: str | None = None
     zorder: int | None = None
 
 
@@ -97,6 +98,7 @@ class PointsRenderParams:
     alpha: float = 1.0
     size: float = 1.0
     transfunc: Callable[[float], float] | None = None
+    method: str | None = None
     zorder: int | None = None