Skip to content

Commit

Permalink
Add allocation network plotting
Browse files Browse the repository at this point in the history
  • Loading branch information
SouthEndMusic committed Dec 20, 2023
1 parent c3ef344 commit f7e81a1
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 2 deletions.
44 changes: 44 additions & 0 deletions python/ribasim/ribasim/geometry/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from numpy.typing import NDArray
from pandera.typing import Series
from pandera.typing.geopandas import GeoSeries
from scipy.spatial import ConvexHull
from shapely.geometry import Point

from ribasim.input_base import SpatialTableModel
Expand Down Expand Up @@ -158,6 +159,49 @@ def connectivity_from_geometry(
to_id = node_index[edge_node_id[:, 1]].to_numpy()
return from_id, to_id

def plot_allocation_networks(self, ax=None, zorder=None) -> Any:
if ax is None:
_, ax = plt.subplots()
ax.axis("off")

for allocation_subnetwork_id, df_subnetwork in self.df.groupby(
"allocation_network_id"
):
points = np.stack(
(
df_subnetwork.geometry.x.to_numpy(),
df_subnetwork.geometry.y.to_numpy(),
),
axis=1,
)
hull = ConvexHull(points)
bounding_polygon_vertices = points[hull.vertices, :]
hull_center = np.mean(bounding_polygon_vertices, axis=0)
bounding_polygon_vertices = (
bounding_polygon_vertices - hull_center
) * 1.1 + hull_center

ax.fill(
bounding_polygon_vertices[:, 0],
bounding_polygon_vertices[:, 1],
linestyle=":",
facecolor="none",
linewidth=2,
edgecolor="gray",
zorder=zorder,
)
if allocation_subnetwork_id == 1:
text = "Main network"
else:
text = f"Subnetwork {allocation_subnetwork_id}"
ax.text(
*hull_center,
text,
horizontalalignment="center",
color="gray",
zorder=zorder,
)

def plot(self, ax=None, zorder=None) -> Any:
"""
Plot the nodes. Each node type is given a separate marker.
Expand Down
6 changes: 4 additions & 2 deletions python/ribasim/ribasim/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -505,9 +505,9 @@ def plot_control_listen(self, ax):
label="Listen edge" if i == 0 else None,
)

def plot(self, ax=None) -> Any:
def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
"""
Plot the nodes and edges of the model.
Plot the nodes, edges and allocation networks of the model.
Parameters
----------
Expand All @@ -524,6 +524,8 @@ def plot(self, ax=None) -> Any:
self.network.edge.plot(ax=ax, zorder=2)
self.plot_control_listen(ax)
self.network.node.plot(ax=ax, zorder=3)
if indicate_subnetworks:
self.network.node.plot_allocation_networks(ax=ax)

ax.legend(loc="lower left", bbox_to_anchor=(1, 0.5))

Expand Down

0 comments on commit f7e81a1

Please sign in to comment.