Selecting network subsets

src/modelskill/network.py

@@ -18,6 +18,7 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Sequence, overload, TYPE_CHECKING
+from copy import deepcopy
 
 import networkx as nx
 import pandas as pd
@@ -324,10 +325,14 @@ class Network:
     """Network built from a set of edges, with coordinate lookup and data access."""
 
     def __init__(self, edges: Sequence[NetworkEdge]):
-        self._edges: dict[str, NetworkEdge] = {e.id: e for e in edges}
-        self._graph = self._initialize_graph()
-        self._alias_map = self._initialize_alias_map()
-        self._df = self._build_dataframe()
+        graph = self._generate_graph(edges)
+        self._initialize_network_attributes(graph)
+        self._edges = self._generate_edges_dict(edges)
+
+    def _initialize_network_attributes(self, graph: nx.Graph):
+        self._alias_map = self._generate_alias_map(graph)
+        self._df = self._build_dataframe(graph)
+        self._graph = graph.copy()
 
     def __repr__(self) -> str:
         time = self._df.index
@@ -360,12 +365,13 @@ def from_res1d(cls, res: str | Path | Res1D) -> Network:
         >>> network = Network.from_res1d(Res1D("model.res1d"))
         """
         if sys.version_info >= (3, 14):
-            raise NotImplementedError(f"Current version of 'mikeio1d' requires python < 3.14 and {sys.version} is being used.")
-
+            raise NotImplementedError(
+                f"Current version of 'mikeio1d' requires python < 3.14 and {sys.version} is being used."
+            )
+
         from mikeio1d import Res1D as _Res1D
         from modelskill.model.adapters._res1d import Res1DReach
 
-
         if isinstance(res, (str, Path)):
             path = Path(res)
             if path.suffix.lower() != ".res1d":
@@ -384,11 +390,17 @@ def from_res1d(cls, res: str | Path | Res1D) -> Network:
         ]
         return cls(edges)
 
-    def _initialize_alias_map(self) -> dict[str | tuple[str, float], int]:
-        return {self.graph.nodes[id]["alias"]: id for id in self.graph.nodes()}
+    @staticmethod
+    def _generate_alias_map(g: nx.Graph) -> dict[str | tuple[str, float], int]:
+        return {g.nodes[id]["alias"]: id for id in g.nodes()}
+
+    @staticmethod
+    def _generate_edges_dict(edges: Sequence[NetworkEdge]) -> dict[str, NetworkEdge]:
+        return {e.id: e for e in edges}
 
-    def _build_dataframe(self) -> pd.DataFrame:
-        df = pd.concat({k: v["data"] for k, v in self._graph.nodes.items()}, axis=1)
+    @staticmethod
+    def _build_dataframe(g: nx.Graph) -> pd.DataFrame:
+        df = pd.concat({k: v["data"] for k, v in g.nodes.items()}, axis=1)
         df.columns = df.columns.set_names(["node", "quantity"])
         df.index.name = "time"
         return df.copy()
@@ -445,9 +457,10 @@ def quantities(self) -> list[str]:
         """
         return list(self.to_dataframe().columns.get_level_values(1).unique())
 
-    def _initialize_graph(self) -> nx.Graph:
+    @staticmethod
+    def _generate_graph(edges: Sequence[NetworkEdge]) -> nx.Graph:
         g0 = nx.Graph()
-        for edge in self._edges.values():
+        for edge in edges:
             # 1) Add start and end nodes
             for node in [edge.start, edge.end]:
                 node_key = node.id
@@ -486,6 +499,39 @@ def _initialize_graph(self) -> nx.Graph:
 
         return nx.convert_node_labels_to_integers(g0, label_attribute="alias")
 
+    def reduce_around(self, node: int, radius: int = 5, copy: bool = True) -> None | "Network":
+        """Select subset of data around a node.
+
+        Parameters
+        ----------
+        node : int
+            Id of node that represents the center of the graph subset
+        radius : int, default 5
+            Number of hops around the central node of the subset
+        copy : bool, default
+            Return a copy of the network or mutate the network in place
+        """
+
+        network_copy = self.copy()
+        graph_subset: nx.Graph = nx.ego_graph(network_copy.graph, node, radius)
+
+        subset_edges = []
+        for edge in self._edges.values():
+            new_start = self.find(node=edge.start.id)
+            new_end = self.find(node=edge.end.id)
+            if (new_start in graph_subset.nodes) and (new_end in graph_subset.nodes):
+                subset_edges.append(edge)
+
+        if copy:
+            network_copy._initialize_network_attributes(graph_subset)
+            network_copy._edges = network_copy._generate_edges_dict(subset_edges)
+            return network_copy
+        else:
+            self._initialize_network_attributes(graph_subset)
+            self._edges = self._generate_edges_dict(subset_edges)
+            return None
+
+
     @overload
     def find(
         self,
@@ -693,6 +739,17 @@ def recall(self, id: int | list[int]) -> dict[str, Any] | list[dict[str, Any]]:
             return results[0]
         else:
             return results
+
+    def copy(self) -> "Network":
+        """Create a deep copy of the Network.
+
+        Returns
+        -------
+        Network
+            Deep copy of the Network object
+        """
+        return deepcopy(self)
+
 
 
 def _make_basic_network(node_ids, time, data, quantity="WaterLevel"):

tests/test_network.py

@@ -437,8 +437,32 @@ def test_matching_workflow_multiple_nodes(self, sample_network, sample_node_data
             assert comparer.n_points > 0
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14"
+)
 def test_open_res1d():
     path_to_file = "./tests/testdata/network.res1d"
     network = Network.from_res1d(path_to_file)
-    assert network.graph.number_of_nodes() == 259
+    assert network.graph.number_of_nodes() == 259
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14"
+)
+def test_network_subset_copy():
+    path_to_file = "./tests/testdata/network.res1d"
+    network = Network.from_res1d(path_to_file)
+    small_network = network.reduce_around(node=52)
+    assert small_network.graph.number_of_nodes() == 22
+    assert 52 in small_network.graph.nodes()
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14"
+)
+def test_network_subset_inplace():
+    path_to_file = "./tests/testdata/network.res1d"
+    network = Network.from_res1d(path_to_file)
+    network.reduce_around(node=52, copy=False)
+    assert network.graph.number_of_nodes() == 22
+    assert 52 in network.graph.nodes()