diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py index 7148e9ee1..3abd3216e 100644 --- a/python/ribasim/ribasim/geometry/edge.py +++ b/python/ribasim/ribasim/geometry/edge.py @@ -40,12 +40,14 @@ class Edge(SpatialTableModel[EdgeSchema]): """ def get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]: + assert self.df is not None return (self.df.edge_type == edge_type).to_numpy() def plot(self, **kwargs) -> Axes: ax = kwargs.get("ax", None) color_flow = kwargs.get("color_flow", None) color_control = kwargs.get("color_control", None) + assert self.df is not None if ax is None: _, ax = plt.subplots() diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py index f10a2bc22..5c3752b00 100644 --- a/python/ribasim/ribasim/geometry/node.py +++ b/python/ribasim/ribasim/geometry/node.py @@ -80,6 +80,7 @@ def geometry_from_connectivity( edge_geometry : np.ndarray Array of shapely LineStrings. """ + assert self.df is not None geometry = self.df["geometry"] from_points = shapely.get_coordinates(geometry.loc[from_id]) to_points = shapely.get_coordinates(geometry.loc[to_id]) @@ -109,6 +110,7 @@ def connectivity_from_geometry( from_node_id : np.ndarray of int to_node_id : np.ndarray of int """ + assert self.df is not None node_index = self.df.index node_xy = shapely.get_coordinates(self.df.geometry.values) edge_xy = shapely.get_coordinates(lines) @@ -143,6 +145,7 @@ def plot_allocation_networks(self, ax=None, zorder=None) -> Any: contains_main_network = False contains_subnetworks = False + assert self.df is not None for allocation_subnetwork_id, df_subnetwork in self.df.groupby( "allocation_network_id" @@ -223,6 +226,7 @@ def plot(self, ax=None, zorder=None) -> Any: "User": "g", "": "k", } + assert self.df is not None for nodetype, df in self.df.groupby("type"): assert isinstance(nodetype, str) @@ -237,6 +241,7 @@ def plot(self, ax=None, zorder=None) -> Any: label=nodetype, ) + assert self.df is not None geometry = self.df["geometry"] for text, xy in zip(self.df.index, np.column_stack((geometry.x, geometry.y))): ax.annotate(text=text, xy=xy, xytext=(2.0, 2.0), textcoords="offset points") diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py index 13b4509df..05057a58c 100644 --- a/python/ribasim/ribasim/model.py +++ b/python/ribasim/ribasim/model.py @@ -468,6 +468,8 @@ def print_discrete_control_record(self, path: FilePath) -> None: node_attrs, node_instances = zip(*self.nodes().items()) node_clss = [node_cls.get_input_type() for node_cls in node_instances] truth_dict = {"T": ">", "F": "<"} + assert self.network.node.df is not None + assert self.network.edge.df is not None if self.discrete_control.condition.df is None: raise ValueError("This model has no control input.")