perf(shapes): vectorize datashader polygon scaling (#738)

Author: timtreis

src/spatialdata_plot/pl/render.py

@@ -593,6 +593,20 @@ def _circles_render_as_points(shapes: gpd.GeoDataFrame, is_point: Any, render_pa
     return bool(np.isfinite(radius).all() and np.ptp(radius) == 0)
 
 
+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 is pure Python and
+    dominates large polygon renders.
+    """
+    import shapely
+
+    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 _render_shapes(
     sdata: sd.SpatialData,
     render_params: ShapesRenderParams,
@@ -817,10 +831,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 shapely.geometry import MultiPolygon
+
+    from spatialdata_plot.pl.render 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
+        poly = 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(poly)
+    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))