From 593047ac29dbbbc15237b0d59ac53a24b2aa2bc6 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Thu, 30 Nov 2023 16:23:36 +0100
Subject: [PATCH] Fix plotting listen edges

---
 python/ribasim/ribasim/model.py    | 37 ++++++++++++++----------------
 python/ribasim/tests/conftest.py   |  5 ++++
 python/ribasim/tests/test_model.py |  4 ++++
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 390908fb0..8eed1f24d 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -404,26 +404,23 @@ def plot_control_listen(self, ax):
                     x_end.append(point_end.x)
                     y_end.append(point_end.y)
 
-        if self.pid_control.static.df is not None:
-            static = self.pid_control.static.df
-            time = self.pid_control.time.df
-            node_static = self.network.node.static.df
-
-            for table in [static, time]:
-                if table is None:
-                    continue
-
-                for node_id in table.node_id.unique():
-                    for listen_node_id in table.loc[
-                        table.node_id == node_id, "listen_node_id"
-                    ].unique():
-                        point_start = node_static.iloc[listen_node_id - 1].geometry
-                        x_start.append(point_start.x)
-                        y_start.append(point_start.y)
-
-                        point_end = node_static.iloc[node_id - 1].geometry
-                        x_end.append(point_end.x)
-                        y_end.append(point_end.y)
+        for table in [self.pid_control.static.df, self.pid_control.time.df]:
+            if table is None:
+                continue
+
+            node = self.network.node.df
+
+            for node_id in table.node_id.unique():
+                for listen_node_id in table.loc[
+                    table.node_id == node_id, "listen_node_id"
+                ].unique():
+                    point_start = node.iloc[listen_node_id - 1].geometry
+                    x_start.append(point_start.x)
+                    y_start.append(point_start.y)
+
+                    point_end = node.iloc[node_id - 1].geometry
+                    x_end.append(point_end.x)
+                    y_end.append(point_end.y)
 
         if len(x_start) == 0:
             return
diff --git a/python/ribasim/tests/conftest.py b/python/ribasim/tests/conftest.py
index c9acf17cf..4e4c1d09f 100644
--- a/python/ribasim/tests/conftest.py
+++ b/python/ribasim/tests/conftest.py
@@ -27,3 +27,8 @@ def tabulated_rating_curve() -> ribasim.Model:
 @pytest.fixture()
 def backwater() -> ribasim.Model:
     return ribasim_testmodels.backwater_model()
+
+
+@pytest.fixture()
+def discrete_control_of_pid_control() -> ribasim.Model:
+    return ribasim_testmodels.discrete_control_of_pid_control_model()
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index 38cf0e78f..0b7069fca 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -131,3 +131,7 @@ def test_tabulated_rating_curve_model(tabulated_rating_curve, tmp_path):
     model_orig = tabulated_rating_curve
     model_orig.write(tmp_path / "tabulated_rating_curve/ribasim.toml")
     Model.read(tmp_path / "tabulated_rating_curve/ribasim.toml")
+
+
+def test_plot(discrete_control_of_pid_control):
+    discrete_control_of_pid_control.plot()