diff --git a/.github/workflows/draft-pdf.yml b/.github/workflows/draft-pdf.yml index a8f2aae1..60fe70ab 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 00000000..1fd9b2af --- /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_demand.py b/tests/test_demand.py index 6032c953..c6d421e7 100644 --- a/tests/test_demand.py +++ b/tests/test_demand.py @@ -119,6 +119,48 @@ def test_garden_demand(self): d1 = {"volume": 0.03 * 0.4, "temperature": 0, "phosphate": 0} self.assertDictAlmostEqual(d1, reply) + def test_demand_overrides(self): + demand = Demand( + name="", + constant_demand=10, + pollutant_load={"phosphate": 0.1, "temperature": 12}, + ) + demand.apply_overrides( + {"constant_demand": 20, "pollutant_load": {"phosphate": 0.5}} + ) + self.assertEqual(demand.constant_demand, 20) + self.assertDictEqual( + demand.pollutant_load, {"phosphate": 0.5, "temperature": 12} + ) + + def test_residentialdemand_overrides(self): + demand = ResidentialDemand( + name="", + gardening_efficiency=0.4, + pollutant_load={"phosphate": 0.1, "temperature": 12}, + ) + demand.apply_overrides( + { + "gardening_efficiency": 0.5, + "population": 153.2, + "per_capita": 32.4, + "constant_weighting": 47.5, + "constant_temp": 0.71, + "constant_demand": 20, + "pollutant_load": {"phosphate": 0.5}, + } + ) + self.assertEqual(demand.gardening_efficiency, 0.5) + self.assertEqual(demand.population, 153.2) + self.assertEqual(demand.per_capita, 32.4) + self.assertEqual(demand.constant_weighting, 47.5) + self.assertEqual(demand.constant_temp, 0.71) + + self.assertEqual(demand.constant_demand, 20) + self.assertDictEqual( + demand.pollutant_load, {"phosphate": 0.5, "temperature": 12} + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_distribution.py b/tests/test_distribution.py index aeaabc74..fca95a2f 100644 --- a/tests/test_distribution.py +++ b/tests/test_distribution.py @@ -67,6 +67,11 @@ def test_leakage(self): self.assertEqual(v2, arc1.vqip_in["volume"]) self.assertEqual(v2 * 0.2, arc2.vqip_in["volume"]) + def test_distribution_overrides(self): + distribution = Distribution(name="", leakage=0.2) + distribution.apply_overrides({"leakage": 0}) + self.assertEqual(distribution.leakage, 0) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_extensions.py b/tests/test_extensions.py new file mode 100644 index 00000000..7f76f681 --- /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 a0042246..3ffdbf7e 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -4,10 +4,10 @@ @author: Barney """ - -# import pytest +import os +import pytest import unittest -from unittest import TestCase +from unittest import TestCase, mock from wsimod.arcs.arcs import Arc from wsimod.nodes.land import Land @@ -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,60 @@ 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) + + +class TestLoadExtensionFiles: + def test_load_extension_files_valid(self, tmp_path_factory): + from wsimod.orchestration.model import load_extension_files + + with tmp_path_factory.mktemp("extensions") as tempdir: + valid_file = os.path.join(tempdir, "valid_extension.py") + with open(valid_file, "w") as f: + f.write("def test_func(): pass") + + load_extension_files([valid_file]) + + def test_load_extension_files_invalid_extension(self, tmp_path_factory): + from wsimod.orchestration.model import load_extension_files + + with tmp_path_factory.mktemp("extensions") as tempdir: + invalid_file = os.path.join(tempdir, "invalid_extension.txt") + with open(invalid_file, "w") as f: + f.write("This is a text file") + + with pytest.raises(ValueError, match="Only .py files are supported"): + load_extension_files([invalid_file]) + + def test_load_extension_files_nonexistent_file(self): + from wsimod.orchestration.model import load_extension_files + + with pytest.raises( + FileNotFoundError, match="File nonexistent_file.py does not exist" + ): + load_extension_files(["nonexistent_file.py"]) + + def test_load_extension_files_import_error(self, tmp_path_factory): + from wsimod.orchestration.model import load_extension_files + + with tmp_path_factory.mktemp("extensions") as tempdir: + valid_file = os.path.join(tempdir, "valid_extension.py") + with open(valid_file, "w") as f: + f.write("raise ImportError") + + with pytest.raises(ImportError): + load_extension_files([valid_file]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 32f73cfc..576ef7b9 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,317 +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) - - def test_overrides(self): - # node - no need to test - # tank - tank = Tank(capacity=10, area=8, datum=4) - tank.apply_overrides({"capacity": 3, "area": 2, "datum": 3.5}) - self.assertEqual(tank.capacity, 3) - self.assertEqual(tank.area, 2) - self.assertEqual(tank.datum, 3.5) - # residence tank - tank = ResidenceTank(capacity=10, area=8, datum=4, residence_time=8) - tank.apply_overrides( - {"capacity": 3, "area": 2, "datum": 3.5, "residence_time": 6} - ) - self.assertEqual(tank.capacity, 3) - self.assertEqual(tank.area, 2) - self.assertEqual(tank.datum, 3.5) - self.assertEqual(tank.residence_time, 6) - # decay tank - tank = DecayTank( - capacity=10, - area=8, - datum=4, - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - ) - tank.apply_overrides( - { - "capacity": 3, - "area": 2, - "datum": 3.5, - "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, - } - ) - self.assertEqual(tank.capacity, 3) - self.assertEqual(tank.area, 2) - self.assertEqual(tank.datum, 3.5) - self.assertDictEqual( - tank.decays, {"phosphate": {"constant": 1.001, "exponent": 10.005}} - ) - # queue tank - tank = QueueTank(capacity=10, area=8, datum=4, number_of_timesteps=8) - tank.apply_overrides( - {"capacity": 3, "area": 2, "datum": 3.5, "number_of_timesteps": 6} - ) - self.assertEqual(tank.capacity, 3) - self.assertEqual(tank.area, 2) - self.assertEqual(tank.datum, 3.5) - self.assertEqual(tank.number_of_timesteps, 6) - self.assertEqual(tank.internal_arc.number_of_timesteps, 6) - # decay queue tank - tank = DecayQueueTank( - capacity=10, - area=8, - datum=4, - number_of_timesteps=8, - decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, - ) - tank.apply_overrides( - { - "capacity": 3, - "area": 2, - "datum": 3.5, - "number_of_timesteps": 6, - "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, - } - ) - self.assertEqual(tank.capacity, 3) - self.assertEqual(tank.area, 2) - self.assertEqual(tank.datum, 3.5) - self.assertEqual(tank.number_of_timesteps, 6) - self.assertEqual(tank.internal_arc.number_of_timesteps, 6) - self.assertDictEqual( - tank.internal_arc.decays, - {"phosphate": {"constant": 1.001, "exponent": 10.005}}, - ) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_sewer.py b/tests/test_sewer.py index 0ca4a289..3299c789 100644 --- a/tests/test_sewer.py +++ b/tests/test_sewer.py @@ -91,6 +91,26 @@ def test_make_discharge(self): d3 = {"volume": 7, "phosphate": 7 * 2 / 9, "temperature": 10} self.assertDictAlmostEqual(d3, sewer.sewer_tank.storage, 15) + def test_sewer_overrides(self): + sewer = Sewer(name="", capacity=10, pipe_timearea={0: 0.3, 1: 0.7}) + sewer.apply_overrides( + { + "capacity": 3, + "chamber_area": 2, + "chamber_floor": 3.5, + "pipe_time": 8.4, + "pipe_timearea": {0: 0.5, 1: 0.5}, + } + ) + self.assertEqual(sewer.capacity, 3) + self.assertEqual(sewer.sewer_tank.capacity, 3) + self.assertEqual(sewer.chamber_area, 2) + self.assertEqual(sewer.sewer_tank.area, 2) + self.assertEqual(sewer.chamber_floor, 3.5) + self.assertEqual(sewer.sewer_tank.datum, 3.5) + self.assertEqual(sewer.pipe_time, 8.4) + self.assertEqual(sewer.pipe_timearea, {0: 0.5, 1: 0.5}) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_tanks.py b/tests/test_tanks.py new file mode 100644 index 00000000..0b1f9456 --- /dev/null +++ b/tests/test_tanks.py @@ -0,0 +1,347 @@ +"""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) + + def test_overrides(self): + # node - no need to test + # tank + tank = Tank(capacity=10, area=8, datum=4) + tank.apply_overrides({"capacity": 3, "area": 2, "datum": 3.5}) + self.assertEqual(tank.capacity, 3) + self.assertEqual(tank.area, 2) + self.assertEqual(tank.datum, 3.5) + # residence tank + tank = ResidenceTank(capacity=10, area=8, datum=4, residence_time=8) + tank.apply_overrides( + {"capacity": 3, "area": 2, "datum": 3.5, "residence_time": 6} + ) + self.assertEqual(tank.capacity, 3) + self.assertEqual(tank.area, 2) + self.assertEqual(tank.datum, 3.5) + self.assertEqual(tank.residence_time, 6) + # decay tank + tank = DecayTank( + capacity=10, + area=8, + datum=4, + decays={"nitrate": {"constant": 0.001, "exponent": 1.005}}, + ) + tank.apply_overrides( + { + "capacity": 3, + "area": 2, + "datum": 3.5, + "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + } + ) + self.assertEqual(tank.capacity, 3) + self.assertEqual(tank.area, 2) + self.assertEqual(tank.datum, 3.5) + self.assertDictEqual( + tank.decays, + { + "nitrate": {"constant": 0.001, "exponent": 1.005}, + "phosphate": {"constant": 1.001, "exponent": 10.005}, + }, + ) + # queue tank + tank = QueueTank(capacity=10, area=8, datum=4, number_of_timesteps=8) + tank.apply_overrides( + {"capacity": 3, "area": 2, "datum": 3.5, "number_of_timesteps": 6} + ) + self.assertEqual(tank.capacity, 3) + self.assertEqual(tank.area, 2) + self.assertEqual(tank.datum, 3.5) + self.assertEqual(tank.number_of_timesteps, 6) + self.assertEqual(tank.internal_arc.number_of_timesteps, 6) + # decay queue tank + tank = DecayQueueTank( + capacity=10, + area=8, + datum=4, + number_of_timesteps=8, + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + ) + tank.apply_overrides( + { + "capacity": 3, + "area": 2, + "datum": 3.5, + "number_of_timesteps": 6, + "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + } + ) + self.assertEqual(tank.capacity, 3) + self.assertEqual(tank.area, 2) + self.assertEqual(tank.datum, 3.5) + self.assertEqual(tank.number_of_timesteps, 6) + self.assertEqual(tank.internal_arc.number_of_timesteps, 6) + self.assertDictEqual( + tank.internal_arc.decays, + {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_wtw.py b/tests/test_wtw.py index 447145e7..44fe4869 100644 --- a/tests/test_wtw.py +++ b/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() diff --git a/wsimod/extensions.py b/wsimod/extensions.py new file mode 100644 index 00000000..6d4a1be4 --- /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/demand.py b/wsimod/nodes/demand.py index 131ac80d..5778186a 100644 --- a/wsimod/nodes/demand.py +++ b/wsimod/nodes/demand.py @@ -5,6 +5,8 @@ Converted to totals BD 2022-05-03 """ +from typing import Any, Dict + from wsimod.core import constants from wsimod.nodes.nodes import Node @@ -28,7 +30,7 @@ def __init__( is used. Defaults to 0. pollutant_load (dict, optional): Pollutant mass per timestep of constant_demand. - Defaults to {}. + Defaults to 0. data_input_dict (dict, optional): Dictionary of data inputs relevant for the node (temperature). Keys are tuples where first value is the name of the variable to read from the dict and the second value is the time. @@ -61,6 +63,19 @@ def __init__( self.mass_balance_out.append(lambda: self.total_backup) self.mass_balance_out.append(lambda: self.total_received) + def apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the sewer. + + Enables a user to override any of the following parameters: + constant_demand, pollutant_load. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.constant_demand = overrides.pop("constant_demand", self.constant_demand) + self.pollutant_load.update(overrides.pop("pollutant_load", {})) + super().apply_overrides(overrides) + def create_demand(self): """Function to call get_demand, which should return a dict with keys that match the keys in directions. @@ -198,6 +213,26 @@ def __init__( # Label as Demand class so that other nodes treat it the same self.__class__.__name__ = "Demand" + def apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the sewer. + + Enables a user to override any of the following parameters: + gardening_efficiency, population, per_capita, constant_weighting, constant_temp. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.gardening_efficiency = overrides.pop( + "gardening_efficiency", self.gardening_efficiency + ) + self.population = overrides.pop("population", self.population) + self.per_capita = overrides.pop("per_capita", self.per_capita) + self.constant_weighting = overrides.pop( + "constant_weighting", self.constant_weighting + ) + self.constant_temp = overrides.pop("constant_temp", self.constant_temp) + super().apply_overrides(overrides) + def get_demand(self): """Overwrite get_demand and replace with custom functions. diff --git a/wsimod/nodes/distribution.py b/wsimod/nodes/distribution.py index 63a7cc0b..66baff45 100644 --- a/wsimod/nodes/distribution.py +++ b/wsimod/nodes/distribution.py @@ -4,6 +4,8 @@ @author: bdobson """ +from typing import Any, Dict + from wsimod.core import constants from wsimod.nodes.nodes import Node @@ -125,8 +127,11 @@ def __init__(self, leakage=0, **kwargs): # Update handlers self.push_set_handler["default"] = self.push_set_deny self.push_check_handler["default"] = self.push_check_deny + self.decorate_pull_handlers() - if leakage > 0: + def decorate_pull_handlers(self): + """Decorate handlers if there is leakage ratio.""" + if self.leakage > 0: self.pull_set_handler["default"] = decorate_leakage_set( self, self.pull_set_handler["default"] ) @@ -134,6 +139,19 @@ def __init__(self, leakage=0, **kwargs): self, self.pull_check_handler["default"] ) + def apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the sewer. + + Enables a user to override any of the following parameters: + leakage. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.leakage = overrides.pop("leakage", self.leakage) + self.decorate_pull_handlers() + super().apply_overrides(overrides) + class UnlimitedDistribution(Distribution): """""" diff --git a/wsimod/nodes/land.py b/wsimod/nodes/land.py index f34fb750..7bcaf86a 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 9f9b6f06..ba7ade47 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,697 +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 apply_overrides(self, overrides: Dict[str, Any] = {}): - """Apply overrides to the tank. - - Enables a user to override any of the following parameters: - area, capacity, datum. - - Args: - overrides (dict, optional): Dictionary of overrides. Defaults to {}. - """ - self.capacity = overrides.pop("capacity", self.capacity) - if "area" in overrides.keys(): - print( - "warning: when overriding area, capacity is not automatically \ - recalculated" - ) - self.area = overrides.pop("area", self.area) - self.datum = overrides.pop("datum", self.datum) - if len(overrides) > 0: - print(f"No override behaviour defined for: {overrides.keys()}") - - 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 apply_overrides(self, overrides: Dict[str, Any] = {}): - """Apply overrides to the residencetank. - - Enables a user to override any of the following parameters: - residence_time. - - Args: - overrides (dict, optional): Dictionary of overrides. Defaults to {}. - """ - self.residence_time = overrides.pop("residence_time", self.residence_time) - super().apply_overrides(overrides) - - 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 apply_overrides(self, overrides: Dict[str, Any] = {}): - """Apply overrides to the decaytank. - - Enables a user to override any of the following parameters: - decays. - - Args: - overrides (dict, optional): Dictionary of overrides. Defaults to {}. - """ - self.decays.update(overrides.pop("decays", {})) - super().apply_overrides(overrides) - - 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 apply_overrides(self, overrides: Dict[str, Any] = {}): - """Apply overrides to the queuetank. - - Enables a user to override any of the following parameters: - number_of_timesteps. - - Args: - overrides (dict, optional): Dictionary of overrides. Defaults to {}. - """ - self.number_of_timesteps = overrides.pop( - "number_of_timesteps", self.number_of_timesteps - ) - self.internal_arc.number_of_timesteps = self.number_of_timesteps - super().apply_overrides(overrides) - - 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 apply_overrides(self, overrides: Dict[str, Any] = {}): - """Apply overrides to the decayqueuetank. - - Enables a user to override any of the following parameters: - number_of_timesteps, decays. - - Args: - overrides (dict, optional): Dictionary of overrides. Defaults to {}. - """ - 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.update(overrides.pop("decays", {})) - super().apply_overrides(overrides) - - 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 af566658..b49546be 100644 --- a/wsimod/nodes/sewer.py +++ b/wsimod/nodes/sewer.py @@ -4,8 +4,11 @@ @author: bdobson Converted to totals on 2022-05-03 """ +from typing import Any, Dict + 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): @@ -112,6 +115,33 @@ def __init__( # Mass balance self.mass_balance_ds.append(lambda: self.sewer_tank.ds()) + def apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the sewer. + + Enables a user to override any of the following parameters: + capacity, chamber_area, chamber_floor, pipe_time, pipe_timearea. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.capacity = overrides.pop("capacity", self.capacity) + self.chamber_area = overrides.pop("chamber_area", self.chamber_area) + self.chamber_floor = overrides.pop("chamber_floor", self.chamber_floor) + self.sewer_tank.capacity = self.capacity + self.sewer_tank.area = self.chamber_area + self.sewer_tank.datum = self.chamber_floor + + self.pipe_time = overrides.pop("pipe_time", self.pipe_time) + if "pipe_timearea" in overrides.keys(): + pipe_timearea_sum = sum([v for k, v in overrides["pipe_timearea"].items()]) + if pipe_timearea_sum != 1: + print( + "ERROR: the sum of pipe_timearea in the overrides dict \ + is not equal to 1, please check it" + ) + self.pipe_timearea = overrides.pop("pipe_timearea", self.pipe_timearea) + super().apply_overrides(overrides) + def push_check_sewer(self, vqip=None): """Generic push check, simply looks at excess. diff --git a/wsimod/nodes/storage.py b/wsimod/nodes/storage.py index d1189f39..546201a4 100644 --- a/wsimod/nodes/storage.py +++ b/wsimod/nodes/storage.py @@ -7,7 +7,8 @@ from typing import Any, Dict 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 00000000..d4a3e93e --- /dev/null +++ b/wsimod/nodes/tanks.py @@ -0,0 +1,696 @@ +"""Module for defining tanks.""" + +from typing import Any, Dict + +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 apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the tank. + + Enables a user to override any of the following parameters: + area, capacity, datum. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.capacity = overrides.pop("capacity", self.capacity) + self.area = overrides.pop("area", self.area) + self.datum = overrides.pop("datum", self.datum) + if len(overrides) > 0: + print(f"No override behaviour defined for: {overrides.keys()}") + + 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 apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the residencetank. + + Enables a user to override any of the following parameters: + residence_time. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.residence_time = overrides.pop("residence_time", self.residence_time) + super().apply_overrides(overrides) + + 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 apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the decaytank. + + Enables a user to override any of the following parameters: + decays. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.decays.update(overrides.pop("decays", {})) + super().apply_overrides(overrides) + + 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 apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the queuetank. + + Enables a user to override any of the following parameters: + number_of_timesteps. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + self.number_of_timesteps = overrides.pop( + "number_of_timesteps", self.number_of_timesteps + ) + self.internal_arc.number_of_timesteps = self.number_of_timesteps + super().apply_overrides(overrides) + + 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 apply_overrides(self, overrides: Dict[str, Any] = {}): + """Apply overrides to the decayqueuetank. + + Enables a user to override any of the following parameters: + number_of_timesteps, decays. + + Args: + overrides (dict, optional): Dictionary of overrides. Defaults to {}. + """ + 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.update(overrides.pop("decays", {})) + super().apply_overrides(overrides) + + 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 e40f0a5d..4a957c2a 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): @@ -542,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.""" diff --git a/wsimod/orchestration/model.py b/wsimod/orchestration/model.py index 8e890175..5b3347f1 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,8 +175,10 @@ 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) + data: dict = yaml.safe_load(file) for key, item in overrides.items(): data[key] = item @@ -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,9 @@ def load(self, address, config_name="config.yml", overrides={}): if "dates" in data.keys(): self.dates = [to_datetime(x) for x in data["dates"]] + load_extension_files(data.get("extensions", [])) + 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 +316,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 +747,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: @@ -1281,3 +1270,27 @@ def yaml2csv(address, config_name="config.yml", csv_folder_name="csv"): writer.writerow( [str(value_[x]) if x in value_.keys() else None for x in fields] ) + + +def load_extension_files(files: list[str]) -> None: + """Load extension files from a list of files. + + Args: + files (list[str]): List of file paths to load + + Raises: + ValueError: If file is not a .py file + FileNotFoundError: If file does not exist + """ + import importlib + from pathlib import Path + + for file in files: + if not file.endswith(".py"): + raise ValueError(f"Only .py files are supported. Invalid file: {file}") + if not Path(file).exists(): + raise FileNotFoundError(f"File {file} does not exist") + + spec = importlib.util.spec_from_file_location("module.name", file) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module)