Skip to content

Commit 2e4f0b3

Browse files
committed
perf(shapes): vectorize datashader polygon scaling
The datashader path scaled polygons with a per-geometry affinity.scale loop (_geometry[is_polygon].apply(lambda g: affinity.scale(...))) — a pure-Python loop that dominates large polygon renders (measured ~60% of a 100k-polygon render; the prototype's 33s at 1M is almost entirely this loop). Replace with _scale_geometries: scale every coordinate about each geometry's bounding-box centre (affinity.scale's default origin) in one vectorized pass via shapely.get_coordinates/set_coordinates. Byte-identical to affinity.scale including asymmetric shapes, multipolygons and holes (verified main-vs-branch, diff_px=0); ~12x polygons / ~6.5x multipolygons -> ~2x+ end-to-end at 100k, growing with n. Only fires for scale != 1.0. The affine transform loop is left as-is (vectorizing it is ~0.9x: shapely coordinate (de)serialization dominates, not the Python callback).
1 parent 55f4970 commit 2e4f0b3

2 files changed

Lines changed: 43 additions & 4 deletions

File tree

src/spatialdata_plot/pl/render.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -593,6 +593,20 @@ def _circles_render_as_points(shapes: gpd.GeoDataFrame, is_point: Any, render_pa
593593
return bool(np.isfinite(radius).all() and np.ptp(radius) == 0)
594594

595595

596+
def _scale_geometries(geometries: np.ndarray, scale: float) -> np.ndarray:
597+
"""Scale each geometry about its bounding-box centre (``shapely.affinity.scale``'s default origin).
598+
599+
Vectorised over all coordinates at once; a per-geometry ``affinity.scale`` loop is pure Python and
600+
dominates large polygon renders.
601+
"""
602+
import shapely
603+
604+
bbox = shapely.bounds(geometries) # (n, 4): minx, miny, maxx, maxy
605+
centre = np.column_stack([(bbox[:, 0] + bbox[:, 2]) / 2, (bbox[:, 1] + bbox[:, 3]) / 2])
606+
coords, idx = shapely.get_coordinates(geometries, return_index=True)
607+
return shapely.set_coordinates(geometries.copy(), (coords - centre[idx]) * scale + centre[idx])
608+
609+
596610
def _render_shapes(
597611
sdata: sd.SpatialData,
598612
render_params: ShapesRenderParams,
@@ -820,10 +834,8 @@ def _draw_centroids(xy: np.ndarray, radius: float | None = None) -> None:
820834
# Handle polygon/multipolygon scaling
821835
is_polygon = _geometry.type.isin(["Polygon", "MultiPolygon"])
822836
if is_polygon.any() and render_params.scale != 1.0:
823-
from shapely import affinity
824-
825-
shapes.loc[is_polygon, "geometry"] = _geometry[is_polygon].apply(
826-
lambda geom: affinity.scale(geom, xfact=render_params.scale, yfact=render_params.scale)
837+
shapes.loc[is_polygon, "geometry"] = _scale_geometries(
838+
_geometry[is_polygon].to_numpy(), render_params.scale
827839
)
828840

829841
# apply transformations to the individual points

tests/pl/test_render_shapes.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1978,3 +1978,30 @@ def bbox(**kw):
19781978
return xs.min(), xs.max(), ys.min(), ys.max()
19791979

19801980
assert bbox() == bbox(outline_width=1.0, outline_alpha=1.0, outline_color="black")
1981+
1982+
1983+
def test_scale_geometries_matches_affinity_scale():
1984+
# The vectorised datashader polygon scale must equal shapely.affinity.scale's default
1985+
# (bounding-box-centre) origin, including for asymmetric shapes, multipolygons and holes.
1986+
import shapely
1987+
from shapely import affinity
1988+
from shapely.geometry import MultiPolygon
1989+
1990+
from spatialdata_plot.pl.render import _scale_geometries
1991+
1992+
rng = np.random.default_rng(0)
1993+
geoms = []
1994+
for cx, cy in rng.random((50, 2)) * 100:
1995+
# asymmetric exterior (bbox-centre != centroid) with a hole
1996+
poly = Polygon(
1997+
[(cx, cy), (cx + 6, cy + 1), (cx + 5, cy + 4), (cx + 1, cy + 3)],
1998+
[[(cx + 2, cy + 2), (cx + 3, cy + 2), (cx + 3, cy + 3), (cx + 2, cy + 3)][::-1]],
1999+
)
2000+
geoms.append(poly)
2001+
geoms.append(MultiPolygon([geoms[0], affinity.translate(geoms[1], 10, 10)])) # multi-part
2002+
arr = np.array(geoms, dtype=object)
2003+
2004+
for scale in (0.6, 2.0):
2005+
expected = np.array([affinity.scale(g, xfact=scale, yfact=scale) for g in geoms], dtype=object)
2006+
result = _scale_geometries(arr, scale)
2007+
assert all(shapely.equals_exact(a, b, tolerance=1e-9) for a, b in zip(expected, result, strict=True))

0 commit comments

Comments
 (0)