remove need for geometry after merge-nodes

ImperialCollegeLondon · May 3, 2024 · 8bba151 · 8bba151
1 parent 2968a0f
commit 8bba151
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/swmmanywhere/graph_utilities.py b/swmmanywhere/graph_utilities.py
@@ -183,7 +183,8 @@ def iterate_graphfcns(G: nx.Graph,
         G = graphfcns[function](G, addresses = addresses, **params)
         logger.info(f"graphfcn: {function} completed.")
         if verbose:
-            save_graph(G, addresses.model / f"{function}_graph.json")
+            save_graph(graphfcns.fix_geometries(G), 
+                       addresses.model / f"{function}_graph.json")
     return G
 
 @register_graphfcn
@@ -424,7 +425,6 @@ def __call__(self,
                                   new_linestring.coords[1:]):
                 geom = shapely.LineString([start, end])
                 new_edges[(start, end, 0)] = {**d,
-                                           'geometry' : geom,
                                            'length' : geom.length
                                            }
 
@@ -494,12 +494,6 @@ def __call__(self,
         self_loops = list(nx.selfloop_edges(G))
         G.remove_edges_from(self_loops)
 
-        # Recalculate geometries
-        for u, v, data in G.edges(data=True):
-            data['geometry'] = shapely.LineString([(G.nodes[u]['x'], 
-                                                    G.nodes[u]['y']),
-                                                   (G.nodes[v]['x'], 
-                                                    G.nodes[v]['y'])])
         return G
 
 @register_graphfcn
@@ -522,10 +516,16 @@ def __call__(self, G: nx.Graph, **kwargs) -> nx.Graph:
         """
         G = G.copy()
         for u, v, data in G.edges(data=True):
-            start_point_node = (G.nodes[u]['x'], G.nodes[u]['y'])
-            start_point_edge = data['geometry'].coords[0]
+            if not data.get('geometry', None):
+                start_point_edge = (None,None)
+                end_point_edge = (None,None)
+            else:
+                start_point_edge = data['geometry'].coords[0]
+                end_point_edge = data['geometry'].coords[-1]
+
+            start_point_node = (G.nodes[u]['x'], G.nodes[u]['y'])            
             end_point_node = (G.nodes[v]['x'], G.nodes[v]['y'])
-            end_point_edge = data['geometry'].coords[-1]
+
             if (start_point_edge == end_point_node) & \
                     (end_point_edge == start_point_node):
                 data['geometry'] = data['geometry'].reverse()

diff --git a/tests/test_data/demo_config.yml b/tests/test_data/demo_config.yml
@@ -31,6 +31,7 @@ graphfcn_list:
   - identify_outlets
   - derive_topology
   - pipe_by_pipe
+  - fix_geometries
   - assign_id
 metric_list:
   - nc_deltacon0