From aaf52fb2102799bde079291f1fa28a474e0727c9 Mon Sep 17 00:00:00 2001
From: Tom Russell <tom.russell@ouce.ox.ac.uk>
Date: Mon, 13 Nov 2023 16:52:33 +0000
Subject: [PATCH] Add merge_parts option to split_multilinestrings

Closes #46
---
 src/snkit/network.py |  9 ++++++-
 tests/test_init.py   | 60 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/src/snkit/network.py b/src/snkit/network.py
index 202e664..491dab8 100644
--- a/src/snkit/network.py
+++ b/src/snkit/network.py
@@ -279,16 +279,23 @@ def round_geometries(network: Network, precision: int = 3) -> Network:
     return network
 
 
-def split_multilinestrings(network: Network) -> Network:
+def split_multilinestrings(network: Network, merge_parts: bool = False) -> Network:
     """
     Create multiple edges from any MultiLineString edge
 
     Ensures that edge geometries are all LineStrings, duplicates attributes
     over any created multi-edges.
+
+    Parameters
+    ----------
+    merge_parts : default False
+        Merge parts of geometries if they are connected, before splitting to multiple edges
     """
 
     edges = network.edges
     geom_col: str = geometry_column_name(edges)
+    if merge_parts:
+        edges.geometry = edges.geometry.apply(merge_multilinestring)
     split_edges = edges.explode(column=geom_col, ignore_index=True)
 
     geo_types = set(split_edges.geom_type)
diff --git a/tests/test_init.py b/tests/test_init.py
index af150e4..3ae6f77 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -642,7 +642,7 @@ def test_split_multilinestrings():
     mls_coords = [
         (
             (0, 0),
-            (0, 1),
+            (0, 1.1),
         ),
         (
             (1, 1),
@@ -681,6 +681,64 @@ def test_split_multilinestrings():
     assert list(split_network.edges.index.values) == list(range(4))
 
 
+def test_split_multilinestrings_merge():
+    """Explode multilinestrings into linestrings - with parts merged if possible"""
+
+    # point coordinates comprising three linestrings
+    mls_coords = [
+        (
+            (0, 0),
+            (0, 1),  # should link to third part
+        ),
+        (
+            (1, 1),
+            (2, 2),
+            (2, 1),
+        ),
+        (
+            (0, 1),  # should merge with first part
+            (-1, 1),
+            (-1, 2),
+            (-1, 0),
+        ),
+    ]
+    # point coordsinate comprising a single linestring
+    ls_coords = [(4, 0), (4, 1), (4, 2)]
+
+    # make input network edges
+    edges = GeoDataFrame(
+        {
+            "data": ["MLS", "LS"],
+            "geometry": [MultiLineString(mls_coords), LineString(ls_coords)],
+        }
+    )
+    multi_network = snkit.network.Network(edges=edges)
+    # check we have two rows (multilinestring and linestring) in the input network
+    assert len(multi_network.edges) == 2
+
+    # split and check we have three rows (the resulting linestrings) in the output network
+    split_network = snkit.network.split_multilinestrings(multi_network)
+    assert len(split_network.edges) == 4
+
+    # check data is replicated from multilinestring to child linestrings
+    assert list(split_network.edges["data"].values) == ["MLS"] * 3 + ["LS"]
+
+    # check everything is reindexed
+    assert list(split_network.edges.index.values) == list(range(4))
+
+    # Check again, with merge_parts
+    split_network_merged = snkit.network.split_multilinestrings(
+        multi_network, merge_parts=True
+    )
+    assert len(split_network_merged.edges) == 3
+
+    # check data is replicated from multilinestring to child linestrings
+    assert list(split_network_merged.edges["data"].values) == ["MLS"] * 2 + ["LS"]
+
+    # check everything is reindexed
+    assert list(split_network_merged.edges.index.values) == list(range(3))
+
+
 def test_link_nodes_to_edges_within(gap, bridged):
     """Nodes should link to edges within a distance, splitting the node at a new point if
     necessary.