diff --git a/swmmanywhere/graph_utilities.py b/swmmanywhere/graph_utilities.py index af0edc6d..f1601d23 100644 --- 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