Define and test apply_overrides for WTW

tests/test_wtw.py

@@ -40,8 +40,66 @@ 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)
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/wtw.py

@@ -6,7 +6,7 @@
 """
 from wsimod.core import constants
 from wsimod.nodes.nodes import Node, Tank
-
+from typing import Any, Dict
 
 class WTW(Node):
     """"""
@@ -71,18 +71,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 +94,53 @@ 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]):
+        """Apply overrides to the process parameters and liquor multipliers.
+
+        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)
+        if len(overrides) > 0:
+            print(f"No override behaviour defined for: {overrides.keys()}")
 
     def get_excess_throughput(self):
         """How much excess treatment capacity is there.