Enable customising orchestration

.github/workflows/draft-pdf.yml

@@ -6,15 +6,15 @@ jobs:
     name: Paper Draft
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Build draft PDF
         uses: openjournals/openjournals-draft-action@master
         with:
           journal: joss
           # This should be the path to the paper within your repo.
           paper-path: docs/paper/paper.md
       - name: Upload
-        uses: actions/upload-artifact@v1
+        uses: actions/upload-artifact@v4
         with:
           name: paper
           # This is the output path where Pandoc will write the compiled

docs/demo/examples/test_customise_orchestration_example.yaml

@@ -0,0 +1,75 @@
+orchestration:
+- Groundwater: infiltrate
+- Sewer: make_discharge
+
+nodes:
+  Sewer:
+    type_: Sewer
+    name: my_sewer
+    capacity: 0.04
+
+  Groundwater:
+    type_: Groundwater
+    name: my_groundwater
+    capacity: 100
+    area: 100
+
+  River:
+    type_: Node
+    name: my_river
+
+  Waste:
+    type_: Waste
+    name: my_outlet
+
+arcs:
+  storm_outflow:
+    type_: Arc
+    name: storm_outflow
+    in_port: my_sewer
+    out_port: my_river
+
+  baseflow:
+    type_: Arc
+    name: baseflow
+    in_port: my_groundwater
+    out_port: my_river
+
+  catchment_outflow:
+    type_: Arc
+    name: catchment_outflow
+    in_port: my_river
+    out_port: my_outlet
+
+pollutants:
+- org-phosphorus
+- phosphate
+- ammonia
+- solids
+- temperature
+- nitrate
+- nitrite
+- org-nitrogen
+additive_pollutants:
+- org-phosphorus
+- phosphate
+- ammonia
+- solids
+- nitrate
+- nitrite
+- org-nitrogen
+non_additive_pollutants:
+- temperature
+float_accuracy: 1.0e-06
+
+dates:
+- '2000-01-01'
+- '2000-01-02'
+- '2000-01-03'
+- '2000-01-04'
+- '2000-01-05'
+- '2000-01-06'
+- '2000-01-07'
+- '2000-01-08'
+- '2000-01-09'
+- '2000-01-10'

tests/test_model.py

@@ -15,6 +15,7 @@
 from wsimod.nodes.sewer import Sewer
 from wsimod.nodes.waste import Waste
 from wsimod.orchestration.model import Model, to_datetime
+import os
 
 
 class MyTestClass(TestCase):
@@ -291,6 +292,18 @@ def test_run(self):
             0.03, my_model.nodes["my_land"].get_surface("urban").storage["volume"]
         )
 
+    def test_customise_orchestration(self):
+        my_model = Model()
+        my_model.load(
+            os.path.join(os.getcwd(), "docs", "demo", "examples"),
+            config_name="test_customise_orchestration_example.yaml",
+        )
+        revised_orchestration = [
+            {"Groundwater": "infiltrate"},
+            {"Sewer": "make_discharge"},
+        ]
+        self.assertListEqual(my_model.orchestration, revised_orchestration)
+
 
 if __name__ == "__main__":
     unittest.main()

wsimod/orchestration/model.py

@@ -140,6 +140,23 @@ def __init__(self):
         self.nodes = {}
         self.nodes_type = {}
 
+        # Default orchestration
+        self.orchestration = [
+            {"FWTW": "treat_water"},
+            {"Demand": "create_demand"},
+            {"Land": "run"},
+            {"Groundwater": "infiltrate"},
+            {"Sewer": "make_discharge"},
+            {"Foul": "make_discharge"},
+            {"WWTW": "calculate_discharge"},
+            {"Groundwater": "distribute"},
+            {"River": "calculate_discharge"},
+            {"Reservoir": "make_abstractions"},
+            {"Land": "apply_irrigation"},
+            {"WWTW": "make_discharge"},
+            {"Catchment": "route"},
+        ]
+
     def get_init_args(self, cls):
         """Get the arguments of the __init__ method for a class and its superclasses."""
         init_args = []
@@ -171,6 +188,15 @@ def load(self, address, config_name="config.yml", overrides={}):
         constants.FLOAT_ACCURACY = float(data["float_accuracy"])
         self.__dict__.update(Model().__dict__)
 
+        """
+        FLAG:
+            E.G. ADDITION FOR NEW ORCHESTRATION
+        """
+
+        if "orchestration" in data.keys():
+            # Update orchestration
+            self.orchestration = data["orchestration"]
+
         nodes = data["nodes"]
 
         for name, node in nodes.items():
@@ -288,6 +314,7 @@ def save(self, address, config_name="config.yml", compress=False):
         data = {
             "nodes": nodes,
             "arcs": arcs,
+            "orchestration": self.orchestration,
             "pollutants": constants.POLLUTANTS,
             "additive_pollutants": constants.ADDITIVE_POLLUTANTS,
             "non_additive_pollutants": constants.NON_ADDITIVE_POLLUTANTS,
@@ -718,55 +745,11 @@ def enablePrint(stdout):
                 node.t = date
                 node.monthyear = date.to_period("M")
 
-            # Run FWTW
-            for node in self.nodes_type.get("FWTW", {}).values():
-                node.treat_water()
-
-            # Create demand (gets pushed to sewers)
-            for node in self.nodes_type.get("Demand", {}).values():
-                node.create_demand()
-
-            # Create runoff (impervious gets pushed to sewers, pervious to groundwater)
-            for node in self.nodes_type.get("Land", {}).values():
-                node.run()
-
-            # Infiltrate GW
-            for node in self.nodes_type.get("Groundwater", {}).values():
-                node.infiltrate()
-
-            # Discharge sewers (pushed to other sewers or WWTW)
-            for node in self.nodes_type.get("Sewer", {}).values():
-                node.make_discharge()
-
-            # Foul second so that it can discharge any misconnection
-            for node in self.nodes_type.get("Foul", {}).values():
-                node.make_discharge()
-
-            # Discharge WWTW
-            for node in self.nodes_type.get("WWTW", {}).values():
-                node.calculate_discharge()
-
-            # Discharge GW
-            for node in self.nodes_type.get("Groundwater", {}).values():
-                node.distribute()
-
-            # river
-            for node in self.nodes_type.get("River", {}).values():
-                node.calculate_discharge()
-
-            # Abstract
-            for node in self.nodes_type.get("Reservoir", {}).values():
-                node.make_abstractions()
-
-            for node in self.nodes_type.get("Land", {}).values():
-                node.apply_irrigation()
-
-            for node in self.nodes_type.get("WWTW", {}).values():
-                node.make_discharge()
-
-            # Catchment routing
-            for node in self.nodes_type.get("Catchment", {}).values():
-                node.route()
+            # Iterate over orchestration
+            for timestep_item in self.orchestration:
+                for node_type, function in timestep_item.items():
+                    for node in self.nodes_type.get(node_type, {}).values():
+                        getattr(node, function)()
 
             # river
             for node_name in self.river_discharge_order: