perf(shapes): vectorize datashader polygon scaling (re-land of #738) (#740)

Author: timtreis

src/spatialdata_plot/pl/_geometry.py

@@ -40,6 +40,17 @@ def _scale_path_around_centroid(path: mpath.Path, scale_factor: float) -> None:
     path.vertices = centroid + (path.vertices - centroid) * scale_value
 
 
+def _scale_geometries(geometries: np.ndarray, scale: float) -> np.ndarray:
+    """Scale each geometry about its bounding-box centre (``shapely.affinity.scale``'s default origin).
+
+    Vectorised over all coordinates at once — a per-geometry ``affinity.scale`` loop dominates large renders.
+    """
+    bbox = shapely.bounds(geometries)  # (n, 4): minx, miny, maxx, maxy
+    centre = np.column_stack([(bbox[:, 0] + bbox[:, 2]) / 2, (bbox[:, 1] + bbox[:, 3]) / 2])
+    coords, idx = shapely.get_coordinates(geometries, return_index=True)
+    return shapely.set_coordinates(geometries.copy(), (coords - centre[idx]) * scale + centre[idx])
+
+
 def _normalize_geom(geom: Any) -> Any:
     """Canonicalize ring orientation so matplotlib's fill rules render holes correctly.
 

src/spatialdata_plot/pl/render.py

@@ -57,6 +57,7 @@
     _build_shape_paths,
     _convert_shapes,
     _get_collection_shape,
+    _scale_geometries,
     _validate_polygons,
 )
 from spatialdata_plot.pl._validate import (
@@ -817,10 +818,8 @@ def _draw_centroids(xy: np.ndarray, radius: float | None = None) -> None:
         # Handle polygon/multipolygon scaling
         is_polygon = _geometry.type.isin(["Polygon", "MultiPolygon"])
         if is_polygon.any() and render_params.scale != 1.0:
-            from shapely import affinity
-
-            shapes.loc[is_polygon, "geometry"] = _geometry[is_polygon].apply(
-                lambda geom: affinity.scale(geom, xfact=render_params.scale, yfact=render_params.scale)
+            shapes.loc[is_polygon, "geometry"] = _scale_geometries(
+                _geometry[is_polygon].to_numpy(), render_params.scale
             )
 
         # apply transformations to the individual points

tests/pl/test_render_shapes.py

@@ -1978,3 +1978,30 @@ def bbox(**kw):
         return xs.min(), xs.max(), ys.min(), ys.max()
 
     assert bbox() == bbox(outline_width=1.0, outline_alpha=1.0, outline_color="black")
+
+
+def test_scale_geometries_matches_affinity_scale():
+    # The vectorised datashader polygon scale must equal shapely.affinity.scale's default
+    # (bounding-box-centre) origin, including for asymmetric shapes, multipolygons and holes.
+    import shapely
+    from shapely import affinity
+
+    from spatialdata_plot.pl._geometry import _scale_geometries
+
+    rng = np.random.default_rng(0)
+    geoms = []
+    for cx, cy in rng.random((50, 2)) * 100:
+        # asymmetric exterior (bbox-centre != centroid) with a hole
+        geoms.append(
+            Polygon(
+                [(cx, cy), (cx + 6, cy + 1), (cx + 5, cy + 4), (cx + 1, cy + 3)],
+                [[(cx + 2, cy + 2), (cx + 3, cy + 2), (cx + 3, cy + 3), (cx + 2, cy + 3)][::-1]],
+            )
+        )
+    geoms.append(MultiPolygon([geoms[0], affinity.translate(geoms[1], 10, 10)]))  # multi-part
+    arr = np.array(geoms, dtype=object)
+
+    for scale in (0.6, 2.0):
+        expected = np.array([affinity.scale(g, xfact=scale, yfact=scale) for g in geoms], dtype=object)
+        result = _scale_geometries(arr, scale)
+        assert all(shapely.equals_exact(a, b, tolerance=1e-9) for a, b in zip(expected, result, strict=True))