Skip to content

Commit

Permalink
Fix plotting listen edges
Browse files Browse the repository at this point in the history
  • Loading branch information
visr committed Nov 30, 2023
1 parent 958626e commit 593047a
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 20 deletions.
37 changes: 17 additions & 20 deletions python/ribasim/ribasim/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions python/ribasim/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
4 changes: 4 additions & 0 deletions python/ribasim/tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit 593047a

Please sign in to comment.