From 5e03873a5359a18cbe2a54fffa93fbca7e56d207 Mon Sep 17 00:00:00 2001
From: Dobson <bdobson@ic.ac.uk>
Date: Wed, 28 Feb 2024 13:01:18 +0000
Subject: [PATCH] Function to run a swmm model

---
 swmmanywhere/defs/basic_drainage_all_bits.inp |  2 +-
 swmmanywhere/swmmanywhere.py                  | 84 +++++++++++++++++++
 tests/test_swmmanywhere.py                    | 27 +++++-
 3 files changed, 111 insertions(+), 2 deletions(-)
 create mode 100644 swmmanywhere/swmmanywhere.py

diff --git a/swmmanywhere/defs/basic_drainage_all_bits.inp b/swmmanywhere/defs/basic_drainage_all_bits.inp
index 0d8a4f5b..80059277 100644
--- a/swmmanywhere/defs/basic_drainage_all_bits.inp
+++ b/swmmanywhere/defs/basic_drainage_all_bits.inp
@@ -48,7 +48,7 @@ DRY_ONLY         NO
 [RAINGAGES]
 ;;Name           Format    Interval SCF      Source    
 ;;-------------- --------- ------ ------ ----------
-1                INTENSITY 00:05    1.0      FILE       "january.dat"    1          MM   
+1                INTENSITY 00:05    1.0      FILE       "storm.dat"    1          MM   
 
 [SUBCATCHMENTS]
 ;;Name           Rain Gage        Outlet           Area     %Imperv  Width    %Slope   CurbLen  SnowPack        
diff --git a/swmmanywhere/swmmanywhere.py b/swmmanywhere/swmmanywhere.py
new file mode 100644
index 00000000..825d9f9c
--- /dev/null
+++ b/swmmanywhere/swmmanywhere.py
@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+"""Created on 2024-01-26.
+
+@author: Barney
+"""
+from pathlib import Path
+
+import pandas as pd
+import pyswmm
+
+
+def run(model: Path,
+        reporting_iters: int = 50,
+        duration: int = 86400,
+        storevars: list[str] = ['flooding','flow']):
+    """Run a SWMM model and store the results.
+
+    Args:
+        model (Path): The path to the SWMM model .inp file.
+        reporting_iters (int, optional): The number of iterations between
+            storing results. Defaults to 50.
+        duration (int, optional): The duration of the simulation in seconds.
+            Starts at the 'START_DATE' and 'START_TIME' defined in the 'model'
+            .inp file Defaults to 86400.
+        storevars (list[str], optional): The variables to store. Defaults to
+            ['flooding','flow'].
+
+    Returns:
+        pd.DataFrame: A DataFrame containing the results.
+    """
+    with pyswmm.Simulation(str(model)) as sim:
+        sim.start()
+
+        # Define the variables to store
+        variables = {
+            'flooding': {'class': pyswmm.Nodes, 'id': '_nodeid'},
+            'depth': {'class': pyswmm.Nodes, 'id': '_nodeid'},
+            'flow': {'class': pyswmm.Links, 'id': '_linkid'},
+            'runoff': {'class': pyswmm.Subcatchments, 'id': '_subcatchmentid'}
+        }
+
+        results_list = []
+        for var, info in variables.items():
+            if var in storevars:
+                # Rather than calling eg Nodes or Links, only call them if they
+                # are needed for storevars because they carry a significant 
+                # overhead
+                pobjs = info['class'](sim)
+                results_list += [{'object': x, 
+                                'variable': var, 
+                                'id': info['id']} for x in pobjs]
+        
+        # Iterate the model
+        results = []
+        t_ = sim.current_time
+        ind = 0
+        while ((sim.current_time - t_).total_seconds() <= duration) & \
+            (sim.current_time < sim.end_time) & (not sim._terminate_request):
+            
+            ind+=1
+
+            # Iterate the main model timestep
+            time = sim._model.swmm_step()
+            
+            # Break condition
+            if time < 0:
+                sim._terminate_request = True
+                break
+            
+            # Check whether to save results
+            if ind % reporting_iters != 1:
+                continue
+
+            # Store results in a list of dictionaries
+            for storevar in results_list:
+                results.append({'date' : sim.current_time,
+                                'value' : getattr(storevar['object'],
+                                                  storevar['variable']),
+                                'variable' : storevar['variable'],
+                                'id' : getattr(storevar['object'],
+                                               storevar['id'])})
+            
+            
+    return pd.DataFrame(results)
\ No newline at end of file
diff --git a/tests/test_swmmanywhere.py b/tests/test_swmmanywhere.py
index a3d619be..77713d51 100644
--- a/tests/test_swmmanywhere.py
+++ b/tests/test_swmmanywhere.py
@@ -1,7 +1,32 @@
 """Tests for the main module."""
-from swmmanywhere import __version__
+from pathlib import Path
+
+from swmmanywhere import __version__, swmmanywhere
 
 
 def test_version():
     """Check that the version is acceptable."""
     assert __version__ == "0.1.0"
+
+
+def test_run():
+    """Test the run function."""
+    model = Path(__file__).parent.parent / 'swmmanywhere' / 'defs' /\
+          'basic_drainage_all_bits.inp'
+    storevars = ['flooding','flow','runoff','depth']
+    results = swmmanywhere.run(model,
+                               reporting_iters = 50,
+                               storevars = storevars)
+    assert set(results.variable.unique()) == set(storevars)
+    
+    # Ensure more reporting iterations results in more results
+    results_ = swmmanywhere.run(model,
+                               reporting_iters = 25,
+                               storevars = storevars)
+    assert results_.shape[0] > results.shape[0]
+
+    # Ensure a shorter duration results in fewer results
+    results_ = swmmanywhere.run(model,
+                               duration = 10000,
+                               storevars = storevars)
+    assert results_.shape[0] < results.shape[0]
\ No newline at end of file