Add calculate_contributing_area to graph_utilities

- Introduce Addresses into parameters
- Test

swmmanywhere/graph_utilities.py

@@ -4,13 +4,16 @@
 @author: Barney
 """
 import json
+import tempfile
 from pathlib import Path
 from typing import Callable
 
+import geopandas as gpd
 import networkx as nx
 import osmnx as ox
 from shapely import geometry as sgeom
 
+from swmmanywhere import geospatial_utilities as go
 from swmmanywhere import parameters
 
 
@@ -289,3 +292,67 @@ def create_new_edge_data(line, data, id_):
         graph.add_edge(edge[0], edge[1], **edge[2])
 
     return graph
+
+@register_graphfcn
+def calculate_contributing_area(G: nx.Graph, 
+                         subcatchment_derivation: parameters.SubcatchmentDerivation,
+                         addresses: parameters.Addresses,
+                         **kwargs):
+    """Calculate the contributing area for each edge.
+    
+    This function calculates the contributing area for each edge. The
+    contributing area is the area of the subcatchment that drains to the
+    edge. The contributing area is calculated from the elevation data.
+
+    Also writes the file 'subcatchments.geojson' to addresses.subcatchments.
+
+    Requires a graph with edges that have:
+        - 'geometry' (shapely LineString)
+        - 'id' (str)
+        - 'width' (float)
+
+    Adds the attributes:
+        - 'contributing_area' (float)
+
+    Args:
+        G (nx.Graph): A graph
+        subcatchment_derivation (parameters.SubcatchmentDerivation): A
+            SubcatchmentDerivation parameter object
+        addresses (parameters.Addresses): An Addresses parameter object
+        **kwargs: Additional keyword arguments are ignored.
+
+    Returns:
+        G (nx.Graph): A graph
+    """
+    G = G.copy()
+
+    # Carve
+    # TODO I guess we don't need to keep this 'carved' file..
+    # maybe could add verbose/debug option to keep it
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_fid = Path(temp_dir) / "carved.tif"
+        go.burn_shape_in_raster([d['geometry'] for u,v,d in G.edges(data=True)],
+                subcatchment_derivation.carve_depth,
+                addresses.elevation,
+                temp_fid)
+
+        # Derive
+        subs_gdf = go.derive_subcatchments(G,temp_fid)
+
+    # RC
+    buildings = gpd.read_file(addresses.building)
+    subs_rc = go.derive_rc(subs_gdf, G, buildings)
+
+    # Write subs
+    # TODO - could just attach subs to nodes where each node has a list of subs
+    subs_rc.to_file(addresses.subcatchments, driver='GeoJSON')
+
+    # Assign contributing area
+    imperv_lookup = subs_rc.set_index('id').impervious_area.to_dict()
+    for u,v,d in G.edges(data=True):
+        if u in imperv_lookup.keys():
+            d['contributing_area'] = imperv_lookup[u]
+        else:
+            d['contributing_area'] = 0.0
+    return G
+

swmmanywhere/parameters.py

@@ -4,6 +4,8 @@
 @author: Barney
 """
 
+from pathlib import Path
+
 from pydantic import BaseModel, Field
 
 
@@ -26,3 +28,48 @@ class SubcatchmentDerivation(BaseModel):
             le = 100.0,
             unit = "m", 
             description = "Distance to split streets into segments.")
+
+class Addresses:
+    """Parameters for address lookup.
+
+    TODO: this doesn't validate addresses to allow for un-initialised data
+    (e.g., subcatchments are created by a graph and so cannot be validated).
+    """
+
+    def __init__(self, 
+                 base_dir: Path, 
+                 project_name: str, 
+                 bbox_number: int, 
+                 model_number: str, 
+                 extension: str='json'):
+        """Initialise the class."""
+        self.base_dir = base_dir
+        self.project_name = project_name
+        self.bbox_number = bbox_number
+        self.model_number = model_number
+        self.extension = extension
+
+    def _generate_path(self, *subdirs):
+        return self.base_dir.joinpath(*subdirs)
+
+    def _generate_property(self, folder_name, location):
+        return property(lambda self: self._generate_path(self.project_name, 
+                                                         location, 
+                                                         folder_name))
+
+    def _generate_properties(self):
+        self.project = self._generate_path(self.project_name)
+        self.national = self._generate_property('national', 
+                                                'project')
+        self.bbox = self._generate_property(f'bbox_{self.bbox_number}', 
+                                            'project')
+        self.model = self._generate_property(f'model_{self.model_number}', 
+                                             'bbox')
+        self.subcatchments = self._generate_property('subcatchments', 
+                                                     'model')
+        self.download = self._generate_property('download', 
+                                                'bbox')
+        self.elevation = self._generate_property('elevation.tif', 
+                                                 'download')
+        self.building = self._generate_property(f'building.{self.extension}', 
+                                                'download')

tests/test_graph_utilities.py

@@ -3,39 +3,41 @@
 
 @author: Barney
 """
+import tempfile
+from pathlib import Path
 
+import geopandas as gpd
+from shapely import geometry as sgeom
 
-from swmmanywhere import geospatial_utilities as go
 from swmmanywhere import graph_utilities as gu
 from swmmanywhere import parameters
-from swmmanywhere.prepare_data import download_street
 
 
-def generate_street_graph():
-    """Generate a street graph."""
+def load_street_network():
+    """Load a street network."""
     bbox = (-0.11643,51.50309,-0.11169,51.50549)
-    G = download_street(bbox)
-    return G
+    G = gu.load_graph(Path(__file__).parent / 'test_data' / 'street_graph.json')
+    return G, bbox
 
 def test_assign_id():
     """Test the assign_id function."""
-    G = generate_street_graph()
+    G, _ = load_street_network()
     G = gu.assign_id(G)
     for u, v, data in G.edges(data=True):
         assert 'id' in data.keys()
         assert isinstance(data['id'], int)
 
 def test_double_directed():
     """Test the double_directed function."""
-    G = generate_street_graph()
+    G, _ = load_street_network()
     G = gu.assign_id(G)
     G = gu.double_directed(G)
     for u, v in G.edges():
         assert (v,u) in G.edges
 
 def test_format_osmnx_lanes():
     """Test the format_osmnx_lanes function."""
-    G = generate_street_graph()
+    G, _ = load_street_network()
     params = parameters.SubcatchmentDerivation()
     G = gu.format_osmnx_lanes(G, params)
     for u, v, data in G.edges(data=True):
@@ -46,16 +48,40 @@ def test_format_osmnx_lanes():
 
 def test_split_long_edges():
     """Test the split_long_edges function."""
-    G = generate_street_graph()
+    G, _ = load_street_network()
     G = gu.assign_id(G)
-    id_ = list(G.nodes)[0]
-    G = go.reproject_graph(G, 
-                           'EPSG:4326', 
-                           go.get_utm_epsg(G.nodes[id_]['x'], 
-                                           G.nodes[id_]['y'])
-                            )
     max_length = 20
     params = parameters.SubcatchmentDerivation(max_street_length = max_length)
     G = gu.split_long_edges(G, params)
     for u, v, data in G.edges(data=True):
-        assert data['length'] <= (max_length * 2)
+        assert data['length'] <= (max_length * 2)
+
+def test_derive_subcatchments():
+    """Test the derive_subcatchments function."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+        addresses = parameters.Addresses(base_dir = temp_path, 
+                            project_name = 'test', 
+                            bbox_number = 1,
+                            extension = 'json',
+                            model_number = 1)
+        addresses.elevation = Path(__file__).parent / 'test_data' / 'elevation.tif'
+        addresses.building = temp_path / 'building.geojson'
+        addresses.subcatchments = temp_path / 'subcatchments.geojson'
+        params = parameters.SubcatchmentDerivation()
+        G, bbox = load_street_network()
+
+        # mock up buildings
+        eg_bldg = sgeom.Polygon([(700291.346,5709928.922), 
+                       (700331.206,5709927.815),
+                       (700321.610,5709896.444), 
+                       (700293.192,5709900.503),
+                       (700291.346,5709928.922)])
+        gdf = gpd.GeoDataFrame(geometry = [eg_bldg],
+                               crs = G.graph['crs'])
+        gdf.to_file(addresses.building, driver='GeoJSON')
+
+        G = gu.calculate_contributing_area(G, params, addresses)
+        for u, v, data in G.edges(data=True):
+            assert 'contributing_area' in data.keys()
+            assert isinstance(data['contributing_area'], float)