Apply overrides in FWTW

tests/test_wtw.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 from wsimod.core import constants
-from wsimod.nodes.wtw import WTW, WWTW
+from wsimod.nodes.wtw import WTW, WWTW, FWTW
 
 
 class MyTestClass(TestCase):
@@ -128,6 +128,28 @@ def test_wwtw_overrides(self):
  self.assertEqual(wwtw.process_parameters["volume"]["constant"], vol)
  self.assertEqual(wwtw.stormwater_storage_capacity, 100)
 
+ def test_fwtw_overrides(self):
+ fwtw = FWTW(name="")
+ vol = fwtw.process_parameters["volume"]["constant"]
+ fwtw.apply_overrides(
+ {
+ "treatment_throughput_capacity": 20,
+ "process_parameters": {"phosphate": {"constant": 0.02}},
+ "service_reservoir_storage_capacity": 100,
+ "service_reservoir_storage_area": 34.7,
+ "service_reservoir_storage_elevation": 68.2,
+ }
+ )
+ self.assertEqual(fwtw.treatment_throughput_capacity, 20)
+ self.assertEqual(fwtw.process_parameters["phosphate"]["constant"], 0.02)
+ self.assertEqual(fwtw.process_parameters["volume"]["constant"], vol)
+ self.assertEqual(fwtw.service_reservoir_storage_capacity, 100)
+ self.assertEqual(fwtw.service_reservoir_tank.capacity, 100)
+ self.assertEqual(fwtw.service_reservoir_storage_area, 34.7)
+ self.assertEqual(fwtw.service_reservoir_tank.area, 34.7)
+ self.assertEqual(fwtw.service_reservoir_storage_elevation, 68.2)
+ self.assertEqual(fwtw.service_reservoir_tank.datum, 68.2)
+
 
 if __name__ == "__main__":
  unittest.main()

wsimod/nodes/wtw.py

@@ -543,6 +543,33 @@ def __init__(
  self.mass_balance_ds.append(lambda: self.service_reservoir_tank.ds())
  self.mass_balance_out.append(lambda: self.unpushed_sludge)
 
+ def apply_overrides(self, overrides=Dict[str, Any]):
+ """Apply overrides to the service reservoir tank and FWTW.
+
+ Enables a user to override any parameter of the service reservoir 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.service_reservoir_storage_capacity = overrides.pop(
+ "service_reservoir_storage_capacity",
+ self.service_reservoir_storage_capacity,
+ )
+ self.service_reservoir_storage_area = overrides.pop(
+ "service_reservoir_storage_area", self.service_reservoir_storage_area
+ )
+ self.service_reservoir_storage_elevation = overrides.pop(
+ "service_reservoir_storage_elevation",
+ self.service_reservoir_storage_elevation,
+ )
+
+ self.service_reservoir_tank.capacity = self.service_reservoir_storage_capacity
+ self.service_reservoir_tank.area = self.service_reservoir_storage_area
+ self.service_reservoir_tank.datum = self.service_reservoir_storage_elevation
+ super().apply_overrides(overrides)
+
  def treat_water(self):
  """Pulls water, aiming to fill service reservoirs, calls WTW
  treat_current_input, avoids deficit, sends liquor and solids to sewers."""