diff --git a/.github/workflows/draft-pdf.yml b/.github/workflows/draft-pdf.yml index a8f2aae..60fe70a 100644 --- a/.github/workflows/draft-pdf.yml +++ b/.github/workflows/draft-pdf.yml @@ -6,7 +6,7 @@ jobs: name: Paper Draft steps: - name: Checkout - uses: actions/checkout@v2 + uses: actions/checkout@v4 - name: Build draft PDF uses: openjournals/openjournals-draft-action@master with: @@ -14,7 +14,7 @@ jobs: # This should be the path to the paper within your repo. paper-path: docs/paper/paper.md - name: Upload - uses: actions/upload-artifact@v1 + uses: actions/upload-artifact@v4 with: name: paper # This is the output path where Pandoc will write the compiled diff --git a/docs/demo/examples/test_customise_orchestration_example.yaml b/docs/demo/examples/test_customise_orchestration_example.yaml new file mode 100644 index 0000000..1fd9b2a --- /dev/null +++ b/docs/demo/examples/test_customise_orchestration_example.yaml @@ -0,0 +1,75 @@ +orchestration: +- Groundwater: infiltrate +- Sewer: make_discharge + +nodes: + Sewer: + type_: Sewer + name: my_sewer + capacity: 0.04 + + Groundwater: + type_: Groundwater + name: my_groundwater + capacity: 100 + area: 100 + + River: + type_: Node + name: my_river + + Waste: + type_: Waste + name: my_outlet + +arcs: + 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' \ No newline at end of file diff --git a/tests/test_extensions.py b/tests/test_extensions.py new file mode 100644 index 0000000..7f76f68 --- /dev/null +++ b/tests/test_extensions.py @@ -0,0 +1,174 @@ +from typing import Optional + +import pytest + + +@pytest.fixture +def temp_extension_registry(): + from wsimod.extensions import extensions_registry + + bkp = extensions_registry.copy() + extensions_registry.clear() + yield + extensions_registry.clear() + extensions_registry.update(bkp) + + +def test_register_node_patch(temp_extension_registry): + from wsimod.extensions import extensions_registry, register_node_patch + + # Define a dummy function to patch a node method + @register_node_patch("node_name", "method_name") + def dummy_patch(): + print("Patched method") + + # Check if the patch is registered correctly + assert extensions_registry[("node_name", "method_name", None, False)] == dummy_patch + + # Another function with other arguments + @register_node_patch("node_name", "method_name", item="default", is_attr=True) + def another_dummy_patch(): + print("Another patched method") + + # Check if this other patch is registered correctly + assert ( + extensions_registry[("node_name", "method_name", "default", True)] + == another_dummy_patch + ) + + +def test_apply_patches(temp_extension_registry): + from wsimod.arcs.arcs import Arc + from wsimod.extensions import ( + apply_patches, + extensions_registry, + register_node_patch, + ) + from wsimod.nodes import Node + from wsimod.orchestration.model import Model + + # Create a dummy model + node = Node("dummy_node") + node.dummy_arc = Arc("dummy_arc", in_port=node, out_port=node) + model = Model() + model.nodes[node.name] = node + + # 1. Patch a method + @register_node_patch("dummy_node", "apply_overrides") + def dummy_patch(): + pass + + # 2. Patch an attribute + @register_node_patch("dummy_node", "t", is_attr=True) + def another_dummy_patch(node): + return f"A pathced attribute for {node.name}" + + # 3. Patch a method with an item + @register_node_patch("dummy_node", "pull_set_handler", item="default") + def yet_another_dummy_patch(): + pass + + # 4. Path a method of an attribute + @register_node_patch("dummy_node", "dummy_arc.arc_mass_balance") + def arc_dummy_patch(): + pass + + # Check if all patches are registered + assert len(extensions_registry) == 4 + + # Apply the patches + apply_patches(model) + + # Verify that the patches are applied correctly + assert ( + model.nodes[node.name].apply_overrides.__qualname__ == dummy_patch.__qualname__ + ) + assert ( + model.nodes[node.name]._patched_apply_overrides.__qualname__ + == "Node.apply_overrides" + ) + assert model.nodes[node.name].t == another_dummy_patch(node) + assert model.nodes[node.name]._patched_t == None + assert ( + model.nodes[node.name].pull_set_handler["default"].__qualname__ + == yet_another_dummy_patch.__qualname__ + ) + assert ( + model.nodes[node.name].dummy_arc.arc_mass_balance.__qualname__ + == arc_dummy_patch.__qualname__ + ) + assert ( + model.nodes[node.name].dummy_arc._patched_arc_mass_balance.__qualname__ + == "Arc.arc_mass_balance" + ) + + +def assert_dict_almost_equal(d1: dict, d2: dict, tol: Optional[float] = None): + """Check if two dictionaries are almost equal. + + Args: + d1 (dict): The first dictionary. + d2 (dict): The second dictionary. + tol (float | None, optional): Relative tolerance. Defaults to 1e-6, + `pytest.approx` default. + """ + for key in d1.keys(): + assert d1[key] == pytest.approx(d2[key], rel=tol) + + +def test_path_method_with_reuse(temp_extension_registry): + from wsimod.arcs.arcs import Arc + from wsimod.extensions import apply_patches, register_node_patch + from wsimod.nodes.storage import Reservoir + from wsimod.orchestration.model import Model + + # Create a dummy model + node = Reservoir(name="dummy_node", initial_storage=10, capacity=10) + node.dummy_arc = Arc("dummy_arc", in_port=node, out_port=node) + + vq = node.pull_distributed({"volume": 5}) + assert_dict_almost_equal(vq, node.v_change_vqip(node.empty_vqip(), 5)) + + model = Model() + model.nodes[node.name] = node + + @register_node_patch("dummy_node", "pull_distributed") + def new_pull_distributed(self, vqip, of_type=None, tag="default"): + return self._patched_pull_distributed(vqip, of_type=["Node"], tag=tag) + + # Apply the patches + apply_patches(model) + + # Check appropriate result + assert node.tank.storage["volume"] == 5 + vq = model.nodes[node.name].pull_distributed({"volume": 5}) + assert_dict_almost_equal(vq, node.empty_vqip()) + assert node.tank.storage["volume"] == 5 + + +def test_handler_extensions(temp_extension_registry): + from wsimod.arcs.arcs import Arc + from wsimod.extensions import apply_patches, register_node_patch + from wsimod.nodes import Node + from wsimod.orchestration.model import Model + + # Create a dummy model + node = Node("dummy_node") + node.dummy_arc = Arc("dummy_arc", in_port=node, out_port=node) + model = Model() + model.nodes[node.name] = node + + # 1. Patch a handler + @register_node_patch("dummy_node", "pull_check_handler", item="default") + def dummy_patch(self, *args, **kwargs): + return "dummy_patch" + + # 2. Patch a handler with access to self + @register_node_patch("dummy_node", "pull_set_handler", item="default") + def dummy_patch(self, vqip, *args, **kwargs): + return f"{self.name} - {vqip['volume']}" + + apply_patches(model) + + assert node.pull_check() == "dummy_patch" + assert node.pull_set({"volume": 1}) == "dummy_node - 1" diff --git a/tests/test_model.py b/tests/test_model.py index a004224..e5e404c 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -15,6 +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): @@ -291,6 +292,18 @@ def test_run(self): 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(), "docs", "demo", "examples"), + config_name="test_customise_orchestration_example.yaml", + ) + revised_orchestration = [ + {"Groundwater": "infiltrate"}, + {"Sewer": "make_discharge"}, + ] + self.assertListEqual(my_model.orchestration, revised_orchestration) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index b285f71..576ef7b 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -11,14 +11,7 @@ from wsimod.arcs.arcs import Arc from wsimod.core import constants -from wsimod.nodes.nodes import ( - DecayQueueTank, - DecayTank, - Node, - QueueTank, - ResidenceTank, - Tank, -) +from wsimod.nodes.nodes import Node from wsimod.nodes.storage import Storage from wsimod.nodes.waste import Waste @@ -424,242 +417,6 @@ def test_data_read(self): self.assertEqual(15, node.get_data_input("temperature")) - def test_tank_ds(self): - tank = Tank( - capacity=10, - initial_storage={"volume": 5, "phosphate": 0.4, "temperature": 10}, - ) - tank.end_timestep() - - d1 = {"volume": 2, "phosphate": 0.01, "temperature": 15} - - _ = tank.push_storage(d1) - - diff = tank.ds() - - d2 = {"volume": 2, "phosphate": 0.01, "temperature": 0} - - self.assertDictAlmostEqual(d2, diff, 16) - - def test_ponded(self): - tank = Tank( - capacity=10, - initial_storage={"volume": 15, "phosphate": 0.4, "temperature": 10}, - ) - d1 = {"volume": 5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.pull_ponded() - self.assertDictAlmostEqual(d1, reply) - - def test_tank_get_avail(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - reply = tank.get_avail() - self.assertDictAlmostEqual(d1, reply) - - reply = tank.get_avail({"volume": 2.5}) - d2 = {"volume": 2.5, "phosphate": 0.2, "temperature": 10} - self.assertDictAlmostEqual(d2, reply) - - reply = tank.get_avail({"volume": 10}) - self.assertDictAlmostEqual(d1, reply) - - def test_tank_get_excess(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.get_excess() - self.assertDictAlmostEqual(d2, reply) - - d2 = {"volume": 1, "phosphate": 0.4 * 1 / 7.5, "temperature": 10} - reply = tank.get_excess({"volume": 1}) - self.assertDictAlmostEqual(d2, reply) - - def test_tank_push_storage(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.4, "temperature": 15} - - d3 = {"volume": 2.5, "phosphate": 0.2, "temperature": 15} - reply = tank.push_storage(d2) - self.assertDictAlmostEqual(d3, reply) - - d4 = {"volume": 0, "phosphate": 0, "temperature": 0} - reply = tank.push_storage(d2, force=True) - self.assertDictAlmostEqual(d4, reply) - - def test_tank_pull_storage(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.4 * 5 / 7.5, "temperature": 10} - - reply = tank.pull_storage({"volume": 5}) - self.assertDictAlmostEqual(d2, reply) - - d3 = {"volume": 2.5, "phosphate": 0.4 * 2.5 / 7.5, "temperature": 10} - - reply = tank.pull_storage({"volume": 5}) - - self.assertDictAlmostEqual(d3, reply, 15) - - def test_tank_pull_pollutants(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 5, "phosphate": 0.1, "temperature": 10} - - reply = tank.pull_pollutants(d2) - self.assertDictAlmostEqual(d2, reply) - - reply = tank.pull_pollutants(d2) - d3 = {"volume": 2.5, "phosphate": 0.1, "temperature": 10} - self.assertDictAlmostEqual(d3, reply, 15) - - def test_tank_head(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1, datum=5, area=2.5) - - reply = tank.get_head() - self.assertEqual(8, reply) - - reply = tank.get_head(datum=-1) - self.assertEqual(2, reply) - - reply = tank.get_head(non_head_storage=2) - self.assertEqual(7.2, reply) - - reply = tank.get_head(non_head_storage=10) - self.assertEqual(5, reply) - - def test_evap(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = Tank(capacity=10, initial_storage=d1) - - d2 = {"volume": 0, "phosphate": 0.4, "temperature": 10} - - reply = tank.evaporate(10) - self.assertEqual(7.5, reply) - self.assertDictAlmostEqual(d2, tank.storage) - - def test_residence_tank(self): - d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} - tank = ResidenceTank(residence_time=3, initial_storage=d1) - - d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} - reply = tank.pull_outflow() - self.assertDictAlmostEqual(d2, reply) - - def test_decay_tank(self): - node = Node(name="", data_input_dict={("temperature", 1): 15}) - node.t = 1 - d1 = {"volume": 8, "phosphate": 0.4, "temperature": 10} - - tank = DecayTank( - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - initial_storage=d1, - parent=node, - ) - _ = tank.pull_storage({"volume": 2}) - - d3 = {"volume": -2, "phosphate": -0.1, "temperature": 0} - - diff = tank.decay_ds() - self.assertDictAlmostEqual(d3, diff, 16) - - tank.end_timestep_decay() - - d2 = { - "volume": 6, - "phosphate": 0.3 - 0.3 * 0.001 * 1.005 ** (15 - 20), - "temperature": 10, - } - - self.assertDictAlmostEqual(d2, tank.storage, 16) - - self.assertAlmostEqual( - 0.3 * 0.001 * 1.005 ** (15 - 20), tank.total_decayed["phosphate"] - ) - - def test_queue_push(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - - tank.push_storage(d2) - - d3 = {"volume": 6, "phosphate": 0.5, "temperature": (5 * 10 + 15) / 6} - - self.assertDictAlmostEqual(d3, tank.storage) - self.assertDictAlmostEqual(d1, tank.active_storage) - self.assertDictAlmostEqual(d2, tank.internal_arc.queue[1]) - - tank.push_storage(d2, force=True) - self.assertDictAlmostEqual(d3, tank.active_storage) - - tank.end_timestep() - - d4 = {"volume": 7, "phosphate": 0.6, "temperature": ((5 * 10) + (15 * 2)) / 7} - self.assertDictAlmostEqual(d4, tank.active_storage) - - def test_queue_pull(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - reply = tank.push_storage(d2) - - reply = tank.pull_storage({"volume": 6}) - self.assertDictAlmostEqual(d1, reply) - tank.end_timestep() - self.assertDictAlmostEqual(d2, tank.active_storage) - - def test_queue_pull_exact(self): - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - reply = tank.push_storage(d2) - - reply = tank.pull_storage_exact( - {"volume": 6, "phosphate": 0.1, "temperature": 10} - ) - - d3 = {"volume": 5, "phosphate": 0.1, "temperature": 10} - self.assertDictAlmostEqual(d3, reply) - - reply = tank.pull_storage_exact( - {"volume": 0, "phosphate": 0.6, "temperature": 10} - ) - d4 = {"volume": 0, "phosphate": 0.3, "temperature": 10} - self.assertDictAlmostEqual(d4, reply, 16) - - def test_decay_queue(self): - node = Node(name="", data_input_dict={("temperature", 1): 15}) - node.t = 1 - d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} - tank = DecayQueueTank( - number_of_timesteps=1, - capacity=10, - initial_storage=d1, - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - parent=node, - ) - - d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} - - _ = tank.push_storage(d2) - - tank.end_timestep() - - d4 = { - "volume": 6, - "phosphate": 0.4 + 0.1 * (1 - 0.001 * 1.005 ** (15 - 20)), - "temperature": ((5 * 10) + (15 * 1)) / 6, - } - self.assertDictAlmostEqual(d4, tank.storage, 15) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_tanks.py b/tests/test_tanks.py new file mode 100644 index 0000000..5c6c759 --- /dev/null +++ b/tests/test_tanks.py @@ -0,0 +1,268 @@ +"""Tests for the tanks module.""" + +import unittest +from unittest import TestCase + +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import ( + DecayQueueTank, + DecayTank, + QueueTank, + ResidenceTank, + Tank, +) + + +class MyTestClass(TestCase): + def assertDictAlmostEqual(self, d1, d2, accuracy=19): + """ + + Args: + d1: + d2: + accuracy: + """ + for d in [d1, d2]: + for key, item in d.items(): + d[key] = round(item, accuracy) + self.assertDictEqual(d1, d2) + + def test_tank_ds(self): + tank = Tank( + capacity=10, + initial_storage={"volume": 5, "phosphate": 0.4, "temperature": 10}, + ) + tank.end_timestep() + + d1 = {"volume": 2, "phosphate": 0.01, "temperature": 15} + + _ = tank.push_storage(d1) + + diff = tank.ds() + + d2 = {"volume": 2, "phosphate": 0.01, "temperature": 0} + + self.assertDictAlmostEqual(d2, diff, 16) + + def test_ponded(self): + tank = Tank( + capacity=10, + initial_storage={"volume": 15, "phosphate": 0.4, "temperature": 10}, + ) + d1 = {"volume": 5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.pull_ponded() + self.assertDictAlmostEqual(d1, reply) + + def test_tank_get_avail(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + reply = tank.get_avail() + self.assertDictAlmostEqual(d1, reply) + + reply = tank.get_avail({"volume": 2.5}) + d2 = {"volume": 2.5, "phosphate": 0.2, "temperature": 10} + self.assertDictAlmostEqual(d2, reply) + + reply = tank.get_avail({"volume": 10}) + self.assertDictAlmostEqual(d1, reply) + + def test_tank_get_excess(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.get_excess() + self.assertDictAlmostEqual(d2, reply) + + d2 = {"volume": 1, "phosphate": 0.4 * 1 / 7.5, "temperature": 10} + reply = tank.get_excess({"volume": 1}) + self.assertDictAlmostEqual(d2, reply) + + def test_tank_push_storage(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.4, "temperature": 15} + + d3 = {"volume": 2.5, "phosphate": 0.2, "temperature": 15} + reply = tank.push_storage(d2) + self.assertDictAlmostEqual(d3, reply) + + d4 = {"volume": 0, "phosphate": 0, "temperature": 0} + reply = tank.push_storage(d2, force=True) + self.assertDictAlmostEqual(d4, reply) + + def test_tank_pull_storage(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.4 * 5 / 7.5, "temperature": 10} + + reply = tank.pull_storage({"volume": 5}) + self.assertDictAlmostEqual(d2, reply) + + d3 = {"volume": 2.5, "phosphate": 0.4 * 2.5 / 7.5, "temperature": 10} + + reply = tank.pull_storage({"volume": 5}) + + self.assertDictAlmostEqual(d3, reply, 15) + + def test_tank_pull_pollutants(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 5, "phosphate": 0.1, "temperature": 10} + + reply = tank.pull_pollutants(d2) + self.assertDictAlmostEqual(d2, reply) + + reply = tank.pull_pollutants(d2) + d3 = {"volume": 2.5, "phosphate": 0.1, "temperature": 10} + self.assertDictAlmostEqual(d3, reply, 15) + + def test_tank_head(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1, datum=5, area=2.5) + + reply = tank.get_head() + self.assertEqual(8, reply) + + reply = tank.get_head(datum=-1) + self.assertEqual(2, reply) + + reply = tank.get_head(non_head_storage=2) + self.assertEqual(7.2, reply) + + reply = tank.get_head(non_head_storage=10) + self.assertEqual(5, reply) + + def test_evap(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = Tank(capacity=10, initial_storage=d1) + + d2 = {"volume": 0, "phosphate": 0.4, "temperature": 10} + + reply = tank.evaporate(10) + self.assertEqual(7.5, reply) + self.assertDictAlmostEqual(d2, tank.storage) + + def test_residence_tank(self): + d1 = {"volume": 7.5, "phosphate": 0.4, "temperature": 10} + tank = ResidenceTank(residence_time=3, initial_storage=d1) + + d2 = {"volume": 2.5, "phosphate": 0.4 / 3, "temperature": 10} + reply = tank.pull_outflow() + self.assertDictAlmostEqual(d2, reply) + + def test_decay_tank(self): + node = Node(name="", data_input_dict={("temperature", 1): 15}) + node.t = 1 + d1 = {"volume": 8, "phosphate": 0.4, "temperature": 10} + + tank = DecayTank( + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + initial_storage=d1, + parent=node, + ) + _ = tank.pull_storage({"volume": 2}) + + d3 = {"volume": -2, "phosphate": -0.1, "temperature": 0} + + diff = tank.decay_ds() + self.assertDictAlmostEqual(d3, diff, 16) + + tank.end_timestep_decay() + + d2 = { + "volume": 6, + "phosphate": 0.3 - 0.3 * 0.001 * 1.005 ** (15 - 20), + "temperature": 10, + } + + self.assertDictAlmostEqual(d2, tank.storage, 16) + + self.assertAlmostEqual( + 0.3 * 0.001 * 1.005 ** (15 - 20), tank.total_decayed["phosphate"] + ) + + def test_queue_push(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + + tank.push_storage(d2) + + d3 = {"volume": 6, "phosphate": 0.5, "temperature": (5 * 10 + 15) / 6} + + self.assertDictAlmostEqual(d3, tank.storage) + self.assertDictAlmostEqual(d1, tank.active_storage) + self.assertDictAlmostEqual(d2, tank.internal_arc.queue[1]) + + tank.push_storage(d2, force=True) + self.assertDictAlmostEqual(d3, tank.active_storage) + + tank.end_timestep() + + d4 = {"volume": 7, "phosphate": 0.6, "temperature": ((5 * 10) + (15 * 2)) / 7} + self.assertDictAlmostEqual(d4, tank.active_storage) + + def test_queue_pull(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + reply = tank.push_storage(d2) + + reply = tank.pull_storage({"volume": 6}) + self.assertDictAlmostEqual(d1, reply) + tank.end_timestep() + self.assertDictAlmostEqual(d2, tank.active_storage) + + def test_queue_pull_exact(self): + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = QueueTank(number_of_timesteps=1, capacity=10, initial_storage=d1) + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + reply = tank.push_storage(d2) + + reply = tank.pull_storage_exact( + {"volume": 6, "phosphate": 0.1, "temperature": 10} + ) + + d3 = {"volume": 5, "phosphate": 0.1, "temperature": 10} + self.assertDictAlmostEqual(d3, reply) + + reply = tank.pull_storage_exact( + {"volume": 0, "phosphate": 0.6, "temperature": 10} + ) + d4 = {"volume": 0, "phosphate": 0.3, "temperature": 10} + self.assertDictAlmostEqual(d4, reply, 16) + + def test_decay_queue(self): + node = Node(name="", data_input_dict={("temperature", 1): 15}) + node.t = 1 + d1 = {"volume": 5, "phosphate": 0.4, "temperature": 10} + tank = DecayQueueTank( + number_of_timesteps=1, + capacity=10, + initial_storage=d1, + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + parent=node, + ) + + d2 = {"volume": 1, "phosphate": 0.1, "temperature": 15} + + _ = tank.push_storage(d2) + + tank.end_timestep() + + d4 = { + "volume": 6, + "phosphate": 0.4 + 0.1 * (1 - 0.001 * 1.005 ** (15 - 20)), + "temperature": ((5 * 10) + (15 * 1)) / 6, + } + self.assertDictAlmostEqual(d4, tank.storage, 15) + + +if __name__ == "__main__": + unittest.main() diff --git a/wsimod/extensions.py b/wsimod/extensions.py new file mode 100644 index 0000000..6d4a1be --- /dev/null +++ b/wsimod/extensions.py @@ -0,0 +1,117 @@ +"""This module contains the utilities to extend WSMOD with new features. + +The `register_node_patch` decorator is used to register a function that will be used +instead of a method or attribute of a node. The `apply_patches` function applies all +registered patches to a model. + +Example of patching a method: + +`empty_distributed` will be called instead of `pull_distributed` of "my_node": + + >>> from wsimod.extensions import register_node_patch, apply_patches + >>> @register_node_patch("my_node", "pull_distributed") + >>> def empty_distributed(self, vqip): + >>> return {} + +Attributes, methods of the node, and sub-attributes can be patched. Also, an item of a +list or a dictionary can be patched if the item argument is provided. + +Example of patching an attribute: + +`10` will be assigned to `t`: + + >>> @register_node_patch("my_node", "t", is_attr=True) + >>> def patch_t(node): + >>> return 10 + +Example of patching an attribute item: + +`patch_default_pull_set_handler` will be assigned to +`pull_set_handler["default"]`: + + >>> @register_node_patch("my_node", "pull_set_handler", item="default") + >>> def patch_default_pull_set_handler(self, vqip): + >>> return {} + +If patching a method of an attribute, the `is_attr` argument should be set to `True` and +the target should include the attribute name and the method name, all separated by +periods, eg. `attribute_name.method_name`. + +It should be noted that the patched function should have the same signature as the +original method or attribute, and the return type should be the same as well, otherwise +there will be a runtime error. In particular, the first argument of the patched function +should be the node object itself, which will typically be named `self`. + +The overridden method or attribute can be accessed within the patched function using the +`_patched_{method_name}` attribute of the object, eg. `self._patched_pull_distributed`. +The exception to this is when patching an item, in which case the original item is no +available to be used within the overriding function. + +Finally, the `apply_patches` is called within the `Model.load` method and will apply all +patches in the order they were registered. This means that users need to be careful with +the order of the patches in their extensions files, as they may have interdependencies. + +TODO: Update documentation on extensions files. +""" +from typing import Callable, Hashable + +from .orchestration.model import Model + +extensions_registry: dict[tuple[str, Hashable, bool], Callable] = {} + + +def register_node_patch( + node_name: str, target: str, item: Hashable = None, is_attr: bool = False +) -> Callable: + """Register a function to patch a node method or any of its attributes. + + Args: + node_name (str): The name of the node to patch. + target (str): The target of the object to patch in the form of a string with the + attribute, sub-attribute, etc. and finally method (or attribute) to replace, + sepparated with period, eg. `make_discharge` or + `sewer_tank.pull_storage_exact`. + item (Hashable): Typically a string or an integer indicating the item to replace + in the selected attribue, which should be a list or a dictionary. + is_attr (bool): If True, the decorated function will be called when applying + the patch and the result assigned to the target, instead of assigning the + function itself. In this case, the only argument passed to the function is + the node object. + """ + target_id = (node_name, target, item, is_attr) + if target_id in extensions_registry: + raise ValueError(f"Patch for {target} already registered.") + + def decorator(func): + extensions_registry[target_id] = func + return func + + return decorator + + +def apply_patches(model: Model) -> None: + """Apply all registered patches to the model. + + TODO: Validate signature of the patched methods and type of patched attributes. + + Args: + model (Model): The model to apply the patches to. + """ + for (node_name, target, item, is_attr), func in extensions_registry.items(): + starget = target.split(".") + method = starget.pop() + + # Get the member to patch + node = obj = model.nodes[node_name] + for attr in starget: + obj = getattr(obj, attr) + + # Apply the patch + if item is not None: + obj = getattr(obj, method) + obj[item] = func(node) if is_attr else func.__get__(node, node.__class__) + else: + setattr(obj, f"_patched_{method}", getattr(obj, method)) + setattr( + obj, method, func(node) if is_attr else func.__get__(obj, obj.__class__) + ) diff --git a/wsimod/nodes/land.py b/wsimod/nodes/land.py index f34fb75..7bcaf86 100644 --- a/wsimod/nodes/land.py +++ b/wsimod/nodes/land.py @@ -8,8 +8,9 @@ from math import exp, log, log10, sin from wsimod.core import constants -from wsimod.nodes.nodes import DecayTank, Node, ResidenceTank +from wsimod.nodes.nodes import Node from wsimod.nodes.nutrient_pool import NutrientPool +from wsimod.nodes.tanks import DecayTank, ResidenceTank class Land(Node): diff --git a/wsimod/nodes/nodes.py b/wsimod/nodes/nodes.py index 2e125af..ba7ade4 100644 --- a/wsimod/nodes/nodes.py +++ b/wsimod/nodes/nodes.py @@ -8,9 +8,8 @@ import logging from typing import Any, Dict -from wsimod.arcs.arcs import AltQueueArc, DecayArcAlt from wsimod.core import constants -from wsimod.core.core import DecayObj, WSIObj +from wsimod.core.core import WSIObj class Node(WSIObj): @@ -747,622 +746,3 @@ def general_distribute(self, vqip, of_type = None, tag = 'default', direction = NODES_REGISTRY: dict[str, type[Node]] = {Node.__name__: Node} - - -class Tank(WSIObj): - """""" - - def __init__(self, capacity=0, area=1, datum=10, initial_storage=0): - """A standard storage object. - - Args: - capacity (float, optional): Volumetric tank capacity. Defaults to 0. - area (float, optional): Area of tank. Defaults to 1. - datum (float, optional): Datum of tank base (not currently used in any - functions). Defaults to 10. - initial_storage (optional): Initial storage for tank. - float: Tank will be initialised with zero pollutants and the float - as volume - dict: Tank will be initialised with this VQIP - Defaults to 0 (i.e., no volume, no pollutants). - """ - # Set parameters - self.capacity = capacity - self.area = area - self.datum = datum - self.initial_storage = initial_storage - - WSIObj.__init__(self) # Not sure why I do this rather than super() - - # TODO I don't think the outer if statement is needed - if "initial_storage" in dir(self): - if isinstance(self.initial_storage, dict): - # Assume dict is VQIP describing storage - self.storage = self.copy_vqip(self.initial_storage) - self.storage_ = self.copy_vqip( - self.initial_storage - ) # Lagged storage for mass balance - else: - # Assume number describes initial stroage - self.storage = self.v_change_vqip( - self.empty_vqip(), self.initial_storage - ) - self.storage_ = self.v_change_vqip( - self.empty_vqip(), self.initial_storage - ) # Lagged storage for mass balance - else: - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() # Lagged storage for mass balance - - def ds(self): - """Should be called by parent object to get change in storage. - - Returns: - (dict): Change in storage - """ - return self.ds_vqip(self.storage, self.storage_) - - def pull_ponded(self): - """Pull any volume that is above the tank's capacity. - - Returns: - ponded (vqip): Amount of ponded water that has been removed from the - tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, - 'phosphate' : 0.2}) - >>> print(my_tank.storage) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.pull_ponded()) - {'volume' : 1, 'phosphate' : 0.02} - >>> print(my_tank.storage) - {'volume' : 9, 'phosphate' : 0.18} - """ - # Get amount - ponded = max(self.storage["volume"] - self.capacity, 0) - # Pull from tank - ponded = self.pull_storage({"volume": ponded}) - return ponded - - def get_avail(self, vqip=None): - """Get minimum of the amount of water in storage and vqip (if provided). - - Args: - vqip (dict, optional): Maximum water required (only 'volume' is used). - Defaults to None. - - Returns: - reply (dict): Water available - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, - 'phosphate' : 0.2}) - >>> print(my_tank.storage) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.get_avail()) - {'volume' : 10, 'phosphate' : 0.2} - >>> print(my_tank.get_avail({'volume' : 1})) - {'volume' : 1, 'phosphate' : 0.02} - """ - reply = self.copy_vqip(self.storage) - if vqip is None: - # Return storage - return reply - else: - # Adjust storage pollutants to match volume in vqip - reply = self.v_change_vqip(reply, min(reply["volume"], vqip["volume"])) - return reply - - def get_excess(self, vqip=None): - """Get difference between current storage and tank capacity. - - Args: - vqip (dict, optional): Maximum capacity required (only 'volume' is - used). Defaults to None. - - Returns: - (dict): Difference available - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.get_excess()) - {'volume' : 4, 'phosphate' : 0.16} - >>> print(my_tank.get_excess({'volume' : 2})) - {'volume' : 2, 'phosphate' : 0.08} - """ - vol = max(self.capacity - self.storage["volume"], 0) - if vqip is not None: - vol = min(vqip["volume"], vol) - - # Adjust storage pollutants to match volume in vqip - # TODO the v_change_vqip in the reply here is a weird default (if a VQIP is not - # provided) - return self.v_change_vqip(self.storage, vol) - - def push_storage(self, vqip, force=False): - """Push water into tank, updating the storage VQIP. Force argument can be used - to ignore tank capacity. - - Args: - vqip (dict): VQIP amount to be pushed - force (bool, optional): Argument used to cause function to ignore tank - capacity, possibly resulting in pooling. Defaults to False. - - Returns: - reply (dict): A VQIP of water not successfully pushed to the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> constants.POLLUTANTS = ['phosphate'] - >>> constants.NON_ADDITIVE_POLLUTANTS = [] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> my_push = {'volume' : 10, 'phosphate' : 0.5} - >>> reply = my_tank.push_storage(my_push) - >>> print(reply) - {'volume' : 6, 'phosphate' : 0.3} - >>> print(my_tank.storage) - {'volume': 9.0, 'phosphate': 0.4} - >>> print(my_tank.push_storage(reply, force = True)) - {'phosphate': 0, 'volume': 0} - >>> print(my_tank.storage) - {'volume': 15.0, 'phosphate': 0.7} - """ - if force: - # Directly add request to storage - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - # Check whether request can be met - excess = self.get_excess()["volume"] - - # Adjust accordingly - reply = max(vqip["volume"] - excess, 0) - reply = self.v_change_vqip(vqip, reply) - entered = self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) - - # Update storage - self.storage = self.sum_vqip(self.storage, entered) - - return reply - - def pull_storage(self, vqip): - """Pull water from tank, updating the storage VQIP. Pollutants are removed from - tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). - - Args: - vqip (dict): VQIP amount to be pulled, (only 'volume' key is needed) - - Returns: - reply (dict): A VQIP water successfully pulled from the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.pull_storage({'volume' : 6})) - {'volume': 5.0, 'phosphate': 0.2} - >>> print(my_tank.storage) - {'volume': 0, 'phosphate': 0} - """ - # Pull from Tank by volume (taking pollutants in proportion) - if self.storage["volume"] == 0: - return self.empty_vqip() - - # Adjust based on available volume - reply = min(vqip["volume"], self.storage["volume"]) - - # Update reply to vqip (in proportion to concentration in storage) - reply = self.v_change_vqip(self.storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - - return reply - - def pull_pollutants(self, vqip): - """Pull water from tank, updating the storage VQIP. Pollutants are removed from - tank in according to their values in vqip. - - Args: - vqip (dict): VQIP amount to be pulled - - Returns: - vqip (dict): A VQIP water successfully pulled from the tank - - Examples: - >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] - >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, - 'phosphate' : 0.2}) - >>> print(my_tank.pull_pollutants({'volume' : 2, 'phosphate' : 0.15})) - {'volume': 2.0, 'phosphate': 0.15} - >>> print(my_tank.storage) - {'volume': 3, 'phosphate': 0.05} - """ - # Adjust based on available mass - for pol in constants.ADDITIVE_POLLUTANTS + ["volume"]: - vqip[pol] = min(self.storage[pol], vqip[pol]) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, vqip) - return vqip - - def get_head(self, datum=None, non_head_storage=0): - """Area volume calculation for head calcuations. Datum and storage that does not - contribute to head can be specified. - - Args: - datum (float, optional): Value to add to pressure head in tank. - Defaults to None. - non_head_storage (float, optional): Amount of storage that does - not contribute to generation of head. The tank must exceed - this value to generate any pressure head. Defaults to 0. - - Returns: - head (float): Total head in tank - - Examples: - >>> my_tank = Tank(datum = 10, initial_storage = 5, capacity = 10, area = 2) - >>> print(my_tank.get_head()) - 12.5 - >>> print(my_tank.get_head(non_head_storage = 1)) - 12 - >>> print(my_tank.get_head(non_head_storage = 1, datum = 0)) - 2 - """ - # If datum not provided use object datum - if datum is None: - datum = self.datum - - # Calculate pressure head generating storage - head_storage = max(self.storage["volume"] - non_head_storage, 0) - - # Perform head calculation - head = head_storage / self.area + datum - - return head - - def evaporate(self, evap): - """Wrapper for v_distill_vqip to apply a volumetric subtraction from tank - storage. Volume removed from storage and no change in pollutant values. - - Args: - evap (float): Volume to evaporate - - Returns: - evap (float): Volumetric amount of evaporation successfully removed - """ - avail = self.get_avail()["volume"] - - evap = min(evap, avail) - self.storage = self.v_distill_vqip(self.storage, evap) - return evap - - ##Old function no longer needed (check it is not used anywhere and remove) - def push_total(self, vqip): - """ - - Args: - vqip: - - Returns: - - """ - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - ##Old function no longer needed (check it is not used anywhere and remove) - def push_total_c(self, vqip): - """ - - Args: - vqip: - - Returns: - - """ - # Push vqip to storage where pollutants are given as a concentration rather - # than storage - vqip = self.concentration_to_total(self.vqip) - self.storage = self.sum_vqip(self.storage, vqip) - return self.empty_vqip() - - def end_timestep(self): - """Function to be called by parent object, tracks previously timestep's - storage.""" - self.storage_ = self.copy_vqip(self.storage) - - def reinit(self): - """Set storage to an empty VQIP.""" - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() - - -class ResidenceTank(Tank): - """""" - - def __init__(self, residence_time=2, **kwargs): - """A tank that has a residence time property that limits storage pulled from the - 'pull_outflow' function. - - Args: - residence_time (float, optional): Residence time, in theory given - in timesteps, in practice it just means that storage / - residence time can be pulled each time pull_outflow is called. - Defaults to 2. - """ - self.residence_time = residence_time - super().__init__(**kwargs) - - def pull_outflow(self): - """Pull storage by residence time from the tank, updating tank storage. - - Returns: - outflow (dict): A VQIP with volume of pulled volume and pollutants - proportionate to the tank's pollutants - """ - # Calculate outflow - outflow = self.storage["volume"] / self.residence_time - # Update pollutant amounts - outflow = self.v_change_vqip(self.storage, outflow) - # Remove from tank - outflow = self.pull_storage(outflow) - return outflow - - -class DecayTank(Tank, DecayObj): - """""" - - def __init__(self, decays={}, parent=None, **kwargs): - """A tank that has DecayObj functions. Decay occurs in end_timestep, after - updating state variables. In this sense, decay is occurring at the very - beginning of the timestep. - - Args: - decays (dict): A dict of dicts containing a key for each pollutant that - decays and, within that, a key for each parameter (a constant and - exponent) - parent (object): An object that can be used to read temperature data from - """ - # Store parameters - self.parent = parent - - # Initialise Tank - Tank.__init__(self, **kwargs) - - # Initialise decay object - DecayObj.__init__(self, decays) - - # Update timestep and ds functions - self.end_timestep = self.end_timestep_decay - self.ds = self.decay_ds - - def end_timestep_decay(self): - """Update state variables and call make_decay.""" - self.total_decayed = self.empty_vqip() - self.storage_ = self.copy_vqip(self.storage) - - self.storage = self.make_decay(self.storage) - - def decay_ds(self): - """Track storage and amount decayed. - - Returns: - ds (dict): A VQIP of change in storage and total decayed - """ - ds = self.ds_vqip(self.storage, self.storage_) - ds = self.sum_vqip(ds, self.total_decayed) - return ds - - -class QueueTank(Tank): - """""" - - def __init__(self, number_of_timesteps=0, **kwargs): - """A tank with an internal queue arc, whose queue must be completed before - storage is available for use. The storage that has completed the queue is under - the 'active_storage' property. - - Args: - number_of_timesteps (int, optional): Built in delay for the internal - queue - it is always added to the queue time, although delay can be - provided with pushes only. Defaults to 0. - """ - # Set parameters - self.number_of_timesteps = number_of_timesteps - - super().__init__(**kwargs) - self.end_timestep = self._end_timestep - self.active_storage = self.copy_vqip(self.storage) - - # TODO enable queue to be initialised not empty - self.out_arcs = {} - self.in_arcs = {} - # Create internal queue arc - self.internal_arc = AltQueueArc( - in_port=self, out_port=self, number_of_timesteps=self.number_of_timesteps - ) - # TODO should mass balance call internal arc (is this arc called in arc mass - # balance?) - - def get_avail(self): - """Return the active_storage of the tank. - - Returns: - (dict): VQIP of active_storage - """ - return self.copy_vqip(self.active_storage) - - def push_storage(self, vqip, time=0, force=False): - """Push storage into QueueTank, applying travel time, unless forced. - - Args: - vqip (dict): A VQIP of the amount to push - time (int, optional): Number of timesteps to spend in queue, in addition - to number_of_timesteps property of internal_arc. Defaults to 0. - force (bool, optional): Force property that will ignore tank capacity - and ignore travel time. Defaults to False. - - Returns: - reply (dict): A VQIP of water that could not be received by the tank - """ - if force: - # Directly add request to storage, skipping queue - self.storage = self.sum_vqip(self.storage, vqip) - self.active_storage = self.sum_vqip(self.active_storage, vqip) - return self.empty_vqip() - - # Push to QueueTank - reply = self.internal_arc.send_push_request(vqip, force=force, time=time) - # Update storage - # TODO storage won't be accurately tracking temperature.. - self.storage = self.sum_vqip( - self.storage, self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) - ) - return reply - - def pull_storage(self, vqip): - """Pull storage from the QueueTank, only water in active_storage is available. - Returning water pulled and updating tank states. Pollutants are removed from - tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). - - Args: - vqip (dict): VQIP amount to pull, only 'volume' property is used - - Returns: - reply (dict): VQIP amount that was pulled - """ - # Adjust based on available volume - reply = min(vqip["volume"], self.active_storage["volume"]) - - # Update reply to vqip - reply = self.v_change_vqip(self.active_storage, reply) - - # Extract from active_storage - self.active_storage = self.extract_vqip(self.active_storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - - return reply - - def pull_storage_exact(self, vqip): - """Pull storage from the QueueTank, only water in active_storage is available. - Pollutants are removed from tank in according to their values in vqip. - - Args: - vqip (dict): A VQIP amount to pull - - Returns: - reply (dict): A VQIP amount successfully pulled - """ - # Adjust based on available - reply = self.copy_vqip(vqip) - for pol in ["volume"] + constants.ADDITIVE_POLLUTANTS: - reply[pol] = min(reply[pol], self.active_storage[pol]) - - # Pull from QueueTank - self.active_storage = self.extract_vqip(self.active_storage, reply) - - # Extract from storage - self.storage = self.extract_vqip(self.storage, reply) - return reply - - def push_check(self, vqip=None, tag="default"): - """Wrapper for get_excess but applies comparison to volume in VQIP. - Needed to enable use of internal_arc, which assumes it is connecting nodes . - rather than tanks. - NOTE: this is intended only for use with the internal_arc. Pushing to - QueueTanks should use 'push_storage'. - - Args: - vqip (dict, optional): VQIP amount to push. Defaults to None. - tag (str, optional): Tag, see Node, don't think it should actually be - used for a QueueTank since there are no handlers. Defaults to - 'default'. - - Returns: - excess (dict): a VQIP amount of excess capacity - """ - # TODO does behaviour for volume = None need to be defined? - excess = self.get_excess() - if vqip is not None: - excess["volume"] = min(vqip["volume"], excess["volume"]) - return excess - - def push_set(self, vqip, tag="default"): - """Behaves differently from normal push setting, it assumes sufficient tank - capacity and receives VQIPs that have reached the END of the internal_arc. - NOTE: this is intended only for use with the internal_arc. Pushing to - QueueTanks should use 'push_storage'. - - Args: - vqip (dict): VQIP amount to push - tag (str, optional): Tag, see Node, don't think it should actually be - used for a QueueTank since there are no handlers. Defaults to - 'default'. - - Returns: - (dict): Returns empty VQIP, indicating all water received (since it - assumes capacity was checked before entering the internal arc) - """ - # Update active_storage (since it has reached the end of the internal_arc) - self.active_storage = self.sum_vqip(self.active_storage, vqip) - - return self.empty_vqip() - - def _end_timestep(self): - """Wrapper for end_timestep that also ends the timestep in the internal_arc.""" - self.internal_arc.end_timestep() - self.internal_arc.update_queue() - self.storage_ = self.copy_vqip(self.storage) - - def reinit(self): - """Zeros storages and arc.""" - self.internal_arc.reinit() - self.storage = self.empty_vqip() - self.storage_ = self.empty_vqip() - self.active_storage = self.empty_vqip() - - -class DecayQueueTank(QueueTank): - """""" - - def __init__(self, decays={}, parent=None, number_of_timesteps=1, **kwargs): - """Adds a DecayAltArc in QueueTank to enable decay to occur within the - internal_arc queue. - - Args: - decays (dict): A dict of dicts containing a key for each pollutant and, - within that, a key for each parameter (a constant and exponent) - parent (object): An object that can be used to read temperature data from - number_of_timesteps (int, optional): Built in delay for the internal - queue - it is always added to the queue time, although delay can be - provided with pushes only. Defaults to 0. - """ - # Initialise QueueTank - super().__init__(number_of_timesteps=number_of_timesteps, **kwargs) - # Replace internal_arc with a DecayArcAlt - self.internal_arc = DecayArcAlt( - in_port=self, - out_port=self, - number_of_timesteps=number_of_timesteps, - parent=parent, - decays=decays, - ) - - self.end_timestep = self._end_timestep - - def _end_timestep(self): - """End timestep wrapper that removes decayed pollutants and calls internal - arc.""" - # TODO Should the active storage decay if decays are given (probably.. though - # that sounds like a nightmare)? - self.storage = self.extract_vqip(self.storage, self.internal_arc.total_decayed) - self.storage_ = self.copy_vqip(self.storage) - self.internal_arc.end_timestep() diff --git a/wsimod/nodes/sewer.py b/wsimod/nodes/sewer.py index af56665..545495c 100644 --- a/wsimod/nodes/sewer.py +++ b/wsimod/nodes/sewer.py @@ -5,7 +5,8 @@ Converted to totals on 2022-05-03 """ from wsimod.core import constants -from wsimod.nodes.nodes import Node, QueueTank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import QueueTank class Sewer(Node): diff --git a/wsimod/nodes/storage.py b/wsimod/nodes/storage.py index eab6866..2cd7599 100644 --- a/wsimod/nodes/storage.py +++ b/wsimod/nodes/storage.py @@ -6,7 +6,8 @@ from math import exp from wsimod.core import constants -from wsimod.nodes.nodes import DecayQueueTank, DecayTank, Node, QueueTank, Tank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import DecayQueueTank, DecayTank, QueueTank, Tank class Storage(Node): diff --git a/wsimod/nodes/tanks.py b/wsimod/nodes/tanks.py new file mode 100644 index 0000000..4806d6a --- /dev/null +++ b/wsimod/nodes/tanks.py @@ -0,0 +1,624 @@ +"""Module for defining tanks.""" + +from wsimod.arcs.arcs import AltQueueArc, DecayArcAlt +from wsimod.core import constants +from wsimod.core.core import DecayObj, WSIObj + + +class Tank(WSIObj): + """""" + + def __init__(self, capacity=0, area=1, datum=10, initial_storage=0): + """A standard storage object. + + Args: + capacity (float, optional): Volumetric tank capacity. Defaults to 0. + area (float, optional): Area of tank. Defaults to 1. + datum (float, optional): Datum of tank base (not currently used in any + functions). Defaults to 10. + initial_storage (optional): Initial storage for tank. + float: Tank will be initialised with zero pollutants and the float + as volume + dict: Tank will be initialised with this VQIP + Defaults to 0 (i.e., no volume, no pollutants). + """ + # Set parameters + self.capacity = capacity + self.area = area + self.datum = datum + self.initial_storage = initial_storage + + WSIObj.__init__(self) # Not sure why I do this rather than super() + + # TODO I don't think the outer if statement is needed + if "initial_storage" in dir(self): + if isinstance(self.initial_storage, dict): + # Assume dict is VQIP describing storage + self.storage = self.copy_vqip(self.initial_storage) + self.storage_ = self.copy_vqip( + self.initial_storage + ) # Lagged storage for mass balance + else: + # Assume number describes initial stroage + self.storage = self.v_change_vqip( + self.empty_vqip(), self.initial_storage + ) + self.storage_ = self.v_change_vqip( + self.empty_vqip(), self.initial_storage + ) # Lagged storage for mass balance + else: + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() # Lagged storage for mass balance + + def ds(self): + """Should be called by parent object to get change in storage. + + Returns: + (dict): Change in storage + """ + return self.ds_vqip(self.storage, self.storage_) + + def pull_ponded(self): + """Pull any volume that is above the tank's capacity. + + Returns: + ponded (vqip): Amount of ponded water that has been removed from the + tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, + 'phosphate' : 0.2}) + >>> print(my_tank.storage) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.pull_ponded()) + {'volume' : 1, 'phosphate' : 0.02} + >>> print(my_tank.storage) + {'volume' : 9, 'phosphate' : 0.18} + """ + # Get amount + ponded = max(self.storage["volume"] - self.capacity, 0) + # Pull from tank + ponded = self.pull_storage({"volume": ponded}) + return ponded + + def get_avail(self, vqip=None): + """Get minimum of the amount of water in storage and vqip (if provided). + + Args: + vqip (dict, optional): Maximum water required (only 'volume' is used). + Defaults to None. + + Returns: + reply (dict): Water available + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 10, + 'phosphate' : 0.2}) + >>> print(my_tank.storage) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.get_avail()) + {'volume' : 10, 'phosphate' : 0.2} + >>> print(my_tank.get_avail({'volume' : 1})) + {'volume' : 1, 'phosphate' : 0.02} + """ + reply = self.copy_vqip(self.storage) + if vqip is None: + # Return storage + return reply + else: + # Adjust storage pollutants to match volume in vqip + reply = self.v_change_vqip(reply, min(reply["volume"], vqip["volume"])) + return reply + + def get_excess(self, vqip=None): + """Get difference between current storage and tank capacity. + + Args: + vqip (dict, optional): Maximum capacity required (only 'volume' is + used). Defaults to None. + + Returns: + (dict): Difference available + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.get_excess()) + {'volume' : 4, 'phosphate' : 0.16} + >>> print(my_tank.get_excess({'volume' : 2})) + {'volume' : 2, 'phosphate' : 0.08} + """ + vol = max(self.capacity - self.storage["volume"], 0) + if vqip is not None: + vol = min(vqip["volume"], vol) + + # Adjust storage pollutants to match volume in vqip + # TODO the v_change_vqip in the reply here is a weird default (if a VQIP is not + # provided) + return self.v_change_vqip(self.storage, vol) + + def push_storage(self, vqip, force=False): + """Push water into tank, updating the storage VQIP. Force argument can be used + to ignore tank capacity. + + Args: + vqip (dict): VQIP amount to be pushed + force (bool, optional): Argument used to cause function to ignore tank + capacity, possibly resulting in pooling. Defaults to False. + + Returns: + reply (dict): A VQIP of water not successfully pushed to the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> constants.POLLUTANTS = ['phosphate'] + >>> constants.NON_ADDITIVE_POLLUTANTS = [] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> my_push = {'volume' : 10, 'phosphate' : 0.5} + >>> reply = my_tank.push_storage(my_push) + >>> print(reply) + {'volume' : 6, 'phosphate' : 0.3} + >>> print(my_tank.storage) + {'volume': 9.0, 'phosphate': 0.4} + >>> print(my_tank.push_storage(reply, force = True)) + {'phosphate': 0, 'volume': 0} + >>> print(my_tank.storage) + {'volume': 15.0, 'phosphate': 0.7} + """ + if force: + # Directly add request to storage + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + # Check whether request can be met + excess = self.get_excess()["volume"] + + # Adjust accordingly + reply = max(vqip["volume"] - excess, 0) + reply = self.v_change_vqip(vqip, reply) + entered = self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) + + # Update storage + self.storage = self.sum_vqip(self.storage, entered) + + return reply + + def pull_storage(self, vqip): + """Pull water from tank, updating the storage VQIP. Pollutants are removed from + tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). + + Args: + vqip (dict): VQIP amount to be pulled, (only 'volume' key is needed) + + Returns: + reply (dict): A VQIP water successfully pulled from the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.pull_storage({'volume' : 6})) + {'volume': 5.0, 'phosphate': 0.2} + >>> print(my_tank.storage) + {'volume': 0, 'phosphate': 0} + """ + # Pull from Tank by volume (taking pollutants in proportion) + if self.storage["volume"] == 0: + return self.empty_vqip() + + # Adjust based on available volume + reply = min(vqip["volume"], self.storage["volume"]) + + # Update reply to vqip (in proportion to concentration in storage) + reply = self.v_change_vqip(self.storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + + return reply + + def pull_pollutants(self, vqip): + """Pull water from tank, updating the storage VQIP. Pollutants are removed from + tank in according to their values in vqip. + + Args: + vqip (dict): VQIP amount to be pulled + + Returns: + vqip (dict): A VQIP water successfully pulled from the tank + + Examples: + >>> constants.ADDITIVE_POLLUTANTS = ['phosphate'] + >>> my_tank = Tank(capacity = 9, initial_storage = {'volume' : 5, + 'phosphate' : 0.2}) + >>> print(my_tank.pull_pollutants({'volume' : 2, 'phosphate' : 0.15})) + {'volume': 2.0, 'phosphate': 0.15} + >>> print(my_tank.storage) + {'volume': 3, 'phosphate': 0.05} + """ + # Adjust based on available mass + for pol in constants.ADDITIVE_POLLUTANTS + ["volume"]: + vqip[pol] = min(self.storage[pol], vqip[pol]) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, vqip) + return vqip + + def get_head(self, datum=None, non_head_storage=0): + """Area volume calculation for head calcuations. Datum and storage that does not + contribute to head can be specified. + + Args: + datum (float, optional): Value to add to pressure head in tank. + Defaults to None. + non_head_storage (float, optional): Amount of storage that does + not contribute to generation of head. The tank must exceed + this value to generate any pressure head. Defaults to 0. + + Returns: + head (float): Total head in tank + + Examples: + >>> my_tank = Tank(datum = 10, initial_storage = 5, capacity = 10, area = 2) + >>> print(my_tank.get_head()) + 12.5 + >>> print(my_tank.get_head(non_head_storage = 1)) + 12 + >>> print(my_tank.get_head(non_head_storage = 1, datum = 0)) + 2 + """ + # If datum not provided use object datum + if datum is None: + datum = self.datum + + # Calculate pressure head generating storage + head_storage = max(self.storage["volume"] - non_head_storage, 0) + + # Perform head calculation + head = head_storage / self.area + datum + + return head + + def evaporate(self, evap): + """Wrapper for v_distill_vqip to apply a volumetric subtraction from tank + storage. Volume removed from storage and no change in pollutant values. + + Args: + evap (float): Volume to evaporate + + Returns: + evap (float): Volumetric amount of evaporation successfully removed + """ + avail = self.get_avail()["volume"] + + evap = min(evap, avail) + self.storage = self.v_distill_vqip(self.storage, evap) + return evap + + ##Old function no longer needed (check it is not used anywhere and remove) + def push_total(self, vqip): + """ + + Args: + vqip: + + Returns: + + """ + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + ##Old function no longer needed (check it is not used anywhere and remove) + def push_total_c(self, vqip): + """ + + Args: + vqip: + + Returns: + + """ + # Push vqip to storage where pollutants are given as a concentration rather + # than storage + vqip = self.concentration_to_total(self.vqip) + self.storage = self.sum_vqip(self.storage, vqip) + return self.empty_vqip() + + def end_timestep(self): + """Function to be called by parent object, tracks previously timestep's + storage.""" + self.storage_ = self.copy_vqip(self.storage) + + def reinit(self): + """Set storage to an empty VQIP.""" + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() + + +class ResidenceTank(Tank): + """""" + + def __init__(self, residence_time=2, **kwargs): + """A tank that has a residence time property that limits storage pulled from the + 'pull_outflow' function. + + Args: + residence_time (float, optional): Residence time, in theory given + in timesteps, in practice it just means that storage / + residence time can be pulled each time pull_outflow is called. + Defaults to 2. + """ + self.residence_time = residence_time + super().__init__(**kwargs) + + def pull_outflow(self): + """Pull storage by residence time from the tank, updating tank storage. + + Returns: + outflow (dict): A VQIP with volume of pulled volume and pollutants + proportionate to the tank's pollutants + """ + # Calculate outflow + outflow = self.storage["volume"] / self.residence_time + # Update pollutant amounts + outflow = self.v_change_vqip(self.storage, outflow) + # Remove from tank + outflow = self.pull_storage(outflow) + return outflow + + +class DecayTank(Tank, DecayObj): + """""" + + def __init__(self, decays={}, parent=None, **kwargs): + """A tank that has DecayObj functions. Decay occurs in end_timestep, after + updating state variables. In this sense, decay is occurring at the very + beginning of the timestep. + + Args: + decays (dict): A dict of dicts containing a key for each pollutant that + decays and, within that, a key for each parameter (a constant and + exponent) + parent (object): An object that can be used to read temperature data from + """ + # Store parameters + self.parent = parent + + # Initialise Tank + Tank.__init__(self, **kwargs) + + # Initialise decay object + DecayObj.__init__(self, decays) + + # Update timestep and ds functions + self.end_timestep = self.end_timestep_decay + self.ds = self.decay_ds + + def end_timestep_decay(self): + """Update state variables and call make_decay.""" + self.total_decayed = self.empty_vqip() + self.storage_ = self.copy_vqip(self.storage) + + self.storage = self.make_decay(self.storage) + + def decay_ds(self): + """Track storage and amount decayed. + + Returns: + ds (dict): A VQIP of change in storage and total decayed + """ + ds = self.ds_vqip(self.storage, self.storage_) + ds = self.sum_vqip(ds, self.total_decayed) + return ds + + +class QueueTank(Tank): + """""" + + def __init__(self, number_of_timesteps=0, **kwargs): + """A tank with an internal queue arc, whose queue must be completed before + storage is available for use. The storage that has completed the queue is under + the 'active_storage' property. + + Args: + number_of_timesteps (int, optional): Built in delay for the internal + queue - it is always added to the queue time, although delay can be + provided with pushes only. Defaults to 0. + """ + # Set parameters + self.number_of_timesteps = number_of_timesteps + + super().__init__(**kwargs) + self.end_timestep = self._end_timestep + self.active_storage = self.copy_vqip(self.storage) + + # TODO enable queue to be initialised not empty + self.out_arcs = {} + self.in_arcs = {} + # Create internal queue arc + self.internal_arc = AltQueueArc( + in_port=self, out_port=self, number_of_timesteps=self.number_of_timesteps + ) + # TODO should mass balance call internal arc (is this arc called in arc mass + # balance?) + + def get_avail(self): + """Return the active_storage of the tank. + + Returns: + (dict): VQIP of active_storage + """ + return self.copy_vqip(self.active_storage) + + def push_storage(self, vqip, time=0, force=False): + """Push storage into QueueTank, applying travel time, unless forced. + + Args: + vqip (dict): A VQIP of the amount to push + time (int, optional): Number of timesteps to spend in queue, in addition + to number_of_timesteps property of internal_arc. Defaults to 0. + force (bool, optional): Force property that will ignore tank capacity + and ignore travel time. Defaults to False. + + Returns: + reply (dict): A VQIP of water that could not be received by the tank + """ + if force: + # Directly add request to storage, skipping queue + self.storage = self.sum_vqip(self.storage, vqip) + self.active_storage = self.sum_vqip(self.active_storage, vqip) + return self.empty_vqip() + + # Push to QueueTank + reply = self.internal_arc.send_push_request(vqip, force=force, time=time) + # Update storage + # TODO storage won't be accurately tracking temperature.. + self.storage = self.sum_vqip( + self.storage, self.v_change_vqip(vqip, vqip["volume"] - reply["volume"]) + ) + return reply + + def pull_storage(self, vqip): + """Pull storage from the QueueTank, only water in active_storage is available. + Returning water pulled and updating tank states. Pollutants are removed from + tank in proportion to 'volume' in vqip (pollutant values in vqip are ignored). + + Args: + vqip (dict): VQIP amount to pull, only 'volume' property is used + + Returns: + reply (dict): VQIP amount that was pulled + """ + # Adjust based on available volume + reply = min(vqip["volume"], self.active_storage["volume"]) + + # Update reply to vqip + reply = self.v_change_vqip(self.active_storage, reply) + + # Extract from active_storage + self.active_storage = self.extract_vqip(self.active_storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + + return reply + + def pull_storage_exact(self, vqip): + """Pull storage from the QueueTank, only water in active_storage is available. + Pollutants are removed from tank in according to their values in vqip. + + Args: + vqip (dict): A VQIP amount to pull + + Returns: + reply (dict): A VQIP amount successfully pulled + """ + # Adjust based on available + reply = self.copy_vqip(vqip) + for pol in ["volume"] + constants.ADDITIVE_POLLUTANTS: + reply[pol] = min(reply[pol], self.active_storage[pol]) + + # Pull from QueueTank + self.active_storage = self.extract_vqip(self.active_storage, reply) + + # Extract from storage + self.storage = self.extract_vqip(self.storage, reply) + return reply + + def push_check(self, vqip=None, tag="default"): + """Wrapper for get_excess but applies comparison to volume in VQIP. + Needed to enable use of internal_arc, which assumes it is connecting nodes . + rather than tanks. + NOTE: this is intended only for use with the internal_arc. Pushing to + QueueTanks should use 'push_storage'. + + Args: + vqip (dict, optional): VQIP amount to push. Defaults to None. + tag (str, optional): Tag, see Node, don't think it should actually be + used for a QueueTank since there are no handlers. Defaults to + 'default'. + + Returns: + excess (dict): a VQIP amount of excess capacity + """ + # TODO does behaviour for volume = None need to be defined? + excess = self.get_excess() + if vqip is not None: + excess["volume"] = min(vqip["volume"], excess["volume"]) + return excess + + def push_set(self, vqip, tag="default"): + """Behaves differently from normal push setting, it assumes sufficient tank + capacity and receives VQIPs that have reached the END of the internal_arc. + NOTE: this is intended only for use with the internal_arc. Pushing to + QueueTanks should use 'push_storage'. + + Args: + vqip (dict): VQIP amount to push + tag (str, optional): Tag, see Node, don't think it should actually be + used for a QueueTank since there are no handlers. Defaults to + 'default'. + + Returns: + (dict): Returns empty VQIP, indicating all water received (since it + assumes capacity was checked before entering the internal arc) + """ + # Update active_storage (since it has reached the end of the internal_arc) + self.active_storage = self.sum_vqip(self.active_storage, vqip) + + return self.empty_vqip() + + def _end_timestep(self): + """Wrapper for end_timestep that also ends the timestep in the internal_arc.""" + self.internal_arc.end_timestep() + self.internal_arc.update_queue() + self.storage_ = self.copy_vqip(self.storage) + + def reinit(self): + """Zeros storages and arc.""" + self.internal_arc.reinit() + self.storage = self.empty_vqip() + self.storage_ = self.empty_vqip() + self.active_storage = self.empty_vqip() + + +class DecayQueueTank(QueueTank): + """""" + + def __init__(self, decays={}, parent=None, number_of_timesteps=1, **kwargs): + """Adds a DecayAltArc in QueueTank to enable decay to occur within the + internal_arc queue. + + Args: + decays (dict): A dict of dicts containing a key for each pollutant and, + within that, a key for each parameter (a constant and exponent) + parent (object): An object that can be used to read temperature data from + number_of_timesteps (int, optional): Built in delay for the internal + queue - it is always added to the queue time, although delay can be + provided with pushes only. Defaults to 0. + """ + # Initialise QueueTank + super().__init__(number_of_timesteps=number_of_timesteps, **kwargs) + # Replace internal_arc with a DecayArcAlt + self.internal_arc = DecayArcAlt( + in_port=self, + out_port=self, + number_of_timesteps=number_of_timesteps, + parent=parent, + decays=decays, + ) + + self.end_timestep = self._end_timestep + + def _end_timestep(self): + """End timestep wrapper that removes decayed pollutants and calls internal + arc.""" + # TODO Should the active storage decay if decays are given (probably.. though + # that sounds like a nightmare)? + self.storage = self.extract_vqip(self.storage, self.internal_arc.total_decayed) + self.storage_ = self.copy_vqip(self.storage) + self.internal_arc.end_timestep() diff --git a/wsimod/nodes/wtw.py b/wsimod/nodes/wtw.py index 002313d..4a957c2 100644 --- a/wsimod/nodes/wtw.py +++ b/wsimod/nodes/wtw.py @@ -7,7 +7,8 @@ from typing import Any, Dict from wsimod.core import constants -from wsimod.nodes.nodes import Node, Tank +from wsimod.nodes.nodes import Node +from wsimod.nodes.tanks import Tank class WTW(Node): diff --git a/wsimod/orchestration/model.py b/wsimod/orchestration/model.py index 8e89017..4026a32 100644 --- a/wsimod/orchestration/model.py +++ b/wsimod/orchestration/model.py @@ -19,7 +19,8 @@ from wsimod.core import constants from wsimod.core.core import WSIObj from wsimod.nodes.land import ImperviousSurface -from wsimod.nodes.nodes import NODES_REGISTRY, QueueTank, ResidenceTank, Tank +from wsimod.nodes.nodes import NODES_REGISTRY +from wsimod.nodes.tanks import QueueTank, ResidenceTank, Tank os.environ["USE_PYGEOS"] = "0" @@ -140,6 +141,23 @@ def __init__(self): 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 get_init_args(self, cls): """Get the arguments of the __init__ method for a class and its superclasses.""" init_args = [] @@ -157,6 +175,8 @@ def load(self, address, config_name="config.yml", overrides={}): config_name: overrides: """ + from ..extensions import apply_patches + with open(os.path.join(address, config_name), "r") as file: data = yaml.safe_load(file) @@ -169,6 +189,15 @@ def load(self, address, config_name="config.yml", overrides={}): 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(): @@ -191,6 +220,8 @@ def load(self, address, config_name="config.yml", overrides={}): if "dates" in data.keys(): self.dates = [to_datetime(x) for x in data["dates"]] + apply_patches(self) + def save(self, address, config_name="config.yml", compress=False): """Save the model object to a yaml file and input data to csv.gz format in the directory specified. @@ -284,6 +315,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, @@ -714,55 +746,11 @@ def enablePrint(stdout): node.t = date node.monthyear = date.to_period("M") - # Run FWTW - for node in self.nodes_type.get("FWTW", {}).values(): - node.treat_water() - - # Create demand (gets pushed to sewers) - for node in self.nodes_type.get("Demand", {}).values(): - node.create_demand() - - # Create runoff (impervious gets pushed to sewers, pervious to groundwater) - for node in self.nodes_type.get("Land", {}).values(): - node.run() - - # Infiltrate GW - for node in self.nodes_type.get("Groundwater", {}).values(): - node.infiltrate() - - # Discharge sewers (pushed to other sewers or WWTW) - for node in self.nodes_type.get("Sewer", {}).values(): - node.make_discharge() - - # Foul second so that it can discharge any misconnection - for node in self.nodes_type.get("Foul", {}).values(): - node.make_discharge() - - # Discharge WWTW - for node in self.nodes_type.get("WWTW", {}).values(): - node.calculate_discharge() - - # Discharge GW - for node in self.nodes_type.get("Groundwater", {}).values(): - node.distribute() - - # river - for node in self.nodes_type.get("River", {}).values(): - node.calculate_discharge() - - # Abstract - for node in self.nodes_type.get("Reservoir", {}).values(): - node.make_abstractions() - - for node in self.nodes_type.get("Land", {}).values(): - node.apply_irrigation() - - for node in self.nodes_type.get("WWTW", {}).values(): - node.make_discharge() - - # Catchment routing - for node in self.nodes_type.get("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.get(node_type, {}).values(): + getattr(node, function)() # river for node_name in self.river_discharge_order: