Swmmio

pyproject.toml

@@ -60,6 +60,7 @@ dependencies = [
   "rioxarray",
   "scipy",
   "shapely",
+  "swmmio",
   "tqdm",
   "xarray",
 ]

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]

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, /*]

src/swmmanywhere/parameters.py

@@ -318,7 +318,6 @@ class MetricEvaluation(BaseModel):
         unit="m",
         description="Scale of the grid for metric evaluation",
     )
-
     warmup: float = Field(
         default=0,
         ge=0,
@@ -328,3 +327,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

src/swmmanywhere/post_processing.py

@@ -15,11 +15,104 @@
 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, 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.
 
@@ -51,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

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
 
 

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

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.