diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py index 4c61a26be..21d359276 100644 --- a/python/ribasim/ribasim/input_base.py +++ b/python/ribasim/ribasim/input_base.py @@ -13,7 +13,6 @@ ) import geopandas as gpd -import numpy as np import pandas as pd import pandera as pa from pandera.typing import DataFrame @@ -204,10 +203,10 @@ def node_ids(self) -> set[int]: return node_ids - def offset_node_ids(self, offset_node_id: int) -> "TableModel": + def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]": copy = deepcopy(self) df = copy.df - if copy.df is not None: + if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)): df.index += offset_node_id for name_column in [ "node_id", @@ -220,8 +219,8 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel": return copy def merge_table( - self, table_added: "TableModel", inplace: bool = True - ) -> "TableModel": + self, table_added: "TableModel[TableT]", inplace: bool = True + ) -> "TableModel[TableT]": assert type(self) == type( table_added ), "Can only merge tables of the same type." @@ -492,9 +491,7 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode setattr(node, field, table_added) return node - def delete_by_ids( - self, node_ids: np.ndarray[int], inplace: bool = True - ) -> "NodeModel": + def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel": if inplace: node = self else: @@ -503,11 +500,13 @@ def delete_by_ids( for field in node.fields(): attr = getattr(node, field) if isinstance(attr, TableModel): - df = attr.df[~attr.df.node_id.isin(node_ids)] - if df.empty: - attr.df = None - else: - attr.df = df + df = attr.df + if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)): + df = df[~df.node_id.isin(node_ids)] + if df.empty: + attr.df = None + else: + attr.df = df return node diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py index 7b8fa1e60..1d0699a4d 100644 --- a/python/ribasim/ribasim/model.py +++ b/python/ribasim/ribasim/model.py @@ -87,8 +87,8 @@ def offset_allocation_network_ids( def add_edges( self, - from_node_id: np.ndarray[int], - to_node_id: np.ndarray[int], + from_node_id: list[int], + to_node_id: list[int], edge_type: list[str], inplace: bool = True, ) -> "Network":