Merge branch 'main' into demand-overrides

.github/workflows/draft-pdf.yml

@@ -6,15 +6,15 @@ 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:
  journal: joss
  # 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

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'

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"

tests/test_model.py

@@ -15,6 +15,7 @@
 from wsimod.nodes.sewer import Sewer
 from wsimod.nodes.waste import Waste
 from wsimod.orchestration.model import Model, to_datetime
+import os
 
 
 class MyTestClass(TestCase):
@@ -291,6 +292,18 @@ def test_run(self):
  0.03, my_model.nodes["my_land"].get_surface("urban").storage["volume"]
  )
 
+ def test_customise_orchestration(self):
+ my_model = Model()
+ my_model.load(
+ os.path.join(os.getcwd(), "docs", "demo", "examples"),
+ config_name="test_customise_orchestration_example.yaml",
+ )
+ revised_orchestration = [
+ {"Groundwater": "infiltrate"},
+ {"Sewer": "make_discharge"},
+ ]
+ self.assertListEqual(my_model.orchestration, revised_orchestration)
+
 
 if __name__ == "__main__":
  unittest.main()