ENH: add polygonize_full (#298)

CHANGELOG.rst

@@ -37,6 +37,7 @@ Version 0.10 (unreleased)
 
 * Addition of a ``contains_properly`` function (#267)
 * Addition of a ``polygonize`` function (#275)
+* Addition of a ``polygonize_full`` function (#298)
 * Addition of a ``segmentize`` function for GEOS >= 3.10 (#299)
 
 **Bug fixes**

pygeos/constructive.py

@@ -23,6 +23,7 @@
     "normalize",
     "point_on_surface",
     "polygonize",
+    "polygonize_full",
     "reverse",
     "simplify",
     "snap",
@@ -564,6 +565,7 @@ def polygonize(geometries, **kwargs):
     See Also
     --------
     get_parts, get_geometry
+    polygonize_full
 
     Examples
     --------
@@ -578,6 +580,65 @@ def polygonize(geometries, **kwargs):
     return lib.polygonize(geometries, **kwargs)
 
 
+def polygonize_full(geometries, **kwargs):
+    """Creates polygons formed from the linework of a set of Geometries and
+    return all extra outputs as well.
+
+    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.
+
+    This function performs the same polygonization as `polygonize` but does
+    not only return the polygonal result but all extra outputs as well. The
+    return value consists of 4 elements:
+
+    * The polygonal valid output
+    * **Cut edges**: edges connected on both ends but not part of polygonal output
+    * **dangles**: edges connected on one end but not part of polygonal output
+    * **invalid rings**: polygons formed but which are not valid
+
+    This function returns the geometries within GeometryCollections.
+    Individual geometries can be obtained using `get_geometry` to get
+    a single geometry or `get_parts` to get an array of 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.
+    **kwargs
+        For other keyword-only arguments, see the
+        `NumPy ufunc docs <https://numpy.org/doc/stable/reference/ufuncs.html#ufuncs-kwargs>`_.
+
+    Returns
+    -------
+    (polgyons, cuts, dangles, invalid)
+        tuple of 4 GeometryCollections or arrays of GeometryCollections
+
+    See Also
+    --------
+    polygonize
+
+    Examples
+    --------
+    >>> lines = [
+    ...     Geometry("LINESTRING (0 0, 1 1)"),
+    ...     Geometry("LINESTRING (0 0, 0 1, 1 1)"),
+    ...     Geometry("LINESTRING (0 1, 1 1)"),
+    ... ]
+    >>> polygonize_full(lines)  # doctest: +NORMALIZE_WHITESPACE
+    (<pygeos.Geometry GEOMETRYCOLLECTION (POLYGON ((1 1, 0 0, 0 1, 1 1)))>,
+     <pygeos.Geometry GEOMETRYCOLLECTION EMPTY>,
+     <pygeos.Geometry GEOMETRYCOLLECTION (LINESTRING (0 1, 1 1))>,
+     <pygeos.Geometry GEOMETRYCOLLECTION EMPTY>)
+    """
+    return lib.polygonize_full(geometries, **kwargs)
+
+
 @requires_geos("3.7.0")
 @multithreading_enabled
 def reverse(geometry, **kwargs):

pygeos/test/test_constructive.py

@@ -477,6 +477,113 @@ def test_polygonize_missing():
     assert result == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
 
 
+def test_polygonize_full():
+    lines = [
+        None,
+        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)"),
+        None,
+        pygeos.Geometry("LINESTRING (1 0, 0 0)"),
+        pygeos.Geometry("LINESTRING (5 5, 6 6)"),
+        pygeos.Geometry("LINESTRING (1 1, 100 100)"),
+        pygeos.Geometry("POINT (0 0)"),
+        None,
+    ]
+    result = pygeos.polygonize_full(lines)
+    assert len(result) == 4
+    assert all(pygeos.get_type_id(geom) == 7 for geom in result)  # GeometryCollection
+    polygons, cuts, dangles, invalid = result
+    expected_polygons = pygeos.Geometry(
+        "GEOMETRYCOLLECTION (POLYGON ((0 0, 1 1, 1 0, 0 0)), POLYGON ((1 1, 0 0, 0 1, 1 1)))"
+    )
+    assert polygons == expected_polygons
+    assert cuts == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
+    expected_dangles = pygeos.Geometry(
+        "GEOMETRYCOLLECTION (LINESTRING (1 1, 100 100), LINESTRING (5 5, 6 6))"
+    )
+    assert dangles == expected_dangles
+    assert invalid == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
+
+
+def test_polygonize_full_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_full(np.array(lines))
+    assert len(result) == 4
+    assert all(isinstance(geom, pygeos.Geometry) for geom in result)
+    assert result[0] == expected
+    assert all(
+        geom == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY") for geom in result[1:]
+    )
+
+    result = pygeos.polygonize_full(np.array([lines]))
+    assert len(result) == 4
+    assert all(isinstance(geom, np.ndarray) for geom in result)
+    assert all(geom.shape == (1,) for geom in result)
+    assert result[0][0] == expected
+    assert all(
+        geom[0] == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY") for geom in result[1:]
+    )
+
+    arr = np.array([lines, lines])
+    assert arr.shape == (2, 3)
+    result = pygeos.polygonize_full(arr)
+    assert len(result) == 4
+    assert all(isinstance(arr, np.ndarray) for arr in result)
+    assert all(arr.shape == (2,) for arr in result)
+    assert result[0][0] == expected
+    assert result[0][1] == expected
+    assert all(
+        g == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
+        for geom in result[1:]
+        for g in geom
+    )
+
+    arr = np.array([[lines, lines], [lines, lines], [lines, lines]])
+    assert arr.shape == (3, 2, 3)
+    result = pygeos.polygonize_full(arr)
+    assert len(result) == 4
+    assert all(isinstance(arr, np.ndarray) for arr in result)
+    assert all(arr.shape == (3, 2) for arr in result)
+    for res in result[0].flatten():
+        assert res == expected
+    for arr in result[1:]:
+        for res in arr.flatten():
+            assert res == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
+
+
+@pytest.mark.skipif(
+    np.__version__ < "1.15",
+    reason="axis keyword for generalized ufunc introduced in np 1.15",
+)
+def test_polygonize_full_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_full(arr, axis=1)
+    assert len(result) == 4
+    assert all(arr.shape == (2,) for arr in result)
+    result = pygeos.polygonize_full(arr, axis=0)
+    assert len(result) == 4
+    assert all(arr.shape == (3,) for arr in result)
+
+
+def test_polygonize_full_missing():
+    # set of geometries that is all missing
+    result = pygeos.polygonize_full([None, None])
+    assert len(result) == 4
+    assert all(geom == pygeos.Geometry("GEOMETRYCOLLECTION EMPTY") for geom in result)
+
+
 @pytest.mark.skipif(pygeos.geos_version < (3, 10, 0), reason="GEOS < 3.10")
 @pytest.mark.parametrize("geometry", all_types)
 @pytest.mark.parametrize("tolerance", [-1, 0])

src/fast_loop_macros.h

@@ -65,6 +65,16 @@
   cp1 = ip1;                      \
   for (i_c1 = 0; i_c1 < n_c1; i_c1++, cp1 += cs1)
 
+/** (ip1, cp1) -> (op1, op2, op3, op4) */
+#define SINGLE_COREDIM_LOOP_OUTER_NOUT4                                                \
+  char *ip1 = args[0], *op1 = args[1], *op2 = args[2], *op3 = args[3], *op4 = args[4], \
+       *cp1;                                                                           \
+  npy_intp is1 = steps[0], os1 = steps[1], os2 = steps[2], os3 = steps[3],             \
+           os4 = steps[4], cs1 = steps[5];                                             \
+  npy_intp n = dimensions[0], n_c1 = dimensions[1];                                    \
+  npy_intp i, i_c1;                                                                    \
+  for (i = 0; i < n; i++, ip1 += is1, op1 += os1, op2 += os2, op3 += os3, op4 += os4)
+
 /** (ip1, cp1, cp2) -> (op1) */
 #define DOUBLE_COREDIM_LOOP_OUTER                                          \
   char *ip1 = args[0], *op1 = args[1], *cp1, *cp2;                         \

src/ufuncs.c

@@ -22,6 +22,12 @@
   Py_XDECREF(*out);                                      \
   *out = ret
 
+#define OUTPUT_Y_I(I, RET_PTR)                              \
+  PyObject* ret##I = GeometryObject_FromGEOS(RET_PTR, ctx); \
+  PyObject** out##I = (PyObject**)op##I;                    \
+  Py_XDECREF(*out##I);                                      \
+  *out##I = ret##I
+
 // Fail if inputs output multiple times on the same place in memory. That would
 // lead to segfaults as the same GEOSGeometry would be 'owned' by multiple PyObjects.
 #define CHECK_NO_INPLACE_OUTPUT(N)                                                \
@@ -1229,7 +1235,7 @@ static PyUFuncGenericFunction YYd_Y_funcs[1] = {&YYd_Y_func};
 
 /* Define functions with unique call signatures */
 static char box_dtypes[6] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
-                             NPY_DOUBLE, NPY_BOOL, NPY_OBJECT};
+                             NPY_DOUBLE, NPY_BOOL,   NPY_OBJECT};
 static void box_func(char** args, npy_intp* dimensions, npy_intp* steps, void* data) {
   char *ip1 = args[0], *ip2 = args[1], *ip3 = args[2], *ip4 = args[3], *ip5 = args[4];
   npy_intp is1 = steps[0], is2 = steps[1], is3 = steps[2], is4 = steps[3], is5 = steps[4];
@@ -1844,6 +1850,83 @@ static void polygonize_func(char** args, npy_intp* dimensions, npy_intp* steps,
 }
 static PyUFuncGenericFunction polygonize_funcs[1] = {&polygonize_func};
 
+static char polygonize_full_dtypes[5] = {NPY_OBJECT, NPY_OBJECT, NPY_OBJECT, NPY_OBJECT,
+                                         NPY_OBJECT};
+static void polygonize_full_func(char** args, npy_intp* dimensions, npy_intp* steps,
+                                 void* data) {
+  GEOSGeometry* geom = NULL;
+  GEOSGeometry* geom_copy = NULL;
+  unsigned int n_geoms;
+
+  GEOSGeometry* collection = NULL;
+  GEOSGeometry* cuts = NULL;
+  GEOSGeometry* dangles = NULL;
+  GEOSGeometry* invalidRings = NULL;
+
+  GEOS_INIT;
+
+  GEOSGeometry** geoms = malloc(sizeof(void*) * dimensions[1]);
+  if (geoms == NULL) {
+    errstate = PGERR_NO_MALLOC;
+    goto finish;
+  }
+
+  SINGLE_COREDIM_LOOP_OUTER_NOUT4 {
+    n_geoms = 0;
+    SINGLE_COREDIM_LOOP_INNER {
+      if (!get_geom(*(GeometryObject**)cp1, &geom)) {
+        errstate = PGERR_NOT_A_GEOMETRY;
+        goto finish;
+      }
+      if (geom == NULL) {
+        continue;
+      }
+      // need to copy the input geometries, because the Collection takes ownership
+      geom_copy = GEOSGeom_clone_r(ctx, geom);
+      if (geom_copy == NULL) {
+        // if something went wrong before creating the collection, destroy previously
+        // cloned geoms
+        for (i = 0; i < n_geoms; i++) {
+          GEOSGeom_destroy_r(ctx, geoms[i]);
+        }
+        errstate = PGERR_GEOS_EXCEPTION;
+        goto finish;
+      }
+      geoms[n_geoms] = geom_copy;
+      n_geoms++;
+    }
+    collection =
+        GEOSGeom_createCollection_r(ctx, GEOS_GEOMETRYCOLLECTION, geoms, n_geoms);
+    if (collection == NULL) {
+      errstate = PGERR_GEOS_EXCEPTION;
+      goto finish;
+    }
+
+    GEOSGeometry* ret_ptr =
+        GEOSPolygonize_full_r(ctx, collection, &cuts, &dangles, &invalidRings);
+    if (ret_ptr == NULL) {
+      errstate = PGERR_GEOS_EXCEPTION;
+      goto finish;
+    }
+    OUTPUT_Y_I(1, ret_ptr);
+    OUTPUT_Y_I(2, cuts);
+    OUTPUT_Y_I(3, dangles);
+    OUTPUT_Y_I(4, invalidRings);
+    GEOSGeom_destroy_r(ctx, collection);
+    collection = NULL;
+  }
+
+finish:
+  if (collection != NULL) {
+    GEOSGeom_destroy_r(ctx, collection);
+  }
+  if (geoms != NULL) {
+    free(geoms);
+  }
+  GEOS_FINISH;
+}
+static PyUFuncGenericFunction polygonize_full_funcs[1] = {&polygonize_full_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,
@@ -2790,6 +2873,12 @@ TODO relate functions
                                               SIGNATURE);                                \
   PyDict_SetItemString(d, #NAME, ufunc)
 
+#define DEFINE_GENERALIZED_NOUT4(NAME, N_IN, SIGNATURE)                                  \
+  ufunc = PyUFunc_FromFuncAndDataAndSignature(NAME##_funcs, null_data, NAME##_dtypes, 1, \
+                                              N_IN, 4, PyUFunc_None, #NAME, "", 0,       \
+                                              SIGNATURE);                                \
+  PyDict_SetItemString(d, #NAME, ufunc)
+
 int init_ufuncs(PyObject* m, PyObject* d) {
   PyObject* ufunc;
 
@@ -2885,6 +2974,7 @@ int init_ufuncs(PyObject* m, PyObject* d) {
   DEFINE_CUSTOM(relate, 2);
   DEFINE_CUSTOM(relate_pattern, 3);
   DEFINE_GENERALIZED(polygonize, 1, "(d)->()");
+  DEFINE_GENERALIZED_NOUT4(polygonize_full, 1, "(d)->(),(),(),()");
 
   DEFINE_GENERALIZED(points, 1, "(d)->()");
   DEFINE_GENERALIZED(linestrings, 1, "(i, d)->()");