ENH: add option extend to function centerline

pygeoops · Jun 17, 2024 · 9a46468 · 9a46468
1 parent 14b0b68
commit 9a46468
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 11 deletions.
diff --git a/pygeoops/_centerline.py b/pygeoops/_centerline.py
@@ -9,6 +9,7 @@
 from shapely.geometry.base import BaseGeometry
 
 from pygeoops._general import _extract_0dim_ndarray
+from pygeoops import _extend_line
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +19,7 @@ def centerline(
     densify_distance: float = -1,
     min_branch_length: float = -1,
     simplifytolerance: float = -0.25,
+    extend: bool = False,
 ) -> Union[BaseGeometry, NDArray[BaseGeometry], GeoSeries, None]:
     """
     Calculates an approximated centerline for a polygon.
@@ -54,6 +56,7 @@ def centerline(
               - value = 0: no simplify
               - value > 0: simplify with this value as tolerance
               - value < 0: simplifytolerance = average width of geometry * abs(value)
+        extend (bool, optional): extend the centerline to the edge of the geometry.
 
     Returns:
         geometry, GeoSeries or array_like: the centerline for each of the input
@@ -75,6 +78,7 @@ def centerline(
                     densify_distance=densify_distance,
                     min_branch_length=min_branch_length,
                     simplifytolerance=simplifytolerance,
+                    extend=extend,
                 )
                 for geom in geometry
             ]
@@ -88,6 +92,7 @@ def centerline(
             densify_distance=densify_distance,
             min_branch_length=min_branch_length,
             simplifytolerance=simplifytolerance,
+            extend=extend,
         )
 
 
@@ -96,6 +101,7 @@ def _centerline(
     densify_distance: float = -1,
     min_branch_length: float = -1,
     simplifytolerance: float = -0.25,
+    extend: bool = False,
 ) -> Optional[BaseGeometry]:
     if geom is None or geom.is_empty:
         return None
@@ -152,6 +158,9 @@ def _centerline(
             tol = abs(simplifytolerance) * average_width
         lines = shapely.simplify(lines, tol)
 
+    if extend:
+        lines = _extend_line.extend_line_to_geometry(lines, geom)
+
     # Return result
     lines = shapely.normalize(lines)
     return lines
@@ -201,7 +210,7 @@ def _remove_short_branches_notempty(
             line, min_branch_length, remove_one_by_one=True
         )
 
-    # If still an empty result, return original version
+    # If still an empty result, retain original version
     if line_cleaned is None or line_cleaned.is_empty:
         line_cleaned = line
 

diff --git a/tests/test_centerline.py b/tests/test_centerline.py
@@ -91,37 +91,43 @@ def test_centerline_None_geometry():
 
 
 @pytest.mark.parametrize(
-    "test, min_branch_length, poly_wkt, expected_centerline_wkt",
+    "test, min_branch_length, poly_wkt, "
+    "expected_centerline_wkt, expected_centerline_extend_wkt",
     [
         (
-            "input: elleptical shape, resulted in small branch not being removed",
+            "elleptical shape, resulted in small branch not being removed",
             0.0,
-            "MULTIPOLYGON (((0 1, 1 3.25, 2 4.5, 3 5.75, 3.5 6.25, 5 3.25, 3.75 1.75, 2.5 0.5, 1 0, 0 1)))",  # noqa: E501
+            "POLYGON ((0 1, 1 3.25, 2 4.5, 3 5.75, 3.5 6.25, 5 3.25, 3.75 1.75, 2.5 0.5, 1 0, 0 1))",  # noqa: E501
             "MULTILINESTRING ((3.2641509433962264 3.3726415094339623, 3.7916666666666665 5.458333333333333), (3.2641509433962264 3.3726415094339623, 3.34375 3.359375), (1.375 1.375, 3.2641509433962264 3.3726415094339623))",  # noqa: E501
+            "MULTILINESTRING ((3.2641509433962264 3.3726415094339623, 4.878048780487804 3.1036585365853644), (3.2641509433962264 3.3726415094339623, 3.8266583229036297 5.5966833541927405), (0.5244235436893204 0.4755764563106795, 3.2641509433962264 3.3726415094339623))",  # noqa: E501
         ),
         (
-            "input: elleptical shape, resulted in small branch not being removed",
+            "elleptical shape, resulted in small branch not being removed",
             -1.0,
-            "MULTIPOLYGON (((0 1, 1 3.25, 2 4.5, 3 5.75, 3.5 6.25, 5 3.25, 3.75 1.75, 2.5 0.5, 1 0, 0 1)))",  # noqa: E501
+            "POLYGON ((0 1, 1 3.25, 2 4.5, 3 5.75, 3.5 6.25, 5 3.25, 3.75 1.75, 2.5 0.5, 1 0, 0 1))",  # noqa: E501
             "LINESTRING (1.375 1.375, 3.7916666666666665 5.458333333333333)",
+            "LINESTRING (0.7243589743589742 0.2756410256410258, 3.8481308411214954 5.553738317757009)",  # noqa: E501
         ),
         (
-            "input: fancy L shape, output: L-ish line",
+            "fancy L shape, output: L-ish line",
             0.0,
             "POLYGON ((0 0, 0 8, -2 10, 4 10, 2 8, 2 2, 10 2, 10 0, 0 0))",
             "MULTILINESTRING ((8.87687074829932 0.9829931972789112, 9.2 1.5), (8.87687074829932 0.9829931972789112, 9.166666666666666 0.5), (1.1367816091954022 1.1160919540229888, 8.87687074829932 0.9829931972789112), (1 8.75, 3.25 9.75), (1 8.75, 1.1367816091954022 1.1160919540229888), (0.833333333333333 0.8, 1.1367816091954022 1.1160919540229888), (-1.25 9.75, 1 8.75))",  # noqa: E501
+            "MULTILINESTRING ((8.87687074829932 0.9829931972789112, 9.5125 2), (8.87687074829932 0.9829931972789112, 9.466666666666667 0), (1.1367816091954022 1.1160919540229888, 8.87687074829932 0.9829931972789112), (1 8.75, 3.8125000000000004 10), (1 8.75, 1.1367816091954022 1.1160919540229888), (0.0653333333333331 0, 1.1367816091954022 1.1160919540229888), (-1.8124999999999996 10, 1 8.75))",  # noqa: E501
         ),
         (
-            "input: fancy L shape, output: L-ish line",
+            "fancy L shape, output: L-ish line",
             -1.0,
             "POLYGON ((0 0, 0 8, -2 10, 4 10, 2 8, 2 2, 10 2, 10 0, 0 0))",
             "MULTILINESTRING ((1 8.75, 1.1367816091954022 1.1160919540229888, 8.87687074829932 0.9829931972789112), (1 8.75, 3.25 9.75), (-1.25 9.75, 1 8.75))",  # noqa: E501
+            "MULTILINESTRING ((1 8.75, 1.1367816091954022 1.1160919540229888, 10 0.9636798399806034), (1 8.75, 3.8125000000000004 10), (-1.8124999999999996 10, 1 8.75))",  # noqa: E501
         ),
         (
-            "input: L shape, output: L line",
+            "L shape, output: L line",
             -1.0,
             "POLYGON ((0 0, 0 10, 2 10, 2 2, 10 2, 10 0, 0 0))",
             "LINESTRING (1 9, 1 1, 9 1)",
+            "LINESTRING (1 10, 1 1, 10 1)",
         ),
     ],
 )
@@ -131,8 +137,13 @@ def test_centerline_poly(
     min_branch_length: float,
     poly_wkt: str,
     expected_centerline_wkt: str,
+    expected_centerline_extend_wkt: str,
 ):
-    """More complicated polygon tests."""
+    """More complicated polygon tests.
+
+    Includes tests on extend=True as this makes it easier to compare both options in the
+    output plots.
+    """
     poly = shapely.from_wkt(poly_wkt)
     centerline = pygeoops.centerline(poly, min_branch_length=min_branch_length)
     assert centerline is not None
@@ -141,4 +152,19 @@ def test_centerline_poly(
     test_helper.plot([poly, centerline], output_path)
     assert (
         centerline.wkt == expected_centerline_wkt
-    ), f"test descr: {test}, {min_branch_length}"
+    ), f"test descr: {test}, {min_branch_length}, with extend=False"
+    centerline = None
+
+    # Test same input polygons with extend=True
+    centerline_extend = pygeoops.centerline(
+        poly, min_branch_length=min_branch_length, extend=True
+    )
+    assert centerline_extend is not None
+    assert isinstance(centerline_extend, BaseGeometry)
+    output_path = (
+        tmp_path / f"test_centerline_poly_{test}_{min_branch_length}_extend.png"
+    )
+    test_helper.plot([poly, centerline_extend], output_path)
+    assert (
+        centerline_extend.wkt == expected_centerline_extend_wkt
+    ), f"test descr: {test}, {min_branch_length}, with extend=True"