diff --git a/pyproject.toml b/pyproject.toml index 5c8096e7..21668ae5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ "rioxarray", "scipy", "shapely", + "swmmio", "tqdm", "xarray", ] 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/defs/swmm_conversion.yml b/src/swmmanywhere/defs/swmm_conversion.yml index a9e95a6f..03260ca8 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/parameters.py b/src/swmmanywhere/parameters.py index 75d109f2..baf08dd4 100644 --- a/src/swmmanywhere/parameters.py +++ b/src/swmmanywhere/parameters.py @@ -322,7 +322,6 @@ class MetricEvaluation(BaseModel): json_schema_extra={"unit": "m"}, description="Scale of the grid for metric evaluation", ) - warmup: float = Field( default=0, ge=0, @@ -332,3 +331,14 @@ class MetricEvaluation(BaseModel): is used to exclude the initial part of the simulation from the metric calculations.""", ) + + +@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 71c652b6..c5c78f40 100644 --- a/src/swmmanywhere/post_processing.py +++ b/src/swmmanywhere/post_processing.py @@ -15,12 +15,105 @@ import numpy as np import pandas as pd import yaml +from swmmio import Model from swmmanywhere.filepaths import FilePaths from swmmanywhere.logging import logger from swmmanywhere.utilities import read_df +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) + + data = {} + for swmmio_key, swmmanywhere_key in zip( + conversion_dict[key]["columns"], conversion_dict[key]["iwcolumns"] + ): + if swmmanywhere_key.startswith("/"): + data[swmmio_key] = swmmanywhere_key[1:] + else: + data[swmmio_key] = df[swmmanywhere_key] + + 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) + + if func.__name__ in io_registry: + logger.warning(f"{func.__name__} already in io register, overwriting.") + io_registry[func.__name__] = wrapper + return wrapper + + +@register_io +def apply_nodes(model: Model, addresses: FilePaths, **kw): + """Apply edges to the model. + + Args: + model (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 + + model.inp.storage = _fill_backslash_columns(nodes, "STORAGE") + model.inp.coordinates = _fill_backslash_columns(nodes, "COORDINATES") + return model + + +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 + 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: + raise ValueError(f"""Functions {not_found} not registered in io_registry""") + + for function in io_list: + # Call the function with the model and parameters + model = io_registry[function](model, addresses, **params) + + logger.info(f"io: {function} completed.") + + model.inp.save(str(addresses.model_paths.inp)) + + def synthetic_write(addresses: FilePaths): """Load synthetic data and write to SWMM input file. @@ -48,8 +141,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 diff --git a/src/swmmanywhere/swmmanywhere.py b/src/swmmanywhere/swmmanywhere.py index 2084e5c4..d2b3c20b 100644 --- a/src/swmmanywhere/swmmanywhere.py +++ b/src/swmmanywhere/swmmanywhere.py @@ -388,6 +388,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. @@ -458,6 +468,9 @@ def load_config( # Check and register custom graphfcns config = check_and_register_custom_graphfcns(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 f8e0dfcd..39370c2d 100644 --- a/tests/test_post_processing.py +++ b/tests/test_post_processing.py @@ -9,10 +9,15 @@ 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 +from swmmanywhere.swmmanywhere import import_module fid = ( Path(__file__).parent.parent @@ -23,6 +28,89 @@ ) +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_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.