Apply overrides in Storage and its superclasses

tests/test_storage.py

@@ -181,7 +181,6 @@ def test_river_pull(self):
         constants.set_simple_pollutants()
         river = River(
             name="",
-            depth=2,
             length=200,
             width=20,
             velocity=0.2 * 86400,
@@ -206,7 +205,6 @@ def test_river_depth(self):
         constants.set_simple_pollutants()
         river = River(
             name="",
-            depth=2,
             length=200,
             width=20,
             velocity=0.2 * 86400,
@@ -294,7 +292,88 @@ def test_riverreservoir_environmental(self):
             "temperature": (20 * 10 + 30 * 23) / 40,
         }
         self.assertDictAlmostEqual(d2, reservoir.tank.storage, 15)
-
+
+    def test_storage_overrides(self):
+        constants.set_default_pollutants()
+        storage = Storage(name = '', decays = {"phosphate": {"constant": 1.005, "exponent": 1.005}})
+        storage.apply_overrides({'capacity': 1.43,
+                           'area': 2.36,
+                           'datum': 0.32,
+                           'decays': {"nitrite": {"constant": 10.005, "exponent": 1.105}}
+                           })
+        self.assertEqual(storage.capacity, 1.43)
+        self.assertEqual(storage.tank.capacity, 1.43)
+        self.assertEqual(storage.area, 2.36)
+        self.assertEqual(storage.tank.area, 2.36)
+        self.assertEqual(storage.datum, 0.32)
+        self.assertEqual(storage.tank.datum, 0.32)
+        self.assertDictEqual(storage.decays, {"phosphate": {"constant": 1.005, "exponent": 1.005}, 
+                                              "nitrite": {"constant": 10.005, "exponent": 1.105}})
+        self.assertDictEqual(storage.tank.decays, {"phosphate": {"constant": 1.005, "exponent": 1.005}, 
+                                              "nitrite": {"constant": 10.005, "exponent": 1.105}})
+
+    def test_groundwater_overrides(self):
+        groundwater = Groundwater(name = '', decays = {"phosphate": {"constant": 1.005, "exponent": 1.005}})
+        groundwater.apply_overrides({'residence_time': 27.3,
+                                     'infiltration_threshold': 200,
+                                     'infiltration_pct': 0.667,
+                                     'capacity': 1.43,
+                                     'area': 2.36,
+                                     'datum': 0.32,
+                                     'decays': {"nitrite": {"constant": 10.005, "exponent": 1.105}}
+                                     })
+        self.assertEqual(groundwater.residence_time, 27.3)
+        self.assertEqual(groundwater.infiltration_threshold, 200)
+        self.assertEqual(groundwater.infiltration_pct, 0.667)
+        self.assertEqual(groundwater.capacity, 1.43)
+        self.assertEqual(groundwater.tank.capacity, 1.43)
+        self.assertEqual(groundwater.area, 2.36)
+        self.assertEqual(groundwater.tank.area, 2.36)
+        self.assertEqual(groundwater.datum, 0.32)
+        self.assertEqual(groundwater.tank.datum, 0.32)
+        self.assertDictEqual(groundwater.decays, {"phosphate": {"constant": 1.005, "exponent": 1.005}, 
+                                              "nitrite": {"constant": 10.005, "exponent": 1.105}})
+        self.assertDictEqual(groundwater.tank.decays, {"phosphate": {"constant": 1.005, "exponent": 1.005}, 
+                                              "nitrite": {"constant": 10.005, "exponent": 1.105}})
+
+    def test_river_overrides(self):
+        river = River(name = '')
+        overrides = {'length': 27.3,
+                               'width': 200,
+                               'velocity': 0.667,
+                               'damp': 1.43,
+                               'mrf': 2.36,
+                               'uptake_PNratio': 0.32,
+                               'bulk_density': 5.32,
+                               'denpar_w': 0.432,
+                               'T_wdays': 23.5,
+                               'halfsatINwater': 3.269,
+                               'hsatTP': 2.431,
+                               'limpppar': 6.473,
+                               'prodNpar': 7.821,
+                               'prodPpar': 8.231,
+                               'muptNpar': 6.213,
+                               'muptPpar': 7.021,
+                               'max_temp_lag': 3.213,
+                               'max_phosphorus_lag': 78.321,     
+                               # 'area': 75.2,
+                               # 'capacity': 123
+                               }
+        overrides_to_check = overrides.copy()
+        river.apply_overrides(overrides)
+        for k, v in overrides_to_check.items():
+            if k == 'area':
+                v = 27.3 * 200
+                self.assertEqual(river.tank.area, v)
+            if k == 'capacity':
+                v = constants.UNBOUNDED_CAPACITY
+                self.assertEqual(river.tank.capacity, v)
+            self.assertEqual(getattr(river, k), v)
+
+    def test_riverreservoir_overrides(self):
+        riverreservoir = RiverReservoir(name = '')
+        riverreservoir.apply_overrides({'environmental_flow': 154})
+        self.assertEqual(riverreservoir.environmental_flow, 154)
 
 if __name__ == "__main__":
     unittest.main()

wsimod/nodes/nodes.py

@@ -1182,8 +1182,7 @@ def apply_overrides(self, overrides: Dict[str, Any] = {}):
         Args:
             overrides (dict, optional): Dictionary of overrides. Defaults to {}.
         """
-        self.decays = overrides.pop("decays", 
-                                    self.decays)
+        self.decays.update(overrides.pop("decays", {}))
         super().apply_overrides(overrides)
 
     def end_timestep_decay(self):
@@ -1428,8 +1427,7 @@ def apply_overrides(self, overrides: Dict[str, Any] = {}):
         self.number_of_timesteps = overrides.pop("number_of_timesteps", 
                                                  self.number_of_timesteps)
         self.internal_arc.number_of_timesteps = self.number_of_timesteps
-        self.internal_arc.decays = overrides.pop("decays", 
-                                                 self.internal_arc.decays)
+        self.internal_arc.decays.update(overrides.pop("decays", {}))
         super().apply_overrides(overrides)
 
     def _end_timestep(self):

wsimod/nodes/storage.py

@@ -7,7 +7,7 @@
 
 from wsimod.core import constants
 from wsimod.nodes.nodes import DecayQueueTank, DecayTank, Node, QueueTank, Tank
-
+from typing import Any, Dict
 
 class Storage(Node):
     """"""
@@ -72,6 +72,29 @@ def __init__(
         # Mass balance
         self.mass_balance_ds.append(lambda: self.tank.ds())
 
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        capacity, area, datum, decays.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        # not using pop as these items need to stay in the overrides to be fed into the tank overrides
+        if 'capacity' in overrides.keys():
+            self.capacity = overrides['capacity']
+        if 'area' in overrides.keys():
+            self.area = overrides['area']
+        if 'datum' in overrides.keys():
+            self.datum = overrides['datum']
+        if 'decays' in overrides.keys():
+            self.decays.update(overrides['decays'])
+        # apply tank overrides
+        self.tank.apply_overrides(overrides)
+        super().apply_overrides(overrides)
+
     def push_set_storage(self, vqip):
         """A node wrapper for the tank push_storage.
 
@@ -171,6 +194,24 @@ def __init__(
         self.data_input_dict = data_input_dict
         super().__init__(**kwargs)
 
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        residence_time, infiltration_threshold, infiltration_pct.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.residence_time = overrides.pop("residence_time", 
+                                            self.residence_time)
+        self.infiltration_threshold = overrides.pop("infiltration_threshold", 
+                                                    self.infiltration_threshold)
+        self.infiltration_pct = overrides.pop("infiltration_pct", 
+                                                    self.infiltration_pct)
+        super().apply_overrides(overrides)
+
     def distribute(self):
         """Calculate outflow with residence time and send to Nodes or Rivers."""
         avail = self.tank.get_avail()["volume"] / self.residence_time
@@ -266,6 +307,20 @@ def __init__(self, timearea={0: 1}, data_input_dict={}, **kwargs):
                 initial_storage=self.initial_storage,
             )
 
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        timearea.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.timearea = overrides.pop("timearea", 
+                                      self.timearea)
+        super().apply_overrides(overrides)
+
     def push_set_timearea(self, vqip):
         """Push setting that enables timearea behaviour, (see __init__ for
         description).Used to receive flow that is assumed to occur widely across some
@@ -394,7 +449,6 @@ class River(Storage):
     # TODO non-day timestep
     def __init__(
         self,
-        depth=2,
         length=200,
         width=20,
         velocity=0.2 * constants.M_S_TO_M_DT,
@@ -449,7 +503,6 @@ def __init__(
                 _Units_: m3/day
         """
         # Set parameters
-        self.depth = depth  # [m]
         self.length = length  # [m]
         self.width = width  # [m]
         self.velocity = velocity  # [m/dt]
@@ -490,16 +543,6 @@ def __init__(
         )
         self.muptNpar = 0.001  # [kg/m2/day] nitrogen macrophyte uptake rate
         self.muptPpar = 0.0001  # 0.01, # [kg/m2/day] phosphorus macrophyte uptake rate
-        self.qbank_365_days = [1e6, 1e6]  # [m3/day] store outflow in the previous year
-        self.qbank = (
-            1e6  # [m3/day] bankfull flow = second largest outflow in the previous year
-        )
-        self.qbankcorrpar = 0.001  # [-] correction coefficient for qbank flow
-        self.sedexppar = 1  # [-]
-        self.EPC0 = 0.05 * constants.MG_L_TO_KG_M3  # [mg/l]
-        self.kd_s = 0 * constants.MG_L_TO_KG_M3  # 6 * 1e-6, # [kg/m3]
-        self.kadsdes_s = 2  # 0.9, # [-]
-        self.Dsed = 0.2  # [m]
 
         self.max_temp_lag = 20
         self.lagged_temperatures = []
@@ -550,7 +593,62 @@ def __init__(
     #     self.v_change_vqip(vqip_, min(vqip_['volume'],
     #     self.get_dt_excess()['volume'])) _ = self.tank.push_storage(vqip_, force=True)
     #     return self.extract_vqip(vqip, vqip_)
-
+
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        timearea.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.length = overrides.pop("length", 
+                                      self.length)
+        self.width = overrides.pop("width", 
+                                      self.width)
+        self.velocity = overrides.pop("velocity", 
+                                      self.velocity)
+        self.damp = overrides.pop("damp", 
+                                      self.damp)
+        self.mrf = overrides.pop("mrf", 
+                                      self.mrf)
+        self.uptake_PNratio = overrides.pop("uptake_PNratio", 
+                                      self.uptake_PNratio)
+        self.bulk_density = overrides.pop("bulk_density", 
+                                      self.bulk_density)
+        self.denpar_w = overrides.pop("denpar_w", 
+                                      self.denpar_w)
+        self.T_wdays = overrides.pop("T_wdays", 
+                                      self.T_wdays)
+        self.halfsatINwater = overrides.pop("halfsatINwater", 
+                                      self.halfsatINwater)
+        self.hsatTP = overrides.pop("hsatTP", 
+                                      self.hsatTP)
+        self.limpppar = overrides.pop("limpppar", 
+                                      self.limpppar)
+        self.prodNpar = overrides.pop("prodNpar", 
+                                      self.prodNpar)
+        self.prodPpar = overrides.pop("prodPpar", 
+                                      self.prodPpar)
+        self.muptNpar = overrides.pop("muptNpar", 
+                                      self.muptNpar)
+        self.muptPpar = overrides.pop("muptPpar", 
+                                      self.muptPpar)
+        self.max_temp_lag = overrides.pop("max_temp_lag", 
+                                      self.max_temp_lag)
+        self.max_phosphorus_lag = overrides.pop("max_phosphorus_lag", 
+                                      self.max_phosphorus_lag)
+
+        if 'area' in overrides.keys():
+            print('ERROR: specifying area is depreciated in overrides for river, please specify length and width instead')
+        overrides['area'] = self.length * self.width
+        if 'capacity' in overrides.keys():
+            print('ERROR: specifying capacity is depreciated in overrides for river, it is always set as unbounded capacity')
+        overrides['capacity'] = constants.UNBOUNDED_CAPACITY
+        super().apply_overrides(overrides)
+
     def pull_check_river(self, vqip=None):
         """Check amount of water that can be pulled from river tank and upstream.
 
@@ -936,6 +1034,20 @@ def __init__(self, environmental_flow=0, **kwargs):
 
         self.__class__.__name__ = "Reservoir"
 
+    def apply_overrides(self, overrides = Dict[str, Any]):
+        """Override parameters.
+
+        Enables a user to override any of the following parameters: 
+        environmental_flow.
+
+        Args:
+            overrides (Dict[str, Any]): Dict describing which parameters should
+                be overridden (keys) and new values (values). Defaults to {}.
+        """
+        self.environmental_flow = overrides.pop("environmental_flow", 
+                                                self.environmental_flow)
+        super().apply_overrides(overrides)
+
     def push_set_river_reservoir(self, vqip):
         """Receive water.