Implicit-listen-edge-follow-up (#2969)

fixes #2952

---------

Co-authored-by: Martijn Visser <mgvisser@gmail.com>

Author: verheem

docs/reference/usage.qmd

@@ -321,7 +321,9 @@ There are currently 4 possible link types:
   The LevelDemand and FlowDemand nodes use control links to indicate which nodes it will assign demands to.
 3. "listen": The listen links define which nodes are listened to by control nodes.
   They point from the listened node to the control node.
-  The control node tables define listening behavior (`listen_node_id`, variable, weight, look-ahead), and the Python API automatically adds missing `listen` links when writing a model.
+  The control node tables define listening behavior (`listen_node_id`, variable, weight, look-ahead).
+  Listen links are **automatically created** when a control link is added via `model.link.add`: the Python API reads the `listen_node_id` columns from the control node's tables and adds any missing listen links immediately.
+  If you want custom geometries, listen links can also be added manually with `model.link.add`, where `from_node` is a listenable node (e.g. a Basin) and `to_node` is a control node (e.g. DiscreteControl or PidControl). The link type is automatically inferred as `"listen"` from the node types.
 4. "observation": Observation links point from an Observation node to the node being observed. These links are not used in the simulation core, but allow attaching time series data for reference, validation, or visualization.
 
 column        | type                          | restriction

python/ribasim/ribasim/geometry/link.py

@@ -1,5 +1,6 @@
 import warnings
 from pathlib import Path
+from typing import TYPE_CHECKING, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -26,6 +27,9 @@
 
 from .base import _GeoBaseSchema
 
+if TYPE_CHECKING:
+    from ribasim.model import Model
+
 __all__ = ("LinkTable",)
 
 SPATIALCONTROLNODETYPES = {
@@ -195,6 +199,72 @@ def add(
             )
         self._used_link_ids.add(link_id)
 
+        # When a control link is added from a listen-capable control node,
+        # automatically add listen links based on listen_node_id in its tables.
+        if link_type == "control" and from_node.node_type in LISTENCONTROLNODETYPES:
+            self._add_listen_links_for_control_node(from_node)
+
+    def _collect_listen_node_ids(self, control_node: NodeData) -> set[int]:
+        """Return the set of ``listen_node_id`` values for *control_node*.
+
+        Reads the relevant tables from the parent model based on the control
+        node type.  Returns an empty set when the parent model is unavailable
+        or the tables are empty.
+        """
+        model = cast("Model | None", self._parent)
+        if model is None:
+            return set()
+
+        tables: list[pd.DataFrame | None] = []
+        if control_node.node_type == "PidControl":
+            tables = [model.pid_control.static.df, model.pid_control.time.df]
+        elif control_node.node_type == "DiscreteControl":
+            tables = [model.discrete_control.variable.df]
+        elif control_node.node_type == "ContinuousControl":
+            tables = [model.continuous_control.variable.df]
+
+        listen_node_ids: set[int] = set()
+        for table in tables:
+            if table is None:
+                continue
+            mask = table["node_id"] == control_node.node_id
+            listen_node_ids.update(
+                int(x) for x in table.loc[mask, "listen_node_id"].unique()
+            )
+        return listen_node_ids
+
+    def _add_listen_links_for_control_node(self, control_node: NodeData) -> None:
+        """Add listen links for *control_node* based on its ``listen_node_id`` columns.
+
+        Skips any listen links that already exist in the link table.
+        """
+        model = cast("Model | None", self._parent)
+        if model is None:
+            return
+
+        listen_node_ids = self._collect_listen_node_ids(control_node)
+        if not listen_node_ids:
+            return
+
+        assert self.df is not None
+        assert model.node.df is not None
+        for listen_node_id in sorted(listen_node_ids):
+            already_exists = (
+                (self.df["from_node_id"] == listen_node_id)
+                & (self.df["to_node_id"] == control_node.node_id)
+                & (self.df["link_type"] == "listen")
+            ).any()
+            if already_exists:
+                continue
+
+            row = model.node.df.loc[listen_node_id]
+            listened_node = NodeData(
+                node_id=listen_node_id,
+                node_type=row["node_type"],
+                geometry=row["geometry"],
+            )
+            self.add(listened_node, control_node)
+
     def _remove_link_id(self, link_id: NonNegativeInt):
         if self.df is not None and link_id in self.df.index:
             # Remove from node table

python/ribasim/ribasim/model.py

@@ -400,7 +400,7 @@ def write(
         external : bool, optional
             Write the NetCDF input files. Default is True.
         """
-        self._ensure_listen_links()
+        self.ensure_listen_links()
 
         if self.use_validation:
             self._validate_model()
@@ -455,7 +455,7 @@ def _collect_listen_link_pairs(self) -> pd.DataFrame:
 
         return _concat(parts).drop_duplicates()
 
-    def _ensure_listen_links(self) -> None:
+    def ensure_listen_links(self) -> None:
         """Add any missing listen links inferred from control-node tables.
 
         Compares the pairs from :meth:`_collect_listen_link_pairs` against links
@@ -766,7 +766,7 @@ def plot(
             _, ax = plt.subplots()
             ax.axis("off")
 
-        self._ensure_listen_links()
+        self.ensure_listen_links()
         node = self.node
         self.link.plot(ax=ax, zorder=2)
         node.plot(ax=ax, zorder=3)

python/ribasim/tests/test_link.py

@@ -57,13 +57,16 @@ def test_node_data():
 
 
 def test_listen_link_type_inference(discrete_control_of_pid_control):
+    """Listen links are auto-added when a control link is created."""
     model = discrete_control_of_pid_control
-    model.link.add(model.basin[3], model.pid_control[6])
     assert model.link.df is not None
-    added_link = model.link.df.iloc[-1]
-    assert added_link["from_node_id"] == 3
-    assert added_link["to_node_id"] == 6
-    assert added_link["link_type"] == "listen"
+    listen_links = model.link.df.loc[model.link.df["link_type"] == "listen"]
+    # Basin(3)->PidControl(6) was auto-added when PidControl(6)->Outlet(2) control link was created
+    basin_to_pid = listen_links.loc[
+        (listen_links["from_node_id"] == 3) & (listen_links["to_node_id"] == 6)
+    ]
+    assert len(basin_to_pid) == 1
+    assert basin_to_pid.iloc[0]["link_type"] == "listen"
 
 
 @pytest.mark.parametrize(
@@ -95,12 +98,11 @@ def test_validate_link_rejects_excess_inneighbors(basic):
 
 
 def test_listen_link_allows_reverse_control(discrete_control_of_pid_control):
-    """A listen link should be allowed even when a control link exists in the opposite direction."""
+    """A listen link is auto-added even when a control link exists in the opposite direction."""
     model = discrete_control_of_pid_control
-    # pid_control[6] already has a control link *to* its controlled node;
-    # adding a listen link from basin to pid_control should not trigger the
-    # "opposite link already exists" error.
-    model.link.add(model.basin[3], model.pid_control[6])
+    # pid_control #6 has a control link *to* outlet #2;
+    # the listen link from basin #3 to pid_control #6 was auto-added
+    # without triggering the "opposite link already exists" error.
     assert model.link.df is not None
     listen_links = model.link.df.loc[model.link.df["link_type"] == "listen"]
     assert not listen_links.empty

python/ribasim/tests/test_model.py

@@ -16,7 +16,7 @@
 from ribasim.geometry.node import NodeData
 from ribasim.input_base import esc_id
 from ribasim.model import Model
-from ribasim.nodes import basin
+from ribasim.nodes import basin, pid_control
 from ribasim_testmodels import (
     basic_model,
     outlet_model,
@@ -102,21 +102,65 @@ def test_plot(basic, discrete_control_of_pid_control):
     discrete_control_of_pid_control.plot()
 
 
+def test_control_link_adds_listen_links(discrete_control_of_pid_control):
+    """Adding a control link automatically creates listen links.
+
+    When ``link.add`` creates a control link from a listen-capable control
+    node, it reads the ``listen_node_id`` columns from the control node's
+    tables and adds the corresponding listen links right away.
+    """
+    model = discrete_control_of_pid_control
+
+    assert model.link.df is not None
+    listen_mask = model.link.df["link_type"] == "listen"
+    control_mask = model.link.df["link_type"] == "control"
+
+    # Remove all listen and control links
+    for link_id in model.link.df.index[listen_mask | control_mask].tolist():
+        model.link._remove_link_id(link_id)
+
+    assert model.link.df.loc[model.link.df["link_type"] == "listen"].empty
+    assert model.link.df.loc[model.link.df["link_type"] == "control"].empty
+
+    # Manually add control links back
+    model.link.add(model.pid_control[6], model.outlet[2])
+    model.link.add(model.discrete_control[7], model.pid_control[6])
+
+    # Assert that listen links are now present (auto-added when control links are added)
+    listen_links = model.link.df.loc[model.link.df["link_type"] == "listen"]
+    assert len(listen_links) == 2
+    assert set(
+        zip(listen_links["from_node_id"], listen_links["to_node_id"], strict=False)
+    ) == {
+        (3, 6),
+        (1, 7),
+    }
+
+
 def test_write_adds_missing_listen_links(discrete_control_of_pid_control, tmp_path):
+    """Removing auto-added listen links and writing restores them."""
     model = discrete_control_of_pid_control
 
     assert model.link.df is not None
+    listen_mask = model.link.df["link_type"] == "listen"
+    assert listen_mask.sum() == 2  # auto-added by link.add
+
+    # Remove the auto-added listen links
+    for link_id in model.link.df.index[listen_mask].tolist():
+        model.link._remove_link_id(link_id)
+
     assert model.link.df.loc[model.link.df["link_type"] == "listen"].empty
 
-    model.write(tmp_path / "model_with_listen_links/ribasim.toml")
+    # write() calls ensure_listen_links(), which re-adds them
+    model.write(tmp_path / "restored_listen_links/ribasim.toml")
 
-    with connect(
-        tmp_path / "model_with_listen_links/input/database.gpkg"
-    ) as connection:
+    with connect(tmp_path / "restored_listen_links/input/database.gpkg") as connection:
         query = "select from_node_id, to_node_id from Link where link_type = 'listen'"
         df = pd.read_sql_query(query, connection)
 
-    assert not df.empty
+    assert len(df) == 2
+    written_pairs = set(zip(df["from_node_id"], df["to_node_id"], strict=False))
+    assert written_pairs == {(3, 6), (1, 7)}
 
 
 def test_collect_listen_link_pairs_with_control(discrete_control_of_pid_control):
@@ -134,15 +178,61 @@ def test_collect_listen_link_pairs_without_control(basic):
 
 
 def test_ensure_listen_links_idempotent(discrete_control_of_pid_control):
-    """Calling _ensure_listen_links twice should not duplicate links."""
+    """Calling ensure_listen_links on a model whose listen links are already present.
+
+    (auto-added by link.add) should not duplicate links.
+    """
     model = discrete_control_of_pid_control
-    model._ensure_listen_links()
     assert model.link.df is not None
-    count_after_first = len(model.link.df)
+    count_before = len(model.link.df)
+
+    model.ensure_listen_links()
+    assert len(model.link.df) == count_before
+
+    model.ensure_listen_links()
+    assert len(model.link.df) == count_before
+
+
+def test_manually_add_listen_link(trivial):
+    """Listen links can be added explicitly via link.add().
+
+    When ``link.add`` is called with a listenable node (e.g. Basin) pointing
+    to a control node (e.g. PidControl), the link type is automatically
+    inferred as ``"listen"`` and the link appears in the link DataFrame.
+    This is useful when listen links are needed *before* the control link
+    is added, or for nodes not referenced by ``listen_node_id``.
+    """
+    model = trivial
+
+    # Add a PidControl node that listens to Basin(6)
+    model.pid_control.add(
+        Node(100, Point(400, 300)),
+        [
+            pid_control.Static(
+                listen_node_id=6,
+                target=[0.5],
+                proportional=1e-2,
+                integral=1e-8,
+                derivative=-1e-1,
+            )
+        ],
+    )
+
+    assert model.link.df is not None
+    n_links_before = len(model.link.df)
+
+    # Explicitly add a listen link from Basin(6) -> PidControl(100)
+    model.link.add(model.basin[6], model.pid_control[100])
+
+    listen_links = model.link.df.loc[model.link.df["link_type"] == "listen"]
+    assert len(listen_links) == 1
+    assert listen_links.iloc[0]["from_node_id"] == 6
+    assert listen_links.iloc[0]["to_node_id"] == 100
+    assert len(model.link.df) == n_links_before + 1
 
-    model._ensure_listen_links()
-    count_after_second = len(model.link.df)
-    assert count_after_first == count_after_second
+    # ensure_listen_links does not duplicate the manually added listen link
+    model.ensure_listen_links()
+    assert len(model.link.df.loc[model.link.df["link_type"] == "listen"]) == 1
 
 
 def test_write_adds_fid_in_tables(basic, tmp_path):