ENH: add voronoi_polygons method (#3177)

Co-authored-by: Matt Richards <45483497+m-richards@users.noreply.github.com>
geopandas · May 6, 2024 · 5186e69 · 5186e69
1 parent 45b653e
commit 5186e69
Show file tree

Hide file tree

Showing 4 changed files with 177 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,6 +35,7 @@ New methods:
 - Added `is_ccw` method from shapely to GeoSeries/GeoDataframe (#3027).
 - Added `is_closed` attribute from shapely to GeoSeries/GeoDataframe (#3092).
 - Added `force_2d` and `force_3d` methods from shapely to GeoSeries/GeoDataframe (#3090).
+- Added `voronoi_polygons` method from shapely to GeoSeries/GeoDataframe (#3177).
 - Added `contains_properly` method from shapely to GeoSeries/GeoDataframe (#3105).
 - Added `build_area` method exposing `build_area` shapely to GeoSeries/GeoDataframe (#3202).
 - Added `snap` method from shapely to GeoSeries/GeoDataframe (#3086).

diff --git a/doc/source/docs/reference/geoseries.rst b/doc/source/docs/reference/geoseries.rst
@@ -121,6 +121,7 @@ Constructive methods and attributes
    GeoSeries.simplify
    GeoSeries.snap
    GeoSeries.transform
+   GeoSeries.voronoi_polygons
 
 Affine transformations
 ----------------------

diff --git a/geopandas/base.py b/geopandas/base.py
@@ -943,11 +943,141 @@ def delaunay_triangles(self, tolerance=0.0, only_edges=False):
         1    MULTILINESTRING ((5 3, 10 10), (5 3, 6 3), (6 ...
         2    MULTILINESTRING ((5 3, 10 10), (5 3, 6 3), (6 ...
         dtype: geometry
+
+        See also
+        --------
+        GeoSeries.voronoi_polygons : Voronoi diagram around vertices
         """
         return _delegate_geo_method(
             "delaunay_triangles", self, tolerance=tolerance, only_edges=only_edges
         )
 
+    def voronoi_polygons(self, tolerance=0.0, extend_to=None, only_edges=False):
+        """Returns a ``GeoSeries`` consisting of objects representing
+        the computed Voronoi diagram around the vertices of an input geometry.
+
+        All geometries within the GeoSeries are considered together within a single
+        Voronoi diagram. The resulting geometries therefore do not necessarily map 1:1
+        to input geometries. Note that each vertex of a geometry is considered a site
+        for the Voronoi diagram, so the diagram will be constructed around the vertices
+        of each geometry.
+
+        Notes
+        -----
+        The order of polygons in the output currently does not correspond to the order
+        of input vertices.
+
+        If you want to generate a Voronoi diagram for each geometry separately, use
+        :func:`shapely.voronoi_polygons` instead.
+
+        Parameters
+        ----------
+        tolerance : float, default 0.0
+            Snap input vertices together if their distance is less than this value.
+        extend_to : shapely.Geometry, default None
+            If set, the Voronoi diagram will be extended to cover the
+            envelope of this geometry (unless this envelope is smaller than the input
+            geometry).
+        only_edges : bool (optional, default False)
+            If set to True, the diagram will return LineStrings instead
+            of Polygons.
+
+        Examples
+        --------
+        The most common use case is to generate polygons representing the Voronoi
+        diagram around a set of points:
+
+        >>> from shapely import LineString, MultiPoint, Point, Polygon
+        >>> s = geopandas.GeoSeries(
+        ...     [
+        ...         Point(1, 1),
+        ...         Point(2, 2),
+        ...         Point(1, 3),
+        ...         Point(0, 2),
+        ...     ]
+        ... )
+        >>> s
+        0    POINT (1 1)
+        1    POINT (2 2)
+        2    POINT (1 3)
+        3    POINT (0 2)
+        dtype: geometry
+
+        By default, you get back a GeoSeries of polygons:
+
+        >>> s.voronoi_polygons()
+        0     POLYGON ((-2 5, 1 2, -2 -1, -2 5))
+        1        POLYGON ((4 5, 1 2, -2 5, 4 5))
+        2    POLYGON ((-2 -1, 1 2, 4 -1, -2 -1))
+        3       POLYGON ((4 -1, 1 2, 4 5, 4 -1))
+        dtype: geometry
+
+        If you set only_edges to True, you get back LineStrings representing the
+        edges of the Voronoi diagram:
+
+        >>> s.voronoi_polygons(only_edges=True)
+        0     LINESTRING (-2 5, 1 2)
+        1    LINESTRING (1 2, -2 -1)
+        2      LINESTRING (4 5, 1 2)
+        3     LINESTRING (1 2, 4 -1)
+        dtype: geometry
+
+        You can also extend each diagram to a given geometry:
+
+        >>> limit = Polygon([(-10, -10), (0, 15), (15, 15), (15, 0)])
+        >>> s.voronoi_polygons(extend_to=limit)
+        0              POLYGON ((-10 13, 1 2, -10 -9, -10 13))
+        1    POLYGON ((15 15, 15 -10, 13 -10, 1 2, 14 15, 1...
+        2    POLYGON ((-10 -10, -10 -9, 1 2, 13 -10, -10 -10))
+        3       POLYGON ((-10 15, 14 15, 1 2, -10 13, -10 15))
+        dtype: geometry
+
+        The method supports any geometry type but keep in mind that the underlying
+        algorithm is based on the vertices of the input geometries only and does not
+        consider edge segments between vertices.
+
+        >>> s2 = geopandas.GeoSeries(
+        ...     [
+        ...         Polygon([(0, 0), (1, 1), (0, 1)]),
+        ...         LineString([(1, 0), (2, 1), (1, 2)]),
+        ...         MultiPoint([(2, 3), (2, 0), (3, 1)]),
+        ...     ]
+        ... )
+        >>> s2
+        0      POLYGON ((0 0, 1 1, 0 1, 0 0))
+        1          LINESTRING (1 0, 2 1, 1 2)
+        2    MULTIPOINT ((2 3), (2 0), (3 1))
+        dtype: geometry
+
+        >>> s2.voronoi_polygons()
+        0    POLYGON ((1.5 1.5, 1.5 0.5, 0.5 0.5, 0.5 1.5, ...
+        1    POLYGON ((1.5 0.5, 1.5 1.5, 2 2, 2.5 2, 2.5 0....
+        2    POLYGON ((-3 -3, -3 0.5, 0.5 0.5, 0.5 -3, -3 -3))
+        3    POLYGON ((0.5 -3, 0.5 0.5, 1.5 0.5, 1.5 -3, 0....
+        4     POLYGON ((-3 5, 0.5 1.5, 0.5 0.5, -3 0.5, -3 5))
+        5    POLYGON ((-3 6, -2 6, 2 2, 1.5 1.5, 0.5 1.5, -...
+        6    POLYGON ((1.5 -3, 1.5 0.5, 2.5 0.5, 6 -3, 1.5 ...
+        7       POLYGON ((6 6, 6 3.75, 2.5 2, 2 2, -2 6, 6 6))
+        8       POLYGON ((6 -3, 2.5 0.5, 2.5 2, 6 3.75, 6 -3))
+        dtype: geometry
+
+        See also
+        --------
+        GeoSeries.delaunay_triangles : Delaunay triangulation around vertices
+        """
+        from .geoseries import GeoSeries
+
+        geometry_input = shapely.geometrycollections(self.geometry.values._data)
+
+        voronoi = shapely.voronoi_polygons(
+            geometry_input,
+            tolerance=tolerance,
+            extend_to=extend_to,
+            only_edges=only_edges,
+        )
+
+        return GeoSeries(voronoi, crs=self.crs).explode(ignore_index=True)
+
     @property
     def envelope(self):
         """Returns a ``GeoSeries`` of geometries representing the envelope of

diff --git a/geopandas/tests/test_geom_methods.py b/geopandas/tests/test_geom_methods.py
@@ -1110,22 +1110,22 @@ def test_delaunay_triangles(self):
             [
                 GeometryCollection([Polygon([(0, 0), (1, 0), (1, 1), (0, 0)])]),
                 GeometryCollection([Polygon([(0, 1), (0, 0), (1, 1), (0, 1)])]),
-            ]
+            ],
+            crs=self.g3.crs,
         )
         dlt = self.g3.delaunay_triangles()
-        assert isinstance(dlt, GeoSeries)
-        assert_series_equal(expected, dlt)
+        assert_geoseries_equal(expected, dlt)
 
     def test_delaunay_triangles_pass_kwargs(self):
         expected = GeoSeries(
             [
                 MultiLineString([[(0, 0), (1, 1)], [(0, 0), (1, 0)], [(1, 0), (1, 1)]]),
                 MultiLineString([[(0, 1), (1, 1)], [(0, 0), (0, 1)], [(0, 0), (1, 1)]]),
-            ]
+            ],
+            crs=self.g3.crs,
         )
         dlt = self.g3.delaunay_triangles(only_edges=True)
-        assert isinstance(dlt, GeoSeries)
-        assert_series_equal(expected, dlt)
+        assert_geoseries_equal(expected, dlt)
 
     def test_delaunay_triangles_wrong_index(self):
         with pytest.raises(
@@ -1139,6 +1139,45 @@ def test_delaunay_triangles_wrong_index(self):
         ):
             self.g3.delaunay_triangles(tolerance=Series([0.1, 0.2], index=[99, 98]))
 
+    def test_voronoi_polygons(self):
+        expected = GeoSeries.from_wkt(
+            [
+                "POLYGON ((2 2, 2 0.5, 0.5 0.5, 0.5 2, 2 2))",
+                "POLYGON ((-1 2, 0.5 2, 0.5 0.5, -1 0.5, -1 2))",
+                "POLYGON ((-1 -1, -1 0.5, 0.5 0.5, 0.5 -1, -1 -1))",
+                "POLYGON ((2 -1, 0.5 -1, 0.5 0.5, 2 0.5, 2 -1))",
+            ],
+            crs=self.g1.crs,
+        )
+        vp = self.g1.voronoi_polygons()
+        assert_geoseries_equal(expected, vp)
+
+    def test_voronoi_polygons_only_edges(self):
+        expected = GeoSeries.from_wkt(
+            [
+                "LINESTRING (0.5 0.5, 0.5 2)",
+                "LINESTRING (2 0.5, 0.5 0.5)",
+                "LINESTRING (0.5 0.5, -1 0.5)",
+                "LINESTRING (0.5 0.5, 0.5 -1)",
+            ],
+            crs=self.g1.crs,
+        )
+        vp = self.g1.voronoi_polygons(only_edges=True)
+        assert_geoseries_equal(expected, vp, check_less_precise=True)
+
+    def test_voronoi_polygons_extend_to(self):
+        expected = GeoSeries.from_wkt(
+            [
+                "POLYGON ((3 3, 3 0.5, 0.5 0.5, 0.5 3, 3 3))",
+                "POLYGON ((-2 3, 0.5 3, 0.5 0.5, -2 0.5, -2 3))",
+                "POLYGON ((-2 -1, -2 0.5, 0.5 0.5, 0.5 -1, -2 -1))",
+                "POLYGON ((3 -1, 0.5 -1, 0.5 0.5, 3 0.5, 3 -1))",
+            ],
+            crs=self.g1.crs,
+        )
+        vp = self.g1.voronoi_polygons(extend_to=box(-2, 0, 3, 3))
+        assert_geoseries_equal(expected, vp)
+
     def test_exterior(self):
         exp_exterior = GeoSeries([LinearRing(p.boundary) for p in self.g3])
         for expected, computed in zip(exp_exterior, self.g3.exterior):