Method to split links depending on mode (#247)

* add convenience method for stripping all links of a mode * take out the method of splitting links on mode from the cli command, add tests * add docstrings for splitting, add handling empty prefix * update docstring * update CHANGELOG.md * add more details to the docstring * address PR comments: - add a step to example in docstring - refine assertions in tests
arup-group · Nov 13, 2024 · 8f35b8c · 8f35b8c
1 parent 0bf392d
commit 8f35b8c
Show file tree

Hide file tree

Showing 4 changed files with 145 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 * Convenience method to strip all links of a given mode [#243](https://github.com/arup-group/genet/issues/243)
+* Method to split links on mode. New links are generated of given mode based on existing links [#244](https://github.com/arup-group/genet/issues/244)
 
 ### Fixed
 

diff --git a/src/genet/cli.py b/src/genet/cli.py
@@ -1045,8 +1045,8 @@ def separate_modes_in_network(
         [out] {"id": "bike---LINK_ID", "modes": {"bike"}, "freespeed": 5, ...}
         ```
 
-        In the case when a link already has a single dedicated mode, no updates are made to the link ID.
-        You can assume that all links that were in the network previously are still there, but their allowed modes may have changed.
+        In the case when a link already has a single dedicated mode, it will be replaced aby a link with a new, prefixed ID.
+        Other links modes will have changed also.
         So, any simulation outputs may not be valid with this new network.
     """
     modes = modes.split(",")
@@ -1061,23 +1061,12 @@ def separate_modes_in_network(
 
     for mode in modes:
         logging.info(f"Splitting links for mode: {mode}")
-        df = network.link_attribute_data_under_key("modes")
-        modal_links = network.links_on_modal_condition({mode})
-        # leave the links that have a single mode as they are
-        modal_links = set(modal_links) & set(df[df != {mode}].index)
-        update_mode_links = {k: {"modes": df.loc[k] - {mode}} for k in modal_links}
-        new_links = {
-            f"{mode}---{k}": {**network.link(k), **{"modes": {mode}, "id": f"{mode}---{k}"}}
-            for k in modal_links
-        }
-        network.apply_attributes_to_links(update_mode_links)
-        network.add_links(new_links)
+        new_links = network.split_links_on_mode(mode)
         if increase_capacity:
             logging.info(f"Increasing capacity for link of mode {mode} to 9999")
-            mode_links = network.extract_links_on_edge_attributes({"modes": mode})
-            df_capacity = network.link_attribute_data_under_keys(["capacity"]).loc[mode_links, :]
-            df_capacity["capacity"] = 9999
-            network.apply_attributes_to_links(df_capacity.T.to_dict())
+            network.apply_attributes_to_links(
+                {link_id: {"capacity": 9999} for link_id in new_links}
+            )
 
     logging.info(f"Number of links after separating graph: {len(network.link_id_mapping)}")
 

diff --git a/src/genet/core.py b/src/genet/core.py
@@ -5,7 +5,7 @@
 import traceback
 import uuid
 from copy import deepcopy
-from typing import Any, Callable, Iterator, Literal, Optional, Union
+from typing import Any, Callable, Iterator, Literal, Optional, Set, Union
 
 import geopandas as gpd
 import networkx as nx
@@ -784,6 +784,61 @@ def empty_modes(mode_attrib):
         )
         self.remove_links(no_mode_links)
 
+    def split_links_on_mode(self, mode: str, link_id_prefix: Optional[str] = None) -> Set[str]:
+        """Method to split links depending on mode.
+        Existing links with mode `mode` will have that mode removed.
+        New links will be added with only the mode `mode` and inheriting data from the link they originated from.
+        The IDs of new link IDs will by default identify the mode, but can be changed with `link_id_prefix`.
+
+        Examples:
+            ```python
+            [1] network.link("LINK_ID")
+            [out] {"id": "LINK_ID", "modes": {"car", "bike"}, "freespeed": 5, ...}
+            ```
+
+            ```python
+            [2] network.split_links_on_mode("bike")
+            [out] {"bike---LINK_ID"}
+            ```
+
+            The new bike link will assume all the same attributes apart from the "modes":
+            ```python
+            [3] network.link("bike---LINK_ID")`
+            [out] {"id": "bike---LINK_ID", "modes": {"bike"}, "freespeed": 5, ...}
+            ```
+
+            The old link will have the `bike` mode removed
+            ```python
+            [4] network.link("LINK_ID")
+            [out] {"id": "LINK_ID", "modes": {"car"}, "freespeed": 5, ...}
+            ```
+
+        Args:
+            mode (str): Mode to split from the links.
+            link_id_prefix (str): Optional, you can request what the
+
+        Returns:
+            Set of link IDs of the new links
+        """
+        modal_links = self.links_on_modal_condition({mode})
+        modal_links = list(modal_links)
+
+        if link_id_prefix == "":
+            logging.warning("Empty string was set as prefix, the IDs will be randomly assigned")
+            new_link_ids = self.generate_indices_for_n_edges(len(modal_links))
+        else:
+            if link_id_prefix is None:
+                link_id_prefix = f"{mode}---"
+            new_link_ids = [f"{link_id_prefix}{link_id}" for link_id in modal_links]
+        new_links = {
+            new_link_id: {**self.link(old_link_id), **{"modes": {mode}, "id": new_link_id}}
+            for new_link_id, old_link_id in zip(new_link_ids, modal_links)
+        }
+
+        self.remove_mode_from_links(modal_links, mode)
+        self.add_links(new_links)
+        return set(new_links.keys())
+
     def retain_n_connected_subgraphs(self, n: int, mode: str):
         """Method to remove modes in-place from link which do not belong to largest connected n components.
 

diff --git a/tests/test_core_network.py b/tests/test_core_network.py
@@ -2404,6 +2404,88 @@ def test_removing_mode_from_links_removes_empty_links():
     assert n.has_link("0")
 
 
+def test_splitting_links_removes_required_mode_from_existing_links():
+    n = Network("epsg:27700")
+    n.add_link("0", 1, 2, attribs={"modes": {"car", "bike"}, "length": 1})
+
+    n.split_links_on_mode("bike")
+
+    assert n.link("0")["modes"] == {"car"}, "The link modes have been incorrectly set"
+
+
+def test_splitting_links_creates_new_links_of_required_mode():
+    n = Network("epsg:27700")
+    n.add_link("0", 1, 2, attribs={"modes": {"car", "bike"}, "length": 1})
+
+    new_links = n.split_links_on_mode("bike")
+
+    assert len(new_links) == 1, "A new link was not generated"
+    new_link = list(new_links)[0]
+    assert n.has_link(new_link), "The new link was not found in the Network"
+    assert n.link(new_link)["modes"] == {"bike"}, "The mode of the new link is wrong"
+
+
+def test_data_retained_when_splitting_links():
+    n = Network("epsg:27700")
+    n.add_link("0", 1, 2, attribs={"modes": {"car", "bike"}, "length": 1, "BIG": "DATA"})
+
+    new_links = n.split_links_on_mode("bike")
+
+    assert len(new_links) == 1, "A new link was not generated"
+    new_link = list(new_links)[0]
+    assert "BIG" in n.link(new_link), "Attribute key was not found in the data saved on the link"
+    assert (
+        n.link(new_link)["BIG"] == "DATA"
+    ), "The new link did not inherit the same attribute value"
+
+
+def test_splitting_links_uses_desired_prefix_for_new_link_ids():
+    n = Network("epsg:27700")
+    n.add_link("0", 1, 2, attribs={"modes": {"car", "bike"}, "length": 1})
+
+    new_links = n.split_links_on_mode("bike", link_id_prefix="HEYO-")
+
+    assert len(new_links) == 1, "No new links were generated"
+    assert list(new_links)[0].startswith(
+        "HEYO-"
+    ), "The ID of a new link did not start with the desired prefix"
+
+
+def test_splitting_links_generates_unique_ids_for_new_links():
+    n = Network("epsg:27700")
+    for i in range(10):
+        n.add_link(str(i), i, i + 1, attribs={"modes": {"car", "bike"}, "length": 1})
+
+    new_links = n.split_links_on_mode("bike", link_id_prefix="HEYO-")
+
+    assert isinstance(new_links, set)  # To force uniqueness
+    assert len(new_links) == 10, "The number of link IDs is incorrect"
+    for link_id in new_links:
+        assert (
+            n.link(link_id)["id"] == link_id
+        ), "The ID declared in attributes does not match the ID in the Network object"
+        assert link_id.startswith(
+            "HEYO-"
+        ), "The ID of a new link did not start with the desired prefix"
+
+
+def test_splitting_links_generates_unique_ids_for_new_links_if_given_empty_prefix():
+    n = Network("epsg:27700")
+    link_ids = {str(i) for i in range(10)}
+    for i, link_id in enumerate(link_ids):
+        n.add_link(link_id, i, i + 1, attribs={"modes": {"car", "bike"}, "length": 1})
+
+    new_links = n.split_links_on_mode("bike", link_id_prefix="")
+
+    assert isinstance(new_links, set)  # To force uniqueness
+    assert len(new_links) == 10, "The number of link IDs is incorrect"
+    assert new_links & link_ids == set(), "There is an overlap between IDs already used and new"
+    for link_id in new_links:
+        assert (
+            n.link(link_id)["id"] == link_id
+        ), "The ID declared in attributes does not match the ID in the Network object"
+
+
 def test_find_shortest_path_when_graph_has_no_extra_edge_choices():
     n = Network("epsg:27700")
     n.add_link("0", 1, 2, attribs={"modes": ["car", "bike"], "length": 1})