From 9bfa0ab67d6b852c27f955d9e03ed5391621a5d8 Mon Sep 17 00:00:00 2001 From: Dobson Date: Mon, 22 Jan 2024 14:59:38 +0000 Subject: [PATCH] Add dict formatter and test which runs a sim --- swmmanywhere/swmm_text_tools.py | 183 +++++++++++++++++++++++++++++++- tests/test_swmm_text_tools.py | 89 +++++++++++++++- 2 files changed, 269 insertions(+), 3 deletions(-) diff --git a/swmmanywhere/swmm_text_tools.py b/swmmanywhere/swmm_text_tools.py index 15f01a7a..47f67311 100644 --- a/swmmanywhere/swmm_text_tools.py +++ b/swmmanywhere/swmm_text_tools.py @@ -3,9 +3,12 @@ @author: Barnaby Dobson """ +import os import re import numpy as np +import pandas as pd +import yaml def overwrite_section(data: np.ndarray, @@ -142,7 +145,183 @@ def data_dict_to_inp(data_dict: dict[str, np.ndarray], print(key) start_section = '[{0}]'.format(key) - overwrite_section(data, new_input_file, start_section) + overwrite_section(data, start_section, new_input_file) # Set the flow routing - change_flow_routing(routing, new_input_file) \ No newline at end of file + change_flow_routing(routing, new_input_file) + +def explode_polygon(row): + """Explode a polygon into a DataFrame of coordinates. + + Args: + row (pd.Series): A row of a GeoDataFrame containing a polygon. + + Example: + >>> import geopandas as gpd + >>> from shapely.geometry import Polygon + >>> df = pd.Series({'subcatchment' : '1', + ... 'geometry' : Polygon([(0,0), (1,0), + ... (1,1), (0,1)])}) + >>> explode_polygon(df) + x y subcatchment + 0 0 0 1 + 1 1 0 1 + 2 1 1 1 + 3 0 1 1 + 4 0 0 1 + """ + # Get the vertices of the polygon + vertices = list(row['geometry'].exterior.coords) + + # Create a new DataFrame for this row + df = pd.DataFrame(columns = ['x','y'], + data =vertices) + df['subcatchment'] = row['subcatchment'] + return df + +def format_to_swmm_dict(nodes, + outfalls, + conduits, + subs, + event, + symbol): + """Format data to a dictionary of data arrays with columns matching SWMM. + + These data are the parameters of all assets that are written to the SWMM + input file. More parameters are available to edit (see + defs/swmm_conversion.yml). + + Args: + nodes (pd.DataFrame): GeoDataFrame of nodes. With at least columns: + 'id', 'x', 'y', 'max_depth', 'chamber_floor_elevation', + 'manhole_area'. + outfalls (pd.DataFrame): GeoDataFrame of outfalls. With at least + columns: 'id', 'chamber_floor_elevation'. + conduits (pd.DataFrame): GeoDataFrame of conduits. With at least + columns: 'id', 'u', 'v', 'length', 'roughness', 'shape_swmm', + 'diameter'. + subs (gpd.GeoDataFrame): GeoDataFrame of subcatchments. With at least + columns: 'subcatchment', 'rain_gage', 'id', 'area', 'rc', 'width', + 'geometry', + event (dict): Dict describing storm event. With at least + keys: 'name', 'unit', 'interval', 'fid'. + symbol (dict): Dict with coordinates of rain gage. With at least keys: + 'x', 'y', 'name'. + + Example: + >>> import geopandas as gpd + >>> from shapely.geometry import Point + >>> nodes = gpd.GeoDataFrame({'id' : ['node1', 'node2'], + ... 'x' : [0, 1], + ... 'y' : [0, 1], + ... 'max_depth' : [1, 1], + ... 'chamber_floor_elevation' : [1, 1], + ... 'manhole_area' : [1,1] + ... }) + >>> outfalls = gpd.GeoDataFrame({'id' : ['outfall3'], + ... 'chamber_floor_elevation' : [1], + ... 'x' : [0], + ... 'y' : [0]}) + >>> conduits = gpd.GeoDataFrame({'id' : ['link1','link2'], + ... 'u' : ['node1','node2'], + ... 'v' : ['node2','outfall3'], + ... 'length' : [1,1], + ... 'roughness' : [1,1], + ... 'shape_swmm' : ['CIRCULAR','CIRCULAR'], + ... 'diameter' : [1,1], + ... 'capacity' : [0.1,0.1] + ... }) + >>> subs = gpd.GeoDataFrame({'subcatchment' : ['sub1'], + ... 'rain_gage' : ['1'], + ... 'id' : ['node1'], + ... 'area' : [1], + ... 'rc' : [1], + ... 'width' : [1], + ... 'slope' : [0.001], + ... 'geometry' : [sgeom.Polygon([(0,0), (1,0), + ... (1,1), (0,1)])]}) + >>> rain_fid = os.path.join(os.path.dirname(os.path.abspath(__file__)), + ... '..', + ... 'swmmanywhere', + ... 'defs', + ... 'storm.dat') + >>> event = {'name' : '1', + ... 'unit' : 'mm', + ... 'interval' : 1, + ... 'fid' : rain_fid} + >>> symbol = {'x' : 0, + ... 'y' : 0, + ... 'name' : 'name'} + >>> data_dict = stt.format_to_swmm_dict(nodes, + ... outfalls, + ... conduits, + ... subs, + ... event, + ... symbol) + """ + # Get the directory of the current module + current_dir = os.path.dirname(os.path.abspath(__file__)) + + # TODO use 'load_yaml_from_defs' + # Create the path to iso_converter.yml + iso_path = os.path.join(current_dir, + "defs", + "swmm_conversion.yml") + + + # Load conversion mapping from YAML file + with open(iso_path, "r") as file: + conversion_dict = yaml.safe_load(file) + + ## Create nodes, coordinates and map dimensions + dims = {'x1' : nodes.x.min(), + 'y1' : nodes.y.min(), + 'x2' : nodes.x.max(), + 'y2' : nodes.y.max(), + } + dims = {x : str(y) + ' ' for x,y in dims.items()} + + map_dimensions = pd.Series(dims).reset_index().set_index('index').T + polygons = subs[['subcatchment','geometry']].copy() + + # Format dicts to DataFrames + event = pd.Series(event).reset_index().set_index('index').T + symbol = pd.Series(symbol).reset_index().set_index('index').T + + # Apply the function to each row + polygons = pd.concat(polygons.apply(explode_polygon, axis=1).tolist()) + + ## Specify sections + shps = {'SUBCATCHMENTS' : subs, + 'CONDUITS' : conduits, + 'OUTFALLS' : outfalls, + 'STORAGE' : nodes, + 'XSECTIONS' : conduits, + 'SUBAREAS' : subs, + 'INFILTRATION' : subs, + 'COORDINATES' : pd.concat([nodes, outfalls],axis=0), + 'MAP' : map_dimensions, + 'Polygons' : polygons, + 'PUMPS' : None, + 'ORIFICES' : None, + 'WEIRS' : None, + 'OUTLETS' : None, + 'RAINGAGES' : event, + 'SYMBOLS' : symbol + } + + data_dict = {} + for key, shp in shps.items(): + if shp is not None: + shp = shp.copy() + shp = shp.fillna(0) + shp = shp[[s for s in shp.columns if not s.startswith('/')]] + columns = conversion_dict[key]['iwcolumns'] + for col in columns: + if col[0] == '/': + shp[col] = col.replace('/','') + data = shp[columns].values + data_dict[key] = data + else: + data_dict[key] = None + return data_dict \ No newline at end of file diff --git a/tests/test_swmm_text_tools.py b/tests/test_swmm_text_tools.py index f7725fe7..30111341 100644 --- a/tests/test_swmm_text_tools.py +++ b/tests/test_swmm_text_tools.py @@ -1,6 +1,11 @@ import os import shutil +import geopandas as gpd +import pandas as pd +import pyswmm +from shapely import geometry as sgeom + from swmmanywhere import swmm_text_tools as stt @@ -75,4 +80,86 @@ def test_data_input_dict_to_inp(): content = file.read() assert 'subca_3' in content finally: - os.remove(temp_fid) \ No newline at end of file + os.remove(temp_fid) + +def test_explode_polygon(): + """Test the explode_polygon function.""" + df = pd.Series({'subcatchment' : '1', + 'geometry' : sgeom.Polygon([(0,0), (1,0), + (1,1), (0,1)])}) + result = stt.explode_polygon(df) + assert result.shape[0] == 5 + assert result.loc[3,'y'] == 1 + +def test_format_format_to_swmm_dict(): + """Test the format_format_to_swmm_dict function. + + Writes a formatted dict to a model and checks that it runs without + error. + """ + nodes = gpd.GeoDataFrame({'id' : ['node1', 'node2'], + 'x' : [0, 1], + 'y' : [0, 1], + 'max_depth' : [1, 1], + 'chamber_floor_elevation' : [1, 1], + 'manhole_area' : [1,1] + }) + outfalls = gpd.GeoDataFrame({'id' : ['outfall3'], + 'chamber_floor_elevation' : [1], + 'x' : [0], + 'y' : [0]}) + conduits = gpd.GeoDataFrame({'id' : ['link1','link2'], + 'u' : ['node1','node2'], + 'v' : ['node2','outfall3'], + 'length' : [1,1], + 'roughness' : [1,1], + 'shape_swmm' : ['CIRCULAR','CIRCULAR'], + 'diameter' : [1,1], + 'capacity' : [0.1,0.1] + }) + subs = gpd.GeoDataFrame({'subcatchment' : ['sub1'], + 'rain_gage' : ['1'], + 'id' : ['node1'], + 'area' : [1], + 'rc' : [1], + 'width' : [1], + 'slope' : [0.001], + 'geometry' : [sgeom.Polygon([(0,0), (1,0), + (1,1), (0,1)])]}) + + rain_fid = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', + 'swmmanywhere', + 'defs', + 'storm.dat') + event = {'name' : '1', + 'unit' : 'mm', + 'interval' : 1, + 'fid' : rain_fid} + symbol = {'x' : 0, + 'y' : 0, + 'name' : 'name'} + data_dict = stt.format_to_swmm_dict(nodes, + outfalls, + conduits, + subs, + event, + symbol) + fid = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', + 'swmmanywhere', + 'defs', + 'basic_drainage_all_bits.inp') + temp_fid = 'temp.inp' + shutil.copy(fid, temp_fid) + try: + stt.data_dict_to_inp(data_dict, + fid, + temp_fid) + with pyswmm.Simulation(temp_fid) as sim: + for ind, step in enumerate(sim): + pass + finally: + os.remove(temp_fid.replace('.inp','.rpt')) + os.remove(temp_fid) + \ No newline at end of file