ENH: add polygonize ufunc (#275)

CHANGELOG.rst

@@ -16,6 +16,7 @@ Version 0.10 (unreleased)
 **Added GEOS functions**
 
 * Addition of a ``contains_properly`` function (#267).
+* Addition of a ``polygonize`` function (#275).
 
 **Bug fixes**
 * Fixed portability issue for ARM architecture (#293)

pygeos/constructive.py

@@ -21,6 +21,7 @@
     "make_valid",
     "normalize",
     "point_on_surface",
+    "polygonize",
     "reverse",
     "simplify",
     "snap",
@@ -486,6 +487,54 @@ def point_on_surface(geometry, **kwargs):
     return lib.point_on_surface(geometry, **kwargs)
 
 
+def polygonize(geometries, **kwargs):
+    """Creates polygons formed from the linework of a set of Geometries.
+
+    Polygonizes an array of Geometries that contain linework which
+    represents the edges of a planar graph. Any type of Geometry may be
+    provided as input; only the constituent lines and rings will be used to
+    create the output polygons.
+
+    Lines or rings that when combined do not completely close a polygon
+    will result in an empty GeometryCollection.  Duplicate segments are
+    ignored.
+
+    This function returns the polygons within a GeometryCollection.
+    Individual Polygons can be obtained using `get_geometry` to get
+    a single polygon or `get_parts` to get an array of polygons.
+    MultiPolygons can be constructed from the output using
+    ``pygeos.multipolygons(pygeos.get_parts(pygeos.polygonize(geometries)))``.
+
+    Parameters
+    ----------
+    geometries : array_like
+        An array of geometries.
+    axis : int
+        Axis along which the geometries are polygonized.
+        The default is to perform a reduction over the last dimension
+        of the input array. A 1D array results in a scalar geometry.
+
+    Returns
+    -------
+    GeometryCollection or array of GeometryCollections
+
+    See Also
+    --------
+    get_parts, get_geometry
+
+    Examples
+    --------
+    >>> lines = [
+    ...     Geometry("LINESTRING (0 0, 1 1)"),
+    ...     Geometry("LINESTRING (0 0, 0 1)"),
+    ...     Geometry("LINESTRING (0 1, 1 1)"),
+    ... ]
+    >>> polygonize(lines)
+    <pygeos.Geometry GEOMETRYCOLLECTION (POLYGON ((1 1, 0 0, 0 1, 1 1)))>
+    """
+    return lib.polygonize(geometries, **kwargs)
+
+
 @requires_geos("3.7.0")
 @multithreading_enabled
 def reverse(geometry, **kwargs):

pygeos/test/test_constructive.py

@@ -391,3 +391,80 @@ def test_clip_by_rect_non_scalar_kwargs():
     msg = "only accepts scalar values"
     with pytest.raises(TypeError, match=msg):
         pygeos.clip_by_rect([line_string, line_string], 0, 0, 1, np.array([0, 1]))
+
+
+def test_polygonize():
+    lines = [
+        pygeos.Geometry("LINESTRING (0 0, 1 1)"),
+        pygeos.Geometry("LINESTRING (0 0, 0 1)"),
+        pygeos.Geometry("LINESTRING (0 1, 1 1)"),
+        pygeos.Geometry("LINESTRING (1 1, 1 0)"),
+        pygeos.Geometry("LINESTRING (1 0, 0 0)"),
+        pygeos.Geometry("LINESTRING (5 5, 6 6)"),
+        pygeos.Geometry("POINT (0 0)"),
+        None,
+    ]
+    result = pygeos.polygonize(lines)
+    assert pygeos.get_type_id(result) == 7  # GeometryCollection
+    expected = pygeos.Geometry(
+        "GEOMETRYCOLLECTION (POLYGON ((0 0, 1 1, 1 0, 0 0)), POLYGON ((1 1, 0 0, 0 1, 1 1)))"
+    )
+    assert result == expected
+
+
+def test_polygonize_array():
+    lines = [
+        pygeos.Geometry("LINESTRING (0 0, 1 1)"),
+        pygeos.Geometry("LINESTRING (0 0, 0 1)"),
+        pygeos.Geometry("LINESTRING (0 1, 1 1)"),
+    ]
+    expected = pygeos.Geometry(
+        "GEOMETRYCOLLECTION (POLYGON ((1 1, 0 0, 0 1, 1 1)))"
+    )
+    result = pygeos.polygonize(np.array(lines))
+    assert isinstance(result, pygeos.Geometry)
+    assert result == expected
+
+    result = pygeos.polygonize(np.array([lines]))
+    assert isinstance(result, np.ndarray)
+    assert result.shape == (1,)
+    assert result[0] == expected
+
+    arr = np.array([lines, lines])
+    assert arr.shape == (2, 3)
+    result = pygeos.polygonize(arr)
+    assert isinstance(result, np.ndarray)
+    assert result.shape == (2,)
+    assert result[0] == expected
+    assert result[1] == expected
+
+    arr = np.array([[lines, lines], [lines, lines], [lines, lines]])
+    assert arr.shape == (3, 2, 3)
+    result = pygeos.polygonize(arr)
+    assert isinstance(result, np.ndarray)
+    assert result.shape == (3, 2)
+    for res in result.flatten():
+        assert res == expected
+
+
+@pytest.mark.skipif(
+    np.__version__ < "1.15",
+    reason="axis keyword for generalized ufunc introduced in np 1.15",
+)
+def test_polygonize_array_axis():
+    lines = [
+        pygeos.Geometry("LINESTRING (0 0, 1 1)"),
+        pygeos.Geometry("LINESTRING (0 0, 0 1)"),
+        pygeos.Geometry("LINESTRING (0 1, 1 1)"),
+    ]
+    arr = np.array([lines, lines])  # shape (2, 3)
+    result = pygeos.polygonize(arr, axis=1)
+    assert result.shape == (2,)
+    result = pygeos.polygonize(arr, axis=0)
+    assert result.shape == (3,)
+
+
+def test_polygonize_missing():
+    # set of geometries that is all missing
+    result = pygeos.polygonize([None, None])
+    assert result == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")

src/ufuncs.c

@@ -1776,6 +1776,50 @@ static void relate_pattern_func(char** args, npy_intp* dimensions, npy_intp* ste
 }
 static PyUFuncGenericFunction relate_pattern_funcs[1] = {&relate_pattern_func};
 
+static char polygonize_dtypes[2] = {NPY_OBJECT, NPY_OBJECT};
+static void polygonize_func(char** args, npy_intp* dimensions, npy_intp* steps,
+                            void* data) {
+  GEOSGeometry* geom = NULL;
+  unsigned int n_geoms;
+
+  GEOS_INIT;
+
+  GEOSGeometry** geoms = malloc(sizeof(void*) * dimensions[1]);
+  if (geoms == NULL) {
+    errstate = PGERR_NO_MALLOC;
+    goto finish;
+  }
+
+  SINGLE_COREDIM_LOOP_OUTER {
+    n_geoms = 0;
+    SINGLE_COREDIM_LOOP_INNER {
+      if (!get_geom(*(GeometryObject**)cp1, &geom)) {
+        errstate = PGERR_NOT_A_GEOMETRY;
+        goto finish;
+      }
+      if (geom == NULL) {
+        continue;
+      }
+      geoms[n_geoms] = geom;
+      n_geoms++;
+    }
+
+    GEOSGeometry* ret_ptr = GEOSPolygonize_r(ctx, geoms, n_geoms);
+    if (ret_ptr == NULL) {
+      errstate = PGERR_GEOS_EXCEPTION;
+      goto finish;
+    }
+    OUTPUT_Y;
+  }
+
+finish:
+  if (geoms != NULL) {
+    free(geoms);
+  }
+  GEOS_FINISH;
+}
+static PyUFuncGenericFunction polygonize_funcs[1] = {&polygonize_func};
+
 #if GEOS_SINCE_3_6_0
 static char set_precision_dtypes[4] = {NPY_OBJECT, NPY_DOUBLE, NPY_BOOL, NPY_OBJECT};
 static void set_precision_func(char** args, npy_intp* dimensions, npy_intp* steps,
@@ -2701,6 +2745,8 @@ int init_ufuncs(PyObject* m, PyObject* d) {
   DEFINE_CUSTOM(is_valid_reason, 1);
   DEFINE_CUSTOM(relate, 2);
   DEFINE_CUSTOM(relate_pattern, 3);
+  DEFINE_GENERALIZED(polygonize, 1, "(d)->()");
+
   DEFINE_GENERALIZED(points, 1, "(d)->()");
   DEFINE_GENERALIZED(linestrings, 1, "(i, d)->()");
   DEFINE_GENERALIZED(linearrings, 1, "(i, d)->()");