Convert set_slope to comprehension

ImperialCollegeLondon · Feb 6, 2024 · 83ec08c · 83ec08c
1 parent 85d09a7
commit 83ec08c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/swmmanywhere/graph_utilities.py b/swmmanywhere/graph_utilities.py
@@ -498,9 +498,13 @@ def __call__(self, G: nx.Graph,
             G (nx.Graph): A graph
         """
         G = G.copy()
-        for u,v,d in G.edges(data=True):
-            slope = (G.nodes[u]['elevation'] - G.nodes[v]['elevation']) / d['length']
-            d['surface_slope'] = slope
+        # Compute the slope for each edge
+        slope_dict = {(u, v, k): (G.nodes[u]['elevation'] - G.nodes[v]['elevation']) 
+                        / d['length'] for u, v, k, d in G.edges(data=True,
+                                                                keys=True)}
+
+        # Set the 'surface_slope' attribute for all edges
+        nx.set_edge_attributes(G, slope_dict, 'surface_slope')
         return G
 
 @register_graphfcn