From 144b3cefee3a16738b350e132600e8806a091c97 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 25 Apr 2025 15:37:37 +0100 Subject: [PATCH 01/11] initial commit --- .vscode/launch.json | 25 +++++++++ pyproject.toml | 1 + src/swmmanywhere/parameters.py | 11 ++++ src/swmmanywhere/post_processing.py | 80 +++++++++++++++++++++++++++++ tests/test_post_processing.py | 80 +++++++++++++++++++++++++++++ 5 files changed, 197 insertions(+) create mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..de59c16b --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,25 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Python: pytest", + "type": "debugpy", + "request": "launch", + "program": "${file}", + "args": [ + "-m", + "pytest", + "--maxfail=1", + "--disable-warnings" + ], + "env": { + "PYTEST_ADDOPTS": "--no-cov" + }, + "console": "integratedTerminal", + "purpose": [ + "debug-test" + ], + "justMyCode": false + } + ] +} diff --git a/pyproject.toml b/pyproject.toml index a381dafd..e9d9fbd7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ dependencies = [ "rioxarray", "scipy", "shapely", + "swmmio", "tqdm", "xarray", ] diff --git a/src/swmmanywhere/parameters.py b/src/swmmanywhere/parameters.py index 796b4c9e..be38a08a 100644 --- a/src/swmmanywhere/parameters.py +++ b/src/swmmanywhere/parameters.py @@ -305,3 +305,14 @@ class MetricEvaluation(BaseModel): unit="m", description="Scale of the grid for metric evaluation", ) + + +@register_parameter_group(name="post_processing") +class PostProcessing(BaseModel): + """Parameters for post processing. + + These parameters are applied or used during the post processing steps. They may be + direct SWMM parameters, or control other factors during post processing. + """ + + pass diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index 2e0dfd95..cd9f6d54 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -15,11 +15,91 @@ import numpy as np import pandas as pd import yaml +from swmmio import Model from swmmanywhere.filepaths import FilePaths from swmmanywhere.logging import logger +def _fill_backslash_columns(df: pd.DataFrame | None, key: str) -> pd.DataFrame | None: + if df is None: + return None + + # Load conversion mapping from YAML file + with (Path(__file__).parent / "defs" / "swmm_conversion.yml").open("r") as file: + conversion_dict = yaml.safe_load(file) + + data = {} + for sa, sio in zip( + conversion_dict[key]["columns"], conversion_dict[key]["iwcolumns"] + ): + if sio.startswith("/"): + data[sa] = sio[1:] + else: + data[sa] = df[sio] + + return pd.DataFrame(data) + + +io_registry = {} +"""Registry for input/output functions.""" + + +def register_io(func): + """Decorator to register input/output functions.""" + + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + io_registry[func.__name__] = wrapper + return wrapper + + +@register_io +def apply_nodes(m: Model, addresses: FilePaths, **kw): + """Apply edges to the model. + + Args: + m (Model): The SWMMIO model to apply edges to. + addresses (FilePaths): A dictionary of file paths. + **kw: Additional keyword arguments are ignored. + """ + nodes = gpd.read_file(addresses.model_paths.nodes) + nodes = nodes[["id", "x", "y", "chamber_floor_elevation", "surface_elevation"]] + + # Nodes + nodes["id"] = nodes["id"].astype(str) + nodes["max_depth"] = nodes.surface_elevation - nodes.chamber_floor_elevation + nodes["surcharge_depth"] = 0 + nodes["flooded_area"] = 100 # TODO arbitrary... not sure how to calc this + nodes["manhole_area"] = 0.5 + + m.inp.storage = _fill_backslash_columns(nodes, "STORAGE") + m.inp.coordinates = _fill_backslash_columns(nodes, "COORDINATES") + return m + + +def iterate_io( + io_list: list[str], + params: dict, + addresses: FilePaths, +): + """Iterate a list of input/output functions over a model.""" + # Load a starting model + m = Model(str(Path(__file__).parent / "defs" / "basic_drainage_all_bits.inp")) + + for function in io_list: + if function not in io_registry: + raise ValueError(f"""Function {function} not registered in io_registry""") + + # Call the function with the model and parameters + m = io_registry[function](m, addresses, **params) + + logger.info(f"io: {function} completed.") + + m.inp.save(str(addresses.model_paths.inp)) + + def synthetic_write(addresses: FilePaths): """Load synthetic data and write to SWMM input file. diff --git a/tests/test_post_processing.py b/tests/test_post_processing.py index 1c8825af..0da5f985 100644 --- a/tests/test_post_processing.py +++ b/tests/test_post_processing.py @@ -9,10 +9,14 @@ import geopandas as gpd import pandas as pd import pyswmm +import pytest from shapely import geometry as sgeom +from swmmio import Model from swmmanywhere import post_processing as stt from swmmanywhere.filepaths import FilePaths +from swmmanywhere.parameters import get_full_parameters +from swmmanywhere.post_processing import io_registry fid = ( Path(__file__).parent.parent @@ -23,6 +27,82 @@ ) +def validate_model(m): + """Validate a SWMMIO model.""" + with tempfile.TemporaryDirectory() as temp_dir: + tmp_path = Path(temp_dir) + m.inp.save(str(tmp_path / "model.inp")) + with pyswmm.Simulation(str(tmp_path / "model.inp")) as sim: + sim.start() + + +@pytest.fixture +def filepaths(tmp_path): + """Fixture to create a temporary FilePaths object, with nodes/edges/subs.""" + addresses = FilePaths( + base_dir=tmp_path, + bbox_bounds=[0, 1, 0, 1], + project_name="test", + extension="json", + precipitation="storm.dat", + ) + + nodes = gpd.GeoDataFrame( + { + "id": ["node1", "node2"], + "x": [0, 1], + "y": [0, 1], + "chamber_floor_elevation": [1, 1], + "surface_elevation": [2, 2], + } + ) + nodes.to_file(addresses.model_paths.nodes) + + edges = gpd.GeoDataFrame( + { + "id": ["node1-node2"], + "u": ["node1"], + "v": ["node2"], + "diameter": [1], + } + ) + edges.to_file(addresses.model_paths.edges) + + subs = gpd.GeoDataFrame( + { + "id": ["node1"], + "area": [1], + "rc": [1], + "width": [1], + "slope": [0.001], + "geometry": [sgeom.Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])], + } + ) + subs.to_file(addresses.model_paths.subcatchments) + + return addresses + + +def test_apply_nodes(filepaths): + """Test the apply_nodes function.""" + m = Model(str(fid)) + m = io_registry["apply_nodes"](m, filepaths) + assert list(m.inp.storage["Name"]) == ["node1", "node2"] + assert list(m.inp.storage["InvertElev"]) == [1, 1] + assert list(m.inp.coordinates["X"]) == [0, 1] + assert list(m.inp.coordinates["Y"]) == [0, 1] + + +def test_iterate_io(filepaths): + """Test the iterate_io function.""" + stt.iterate_io( + ["apply_nodes"], + get_full_parameters(), + filepaths, + ) + assert filepaths.model_paths.inp.exists() + + def test_overwrite_section(): """Test the overwrite_section function. From 589a1d438c3ddda60e4bbc99acce8f4db8e70c5b Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 25 Apr 2025 15:38:30 +0100 Subject: [PATCH 02/11] remove vscode launch --- .vscode/launch.json | 25 ------------------------- 1 file changed, 25 deletions(-) delete mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index de59c16b..00000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "version": "0.2.0", - "configurations": [ - { - "name": "Python: pytest", - "type": "debugpy", - "request": "launch", - "program": "${file}", - "args": [ - "-m", - "pytest", - "--maxfail=1", - "--disable-warnings" - ], - "env": { - "PYTEST_ADDOPTS": "--no-cov" - }, - "console": "integratedTerminal", - "purpose": [ - "debug-test" - ], - "justMyCode": false - } - ] -} From cebcad7727472d58027d8265a1686edd8d0331fc Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 25 Apr 2025 15:53:25 +0100 Subject: [PATCH 03/11] register from swmmanywhere --- src/swmmanywhere/defs/schema.yml | 1 + src/swmmanywhere/post_processing.py | 2 ++ src/swmmanywhere/swmmanywhere.py | 13 +++++++++++++ tests/test_data/custom_io.py | 9 +++++++++ tests/test_post_processing.py | 8 ++++++++ 5 files changed, 33 insertions(+) create mode 100644 tests/test_data/custom_io.py diff --git a/src/swmmanywhere/defs/schema.yml b/src/swmmanywhere/defs/schema.yml index f62f60a8..dc452a30 100644 --- a/src/swmmanywhere/defs/schema.yml +++ b/src/swmmanywhere/defs/schema.yml @@ -34,4 +34,5 @@ properties: custom_metric_modules: {type: array, items: {type: string}} custom_graphfcn_modules: {type: array, items: {type: string}} custom_parameters_modules: {type: array, items: {type: string}} + custom_io_modules: {type: array, items: {type: string}} required: [base_dir, project, bbox] \ No newline at end of file diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index cd9f6d54..fd024011 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -51,6 +51,8 @@ def register_io(func): def wrapper(*args, **kwargs): return func(*args, **kwargs) + if func.__name__ in io_registry: + logger.warning(f"{func.__name__} already in io register, overwriting.") io_registry[func.__name__] = wrapper return wrapper diff --git a/src/swmmanywhere/swmmanywhere.py b/src/swmmanywhere/swmmanywhere.py index fde23713..3672e6ce 100644 --- a/src/swmmanywhere/swmmanywhere.py +++ b/src/swmmanywhere/swmmanywhere.py @@ -391,6 +391,16 @@ def register_custom_parameters(config: dict): return config +def register_custom_io(config: dict): + """Register custom IO modules. + + Args: + config (dict): The configuration. + """ + import_modules(config.get("custom_io_modules", [])) + return config + + def save_config(config: dict, config_path: Path): """Save the configuration to a file. @@ -461,6 +471,9 @@ def load_config( # Register custom parameters config = register_custom_parameters(config) + # Register custom IO + config = register_custom_io(config) + return config diff --git a/tests/test_data/custom_io.py b/tests/test_data/custom_io.py new file mode 100644 index 00000000..4384404c --- /dev/null +++ b/tests/test_data/custom_io.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +from swmmanywhere.post_processing import register_io + + +@register_io +def new_io(m, **kw): + """New io function.""" + return m diff --git a/tests/test_post_processing.py b/tests/test_post_processing.py index 0da5f985..3053569e 100644 --- a/tests/test_post_processing.py +++ b/tests/test_post_processing.py @@ -17,6 +17,7 @@ from swmmanywhere.filepaths import FilePaths from swmmanywhere.parameters import get_full_parameters from swmmanywhere.post_processing import io_registry +from swmmanywhere.swmmanywhere import import_module fid = ( Path(__file__).parent.parent @@ -103,6 +104,13 @@ def test_iterate_io(filepaths): assert filepaths.model_paths.inp.exists() +def test_custom_io(): + """Test register_parameter_group.""" + import_module(Path(__file__).parent / "test_data" / "custom_io.py") + + assert "new_io" in io_registry.keys() + + def test_overwrite_section(): """Test the overwrite_section function. From d001de5233d1dc9688f9afa1d4f16344ad1d4007 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 25 Apr 2025 15:59:51 +0100 Subject: [PATCH 04/11] doc --- src/swmmanywhere/post_processing.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index fd024011..28409158 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -21,10 +21,20 @@ from swmmanywhere.logging import logger -def _fill_backslash_columns(df: pd.DataFrame | None, key: str) -> pd.DataFrame | None: - if df is None: - return None +def _fill_backslash_columns(df: pd.DataFrame, key: str) -> pd.DataFrame: + """Format the data into the swmmio columns. + Use the schema set out in defs/swmm_conversion.yml. Not all columns in `df` need to + be used, but all non-backslash columns in the `iwcolumns` list in + swmm_conversion.yml must be present. + + Args: + df (pd.DataFrame): DataFrame to be formatted. + key (str): Key to look up in swmm_conversion.yml. + + Returns: + pd.DataFrame: Formatted DataFrame + """ # Load conversion mapping from YAML file with (Path(__file__).parent / "defs" / "swmm_conversion.yml").open("r") as file: conversion_dict = yaml.safe_load(file) From 09cebdb3d71118c1c87ea672da9388fbfbb5bb02 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 25 Apr 2025 16:00:40 +0100 Subject: [PATCH 05/11] naming --- src/swmmanywhere/post_processing.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index 28409158..6d4ffa41 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -40,13 +40,13 @@ def _fill_backslash_columns(df: pd.DataFrame, key: str) -> pd.DataFrame: conversion_dict = yaml.safe_load(file) data = {} - for sa, sio in zip( + for swmmio_key, swmmanywhere_key in zip( conversion_dict[key]["columns"], conversion_dict[key]["iwcolumns"] ): - if sio.startswith("/"): - data[sa] = sio[1:] + if swmmanywhere_key.startswith("/"): + data[swmmio_key] = swmmanywhere_key[1:] else: - data[sa] = df[sio] + data[swmmio_key] = df[swmmanywhere_key] return pd.DataFrame(data) From 06a572665b6aa1636d9b5be89f0d5369939afddc Mon Sep 17 00:00:00 2001 From: barneydobson Date: Mon, 12 May 2025 11:44:11 +0100 Subject: [PATCH 06/11] Update src/swmmanywhere/post_processing.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Diego Alonso Álvarez <6095790+dalonsoa@users.noreply.github.com> --- src/swmmanywhere/post_processing.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index 6d4ffa41..3d25536b 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -100,9 +100,11 @@ def iterate_io( # Load a starting model m = Model(str(Path(__file__).parent / "defs" / "basic_drainage_all_bits.inp")) + not_found = [f for f in io_list if f not in io_registry] + if not_found: + raise ValueError(f"""Functions {not_found} not registered in io_registry""") + for function in io_list: - if function not in io_registry: - raise ValueError(f"""Function {function} not registered in io_registry""") # Call the function with the model and parameters m = io_registry[function](m, addresses, **params) From de49a685aae241b9f0ac46dca64b0ffafdf4972c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 May 2025 10:45:50 +0000 Subject: [PATCH 07/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/swmmanywhere/post_processing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index 3d25536b..f1ce6394 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -105,7 +105,6 @@ def iterate_io( raise ValueError(f"""Functions {not_found} not registered in io_registry""") for function in io_list: - # Call the function with the model and parameters m = io_registry[function](m, addresses, **params) From 0578f1999658cb51f1e97d822cb781860c28bdf7 Mon Sep 17 00:00:00 2001 From: Dobson Date: Mon, 12 May 2025 11:53:13 +0100 Subject: [PATCH 08/11] update terminology to model --- src/swmmanywhere/post_processing.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index f1ce6394..db549d8f 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -68,11 +68,11 @@ def wrapper(*args, **kwargs): @register_io -def apply_nodes(m: Model, addresses: FilePaths, **kw): +def apply_nodes(model: Model, addresses: FilePaths, **kw): """Apply edges to the model. Args: - m (Model): The SWMMIO model to apply edges to. + model (Model): The SWMMIO model to apply edges to. addresses (FilePaths): A dictionary of file paths. **kw: Additional keyword arguments are ignored. """ @@ -86,9 +86,9 @@ def apply_nodes(m: Model, addresses: FilePaths, **kw): nodes["flooded_area"] = 100 # TODO arbitrary... not sure how to calc this nodes["manhole_area"] = 0.5 - m.inp.storage = _fill_backslash_columns(nodes, "STORAGE") - m.inp.coordinates = _fill_backslash_columns(nodes, "COORDINATES") - return m + model.inp.storage = _fill_backslash_columns(nodes, "STORAGE") + model.inp.coordinates = _fill_backslash_columns(nodes, "COORDINATES") + return model def iterate_io( @@ -98,7 +98,7 @@ def iterate_io( ): """Iterate a list of input/output functions over a model.""" # Load a starting model - m = Model(str(Path(__file__).parent / "defs" / "basic_drainage_all_bits.inp")) + model = Model(str(Path(__file__).parent / "defs" / "basic_drainage_all_bits.inp")) not_found = [f for f in io_list if f not in io_registry] if not_found: @@ -106,11 +106,11 @@ def iterate_io( for function in io_list: # Call the function with the model and parameters - m = io_registry[function](m, addresses, **params) + model = io_registry[function](model, addresses, **params) logger.info(f"io: {function} completed.") - m.inp.save(str(addresses.model_paths.inp)) + model.inp.save(str(addresses.model_paths.inp)) def synthetic_write(addresses: FilePaths): From 2353a23075e5a55793bdb3b5772b7b1c498f381e Mon Sep 17 00:00:00 2001 From: Dobson Date: Mon, 12 May 2025 11:54:38 +0100 Subject: [PATCH 09/11] move surcharge and ponded are ato defaults --- src/swmmanywhere/defs/swmm_conversion.yml | 2 +- src/swmmanywhere/post_processing.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/swmmanywhere/defs/swmm_conversion.yml b/src/swmmanywhere/defs/swmm_conversion.yml index 492f0181..3cb9e4e2 100644 --- a/src/swmmanywhere/defs/swmm_conversion.yml +++ b/src/swmmanywhere/defs/swmm_conversion.yml @@ -11,7 +11,7 @@ INFILTRATION: iwcolumns: [subcatchment, /0, /0, /0, /0, /0] JUNCTIONS: columns: [Name, InvertElev, MaxDepth, InitDepth, SurchargeDepth, PondedArea] - iwcolumns: [id, chamber_floor_elevation, max_depth, /0, surcharge_depth, flooded_area] + iwcolumns: [id, chamber_floor_elevation, max_depth, /0, /0, /100] OUTFALLS: columns: [Name, InvertElev, OutfallType, StageOrTimeseries, TideGate, RouteTo] iwcolumns: [id, chamber_floor_elevation, /FREE, / , /NO, /*] diff --git a/src/swmmanywhere/post_processing.py b/src/swmmanywhere/post_processing.py index db549d8f..f6e96ef0 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -144,8 +144,6 @@ def synthetic_write(addresses: FilePaths): # Nodes nodes["id"] = nodes["id"].astype(str) nodes["max_depth"] = nodes.surface_elevation - nodes.chamber_floor_elevation - nodes["surcharge_depth"] = 0 - nodes["flooded_area"] = 100 # TODO arbitrary... not sure how to calc this nodes["manhole_area"] = 0.5 # Subs From a6be8e9c2ed629749005da94b1367712296d723b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 20 May 2025 08:26:58 +0000 Subject: [PATCH 10/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/swmmanywhere/parameters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swmmanywhere/parameters.py b/src/swmmanywhere/parameters.py index 2dfb05a4..45d05599 100644 --- a/src/swmmanywhere/parameters.py +++ b/src/swmmanywhere/parameters.py @@ -328,7 +328,7 @@ class MetricEvaluation(BaseModel): calculations.""", ) - + @register_parameter_group(name="post_processing") class PostProcessing(BaseModel): """Parameters for post processing. From b2aa145fe1f59092dcde43cd02dfb7c958bbfff5 Mon Sep 17 00:00:00 2001 From: barneydobson Date: Fri, 8 May 2026 09:08:16 +0100 Subject: [PATCH 11/11] Apply suggestion from @barneydobson --- src/swmmanywhere/swmmanywhere.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/swmmanywhere/swmmanywhere.py b/src/swmmanywhere/swmmanywhere.py index cc568cee..d2b3c20b 100644 --- a/src/swmmanywhere/swmmanywhere.py +++ b/src/swmmanywhere/swmmanywhere.py @@ -468,9 +468,6 @@ def load_config( # Check and register custom graphfcns config = check_and_register_custom_graphfcns(config) - # Register custom parameters - config = register_custom_parameters(config) - # Register custom IO config = register_custom_io(config)