Merge pull request #70 from ImperialCollegeLondon/wwtw_override

Wwtw override

tests/test_wtw.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 from wsimod.core import constants
-from wsimod.nodes.wtw import WTW
+from wsimod.nodes.wtw import WTW, WWTW
 
 
 class MyTestClass(TestCase):
@@ -40,8 +40,79 @@ def test_excess(self):
         self.assertEqual(2, wtw.get_excess_throughput())
 
     def test_treat(self):
-        pass
+        constants.set_simple_pollutants()
+        wtw = WTW(
+            name="",
+            treatment_throughput_capacity=10,
+        )
+        wtw.current_input= {'volume' : 8,
+                            'phosphate' : 5,
+                            'temperature' : constants.DECAY_REFERENCE_TEMPERATURE}
+        wtw.treat_current_input()
+        self.assertEqual(8 * wtw.process_parameters['volume']['constant'], 
+                         wtw.treated["volume"])
+        self.assertEqual(5 * wtw.process_parameters['phosphate']['constant'], 
+                         wtw.treated["phosphate"])
+        self.assertEqual(8 * wtw.liquor_multiplier['volume'],
+                         wtw.liquor['volume'])
+        self.assertEqual(5 * wtw.liquor_multiplier['phosphate'],
+                         wtw.liquor['phosphate'])
+        self.assertEqual(5 - 5 * wtw.process_parameters['phosphate']['constant']\
+                         - 5 * wtw.liquor_multiplier['phosphate'],
+                         wtw.solids['phosphate'])
+
+    def test_override(self):
+        wtw = WTW(
+            name="",
+            treatment_throughput_capacity=10,
+            percent_solids = 0.1,
+            liquor_multiplier = {'volume' : 0.05, 'phosphate' : 0.5},
+            process_parameters = {'phosphate' : {'constant' : 0.1,
+                                                 'exponent' : 1.001}}
+        )
+
+        wtw.apply_overrides({'percent_solids' : 0.05})
+        self.assertAlmostEqual(wtw.process_parameters['volume']['constant'],
+                        0.9)
+        self.assertEqual(wtw.percent_solids, 0.05)
+
+        wtw.apply_overrides({'percent_solids' : 0.1,
+                             'liquor_multiplier' : {'volume' : 0.1}})
+
+        self.assertEqual(wtw.process_parameters['volume']['constant'],
+                        0.8)
+        self.assertEqual(wtw.liquor_multiplier['volume'], 0.1)
+        self.assertEqual(wtw.liquor_multiplier['phosphate'], 0.5)
+
+        wtw.apply_overrides({'percent_solids' : 0.1,
+                             'liquor_multiplier' : {'volume' : 0.1,
+                                                    'phosphate' : 0.01}})
+        self.assertEqual(wtw.liquor_multiplier['phosphate'], 0.01)
+
+        wtw.apply_overrides({'process_parameters' : {'phosphate' : {'constant' : 0.01}}})
+        self.assertEqual(wtw.process_parameters['phosphate']['constant'], 0.01)
+        self.assertEqual(wtw.process_parameters['phosphate']['exponent'], 1.001)
+
+        overrides = {'process_parameters' : {'phosphate' : {'exponent' : 1.01}},
+                            'liquor_multiplier' : {'phosphate' : 0.1},
+                            'percent_solids' : 0.1,
+                            'treatment_throughput_capacity' : 20,
+                            'name' : 'new_name'}
+        wtw.apply_overrides(overrides)
+        self.assertSetEqual(set(overrides.keys()), set(['name']))
+        self.assertEqual(wtw.treatment_throughput_capacity, 20)
 
+    def test_wwtw_overrides(self):
+        wwtw = WWTW(name='')
+        vol = wwtw.process_parameters['volume']['constant']
+        wwtw.apply_overrides({'treatment_throughput_capacity' : 20,
+                              'process_parameters' : {'phosphate' : 
+                                                      {'constant' : 0.01}},
+                              'stormwater_storage_capacity': 100})
+        self.assertEqual(wwtw.treatment_throughput_capacity, 20)
+        self.assertEqual(wwtw.process_parameters['phosphate']['constant'], 0.01)
+        self.assertEqual(wwtw.process_parameters['volume']['constant'], vol)
+        self.assertEqual(wwtw.stormwater_storage_capacity, 100)
 
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/nodes.py

@@ -81,10 +81,15 @@ def __init__(self, name, data_input_dict=None):
     def apply_overrides(self, overrides: Dict[str, Any] = {}) -> None:
         """Apply overrides to the node.
 
+        The Node does not have any overwriteable parameters. So if any 
+        overrides are passed up to the node, this means that there are unused 
+        parameters from the Node subclass, which is flagged.
+        
         Args:
             overrides (dict, optional): Dictionary of overrides. Defaults to {}.
         """
-        pass
+        if len(overrides) > 0:
+            print(f"No override behaviour defined for: {overrides.keys()}")
 
     def total_in(self):
         """Sum flow and pollutant amounts entering a node via in_arcs.

wsimod/nodes/wtw.py

@@ -6,10 +6,15 @@
 """
 from wsimod.core import constants
 from wsimod.nodes.nodes import Node, Tank
-
+from typing import Any, Dict
 
 class WTW(Node):
-    """"""
+    """A generic Water Treatment Works (WTW) node.
+    
+    This class is a generic water treatment works node. It is intended to be
+    subclassed into freshwater and wastewater treatment works (FWTW and WWTW
+    respectively).
+    """
 
     def __init__(
         self,
@@ -71,18 +76,18 @@ def __init__(
                 for x in constants.ADDITIVE_POLLUTANTS
             }
         if len(liquor_multiplier) > 0:
-            self.liquor_multiplier = liquor_multiplier
+            self._liquor_multiplier = liquor_multiplier
         else:
-            self.liquor_multiplier = {x: 0.7 for x in constants.ADDITIVE_POLLUTANTS}
-            self.liquor_multiplier["volume"] = 0.03
+            self._liquor_multiplier = {x: 0.7 for x in constants.ADDITIVE_POLLUTANTS}
+            self._liquor_multiplier["volume"] = 0.03
 
-        self.percent_solids = percent_solids
+        self._percent_solids = percent_solids
 
         # Update args
         super().__init__(name)
 
         self.process_parameters["volume"] = {
-            "constant": 1 - self.percent_solids - self.liquor_multiplier["volume"]
+            "constant": self.calculate_volume()
         }
 
         # Update handlers
@@ -94,6 +99,57 @@ def __init__(
         self.treated = self.empty_vqip()
         self.liquor = self.empty_vqip()
         self.solids = self.empty_vqip()
+
+    def calculate_volume(self):
+        """Calculate the volume proportion of treated water.
+
+        Returns:
+            (float): Volume of treated water
+        """
+        return 1 - self._percent_solids - self._liquor_multiplier["volume"]
+
+    @property
+    def percent_solids(self):
+        return self._percent_solids
+
+    @percent_solids.setter
+    def percent_solids(self, value):
+        self._percent_solids = value
+        self.process_parameters["volume"]["constant"] = self.calculate_volume()
+
+    @property
+    def liquor_multiplier(self):
+        return self._liquor_multiplier
+
+    @liquor_multiplier.setter
+    def liquor_multiplier(self, value):
+        self._liquor_multiplier.update(value)
+        self.process_parameters["volume"]["constant"] = self.calculate_volume()
+
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        percent_solids, treatment_throughput_capacity, process_parameters (the
+        entire dict does not need to be redefined, only changed values need to
+        be included), liquor_multiplier (as with process_parameters).
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.percent_solids = overrides.pop("percent_solids", 
+                                            self._percent_solids)
+        self.liquor_multiplier = overrides.pop("liquor_multiplier", 
+                                               self._liquor_multiplier)
+        process_parameters = overrides.pop("process_parameters", {})
+        for key, value in process_parameters.items():
+            self.process_parameters[key].update(value)
+
+        self.treatment_throughput_capacity = overrides.pop(
+            "treatment_throughput_capacity",
+            self.treatment_throughput_capacity)
+        super().apply_overrides(overrides)
 
     def get_excess_throughput(self):
         """How much excess treatment capacity is there.
@@ -164,7 +220,7 @@ def end_timestep(self):
 
 
 class WWTW(WTW):
-    """"""
+    """Wastewater Treatment Works (WWTW) node."""
 
     def __init__(
         self,
@@ -239,6 +295,30 @@ def __init__(
             lambda: self.ds_vqip(self.liquor, self.liquor_)
         )  # Change in liquor
 
+    def apply_overrides(self, overrides=Dict[str, Any]):
+        """Apply overrides to the stormwater tank and WWTW.
+
+        Enables a user to override any parameter of the stormwater tank, and
+        then calls any overrides in WTW.
+        
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.stormwater_storage_capacity = overrides.pop(
+            "stormwater_storage_capacity", 
+            self.stormwater_storage_capacity)
+        self.stormwater_storage_area = overrides.pop(
+            "stormwater_storage_area", 
+            self.stormwater_storage_area)
+        self.stormwater_storage_elevation = overrides.pop(
+            "stormwater_storage_elevation", 
+            self.stormwater_storage_elevation)
+        self.stormwater_tank.area = self.stormwater_storage_area
+        self.stormwater_tank.capacity = self.stormwater_storage_capacity
+        self.stormwater_tank.datum = self.stormwater_storage_elevation
+        super().apply_overrides(overrides)
+
     def calculate_discharge(self):
         """Clear stormwater tank if possible, and call treat_current_input."""
         # Run WWTW model