Add option to add missing nodes to a specific topology

Author: arneboockmeyer

test/split_test.py

@@ -11,7 +11,7 @@
     Topology,
     Wgs84GeoNode,
 )
-from yaramo.operations import Split
+from yaramo.operations import Label, Split
 
 
 def test_simple_split():
@@ -257,16 +257,20 @@ def test_add_non_connected_elements_to_a_partition():
     node_b4 = Node(geo_node=Wgs84GeoNode(120, 10))
     edge_1 = Edge(node_1, node_3)
     edge_2 = Edge(node_2, node_3)
-    edge_3 = Edge(node_4, node_3)
+    edge_3 = Edge(node_3, node_4)
     edge_4 = Edge(node_4, node_5)
     edge_5 = Edge(node_4, node_6)
     edge_b1 = Edge(node_b1, node_b2)
     edge_b2 = Edge(node_b2, node_b3)
     edge_b3 = Edge(node_b2, node_b4)
-    topology.add_nodes([node_1, node_2, node_3, node_4, node_5, node_6, node_b1, node_b2, node_b3, node_b4])
+    topology.add_nodes(
+        [node_1, node_2, node_3, node_4, node_5, node_6, node_b1, node_b2, node_b3, node_b4]
+    )
     topology.add_edges([edge_1, edge_2, edge_3, edge_4, edge_5, edge_b1, edge_b2, edge_b3])
 
-    topology_a, topology_b = Split.split(topology, split_edges={edge_3: 5.0})
+    topology_a, topology_b = Split.split(
+        topology, split_edges={edge_3: 5.0}, add_missing_elements_to_topology=Label.A_Topology
+    )
 
     assert set([node_1, node_2, node_3, node_b1, node_b2, node_b3, node_b4]).issubset(
         topology_a.nodes.values()

yaramo/operations/__init__.py

@@ -1,2 +1,2 @@
-from .split import Split
+from .split import Label, Split
 from .union import Union

yaramo/operations/split.py

@@ -17,7 +17,11 @@ def get_opposite_label(label):
 
 class Split:
     @staticmethod
-    def split(topology: Topology, split_edges: Dict[Edge, float]) -> Tuple[Topology, Topology]:
+    def split(
+        topology: Topology,
+        split_edges: Dict[Edge, float],
+        add_missing_elements_to_topology: None | Label = None,
+    ) -> Tuple[Topology, Topology]:
         Split._validate_split_edges(split_edges)
 
         topology_a = Topology()
@@ -26,6 +30,11 @@ def split(topology: Topology, split_edges: Dict[Edge, float]) -> Tuple[Topology,
         new_end_nodes = Split._split_edges(topology, split_edges)
         node_labels, edge_labels, signal_labels = Split._label_elements(topology, new_end_nodes)
 
+        if add_missing_elements_to_topology is not None:
+            Split._assign_missing_elements_to_label(
+                topology, node_labels, edge_labels, signal_labels, add_missing_elements_to_topology
+            )
+
         # Add nodes, edges and signals to destination topologies.
         def _add_elements_to_topology(
             _element_label_matching, _topology_a_method, _topology_b_method
@@ -188,6 +197,24 @@ def _dfs(_start_node: Node, _label):
 
         return node_labels, edge_labels, signal_labels
 
+    @staticmethod
+    def _assign_missing_elements_to_label(
+        topology: Topology,
+        node_labels: Dict[Node, Label],
+        edge_labels: Dict[Edge, Label],
+        signal_labels: Dict[Signal, Label],
+        label: Label,
+    ):
+        for node in topology.nodes.values():
+            if node not in node_labels:
+                node_labels[node] = label
+        for edge in topology.edges.values():
+            if edge not in edge_labels:
+                edge_labels[edge] = label
+        for signal in topology.signals.values():
+            if signal not in signal_labels:
+                signal_labels[signal] = label
+
     @staticmethod
     def _validate_for_data_loss(
         topology: Topology,