Enable customising orchestration

docs/demo/examples/test_customise_orchestration_example.yaml

@@ -0,0 +1,138 @@
+inputs: docs/demo/data/processed
+outputs: results/quickstart_results
+
+orchestration:
+- Land: run
+- Groundwater: infiltrate
+- Sewer: make_discharge
+- Land: apply_irrigation
+
+data:
+  my_land_data:
+    filename: timeseries_data.csv
+    filter:
+      - where: site
+        is: oxford_land
+    scaling:
+      - where: variable
+        is: precipitation
+        variable: value
+        factor: "MM_TO_M"
+    format: dict
+    index: ['variable', 'date']
+    output: 'value'
+    options: parse_dates=['date']
+
+  dates_data:
+    filename: timeseries_data.csv
+    options: usecols=['date'],parse_dates=['date']
+
+dates: data:dates_data
+
+nodes:
+  Sewer:
+    type_: Sewer
+    name: my_sewer
+    capacity: 0.04
+  Land:
+    type_: Land
+    name: my_land
+    data_input_dict: data:my_land_data
+    surfaces:
+      ImperviousSurface:
+        type_: ImperviousSurface
+        surface: urban
+        area: 10
+        pollutant_load:
+          phosphate: 1.0e-07
+      PerviousSurface:
+        type_: PerviousSurface
+        surface: rural
+        area: 100
+        depth: 0.5
+        pollutant_load:
+          phosphate: 1.0e-07
+
+  Groundwater:
+    type_: Groundwater
+    name: my_groundwater
+    capacity: 100
+    area: 100
+
+  River:
+    type_: Node
+    name: my_river
+
+  Waste:
+    type_: Waste
+    name: my_outlet
+
+arcs:
+  urban_drainage:
+    type_: Arc
+    name: urban_drainage
+    in_port: my_land
+    out_port: my_sewer
+
+  percolation:
+    type_: Arc
+    name: percolation
+    in_port: my_land
+    out_port: my_groundwater
+
+  runoff:
+    type_: Arc
+    name: runoff
+    in_port: my_land
+    out_port: my_river
+
+  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,7 +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):
     def assertDictAlmostEqual(self, d1, d2, accuracy=19):
@@ -290,7 +290,15 @@ def test_run(self):
         self.assertEqual(
             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(), os.pardir, "docs", "demo", "examples"), 
+                      config_name='test_customise_orchestration_example.yaml')
+        revised_orchestration = [{'Land': 'run'}, 
+                                 {'Groundwater': 'infiltrate'}, 
+                                 {'Sewer': 'make_discharge'}, 
+                                 {'Land': 'apply_irrigation'}]
+        self.assertListEqual(my_model.orchestration, revised_orchestration)
 
 if __name__ == "__main__":
     unittest.main()

wsimod/orchestration/model.py

@@ -144,6 +144,22 @@ def __init__(self):
         # self.arcs_type = {} #not sure that this would be necessary
         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 all_subclasses(cls):
             """
@@ -194,7 +210,16 @@ def load(self, address, config_name="config.yml", overrides={}):
         constants.NON_ADDITIVE_POLLUTANTS = data["non_additive_pollutants"]
         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():
@@ -311,6 +336,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,
@@ -742,57 +768,13 @@ def enablePrint(stdout):
             for node in self.nodelist:
                 node.t = date
                 node.monthyear = date.to_period("M")
-
-            # Run FWTW
-            for node in self.nodes_type["FWTW"].values():
-                node.treat_water()
-
-            # Create demand (gets pushed to sewers)
-            for node in self.nodes_type["Demand"].values():
-                node.create_demand()
-
-            # Create runoff (impervious gets pushed to sewers, pervious to groundwater)
-            for node in self.nodes_type["Land"].values():
-                node.run()
-
-            # Infiltrate GW
-            for node in self.nodes_type["Groundwater"].values():
-                node.infiltrate()
-
-            # Discharge sewers (pushed to other sewers or WWTW)
-            for node in self.nodes_type["Sewer"].values():
-                node.make_discharge()
-
-            # Foul second so that it can discharge any misconnection
-            for node in self.nodes_type["Foul"].values():
-                node.make_discharge()
-
-            # Discharge WWTW
-            for node in self.nodes_type["WWTW"].values():
-                node.calculate_discharge()
-
-            # Discharge GW
-            for node in self.nodes_type["Groundwater"].values():
-                node.distribute()
-
-            # river
-            for node in self.nodes_type["River"].values():
-                node.calculate_discharge()
-
-            # Abstract
-            for node in self.nodes_type["Reservoir"].values():
-                node.make_abstractions()
-
-            for node in self.nodes_type["Land"].values():
-                node.apply_irrigation()
-
-            for node in self.nodes_type["WWTW"].values():
-                node.make_discharge()
-
-            # Catchment routing
-            for node in self.nodes_type["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[node_type].values():
+                        getattr(node, function)()
+
             # river
             for node_name in self.river_discharge_order:
                 self.nodes[node_name].distribute()