From 6049ae3939b12d847e5cb508ddcfdbb1bd993d59 Mon Sep 17 00:00:00 2001 From: Bart de Koning Date: Fri, 22 Dec 2023 17:00:53 +0100 Subject: [PATCH] Add docstrings --- python/ribasim/ribasim/geometry/edge.py | 2 ++ python/ribasim/ribasim/geometry/node.py | 2 ++ python/ribasim/ribasim/input_base.py | 5 +++++ python/ribasim/ribasim/model.py | 8 ++++++++ python/ribasim/tests/test_model.py | 1 + 5 files changed, 18 insertions(+) diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py index 83365ab36..8d17f53c4 100644 --- a/python/ribasim/ribasim/geometry/edge.py +++ b/python/ribasim/ribasim/geometry/edge.py @@ -43,6 +43,7 @@ class Edge(SpatialTableModel[EdgeSchema]): def translate_spacially( self, offset_spacial: tuple[float, float], inplace: bool = True ) -> "Edge": + """Add the same spacial offset to all edges.""" if inplace: edge = self else: @@ -61,6 +62,7 @@ def translate_spacially( def offset_allocation_network_ids( self, offset_allocation_network_id: int, inplace: bool = True ) -> "Edge": + """Add the same offset to all node IDs.""" if inplace: edge = self else: diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py index e5c17a4ff..a6ad869a6 100644 --- a/python/ribasim/ribasim/geometry/node.py +++ b/python/ribasim/ribasim/geometry/node.py @@ -66,6 +66,7 @@ def node_ids_and_types(*nodes): def translate_spacially( self, offset_spacial: tuple[float, float], inplace: bool = True ) -> "Node": + """Add the same spacial offset to all nodes.""" if inplace: node = self else: @@ -81,6 +82,7 @@ def translate_spacially( def offset_allocation_network_ids( self, offset_allocation_network_id: int, inplace: bool = True ) -> "Node": + """Add the same offset to all node IDs.""" if inplace: node = self else: diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py index e0a89d1b2..474ab9946 100644 --- a/python/ribasim/ribasim/input_base.py +++ b/python/ribasim/ribasim/input_base.py @@ -217,6 +217,7 @@ def node_ids(self) -> set[int]: return node_ids def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]": + """Add the same offset to all node IDs.""" copy = deepcopy(self) df = copy.df if isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)): @@ -234,6 +235,7 @@ def offset_node_ids(self, offset_node_id: int) -> "TableModel[TableT]": def merge_table( self, table_added: "TableModel[TableT]", inplace: bool = True ) -> "TableModel[TableT]": + """Merge an added table of the same type into this table.""" assert type(self) == type( table_added ), "Can only merge tables of the same type." @@ -472,6 +474,7 @@ def node_ids_and_types(self) -> tuple[list[int], list[str]]: return list(ids), len(ids) * [self.get_input_type()] def offset_node_ids(self, offset_node_id: int) -> "NodeModel": + """Add the same offset to all node IDs in all underlying tables.""" node_copy = deepcopy(self) for field in node_copy.fields(): attr = getattr(node_copy, field) @@ -485,6 +488,7 @@ def offset_node_ids(self, offset_node_id: int) -> "NodeModel": return node_copy def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeModel": + """Merge an added node of the same type into this node.""" assert type(self) == type(node_added), "Can only merge nodes of the same type." if inplace: @@ -505,6 +509,7 @@ def merge_node(self, node_added: "NodeModel", inplace: bool = True) -> "NodeMode return node def delete_by_ids(self, node_ids: list[int], inplace: bool = True) -> "NodeModel": + """Delete all rows of the underlying tables whose node ID is in the given list.""" if inplace: node = self else: diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py index 531ec22d9..e3c0eccbc 100644 --- a/python/ribasim/ribasim/model.py +++ b/python/ribasim/ribasim/model.py @@ -64,6 +64,7 @@ def n_nodes(self): def translate_spacially( self, offset_spacial: tuple[float, float], inplace: bool = True ) -> "Network": + """Add the same spacial offset to all nodes and edges.""" if inplace: network = self else: @@ -92,6 +93,7 @@ def add_edges( edge_type: list[str], inplace: bool = True, ) -> "Network": + """Add new edges to the network of the given type. Assumes no source edges are added.""" if inplace: network = self else: @@ -463,6 +465,12 @@ def merge_model( offset_spacial: tuple[float, float] = (0.0, 0.0), inplace: bool = True, ): + """ + Merge copies of the nodes and edges of an added model into this model. + The added model is not modified, but the following modificadions are made to the added data: + - Node IDs are shifted by at least the maximum node ID of this model + - Allocation network IDs are shifted by at least the maximum allocation network ID of this model. + """ if inplace: model = self else: diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py index 36998dc6f..9ecb2fbfb 100644 --- a/python/ribasim/tests/test_model.py +++ b/python/ribasim/tests/test_model.py @@ -154,6 +154,7 @@ def test_model_merging(basic, subnetwork, tmp_path): model.merge_model(model_added) assert (model.network.node.df.index == range(1, 44)).all() assert model.max_allocation_network_id() == 2 + # Added model should not change for node_type, node_added in model_added.nodes().items(): node_subnetwork = getattr(subnetwork, node_type) for table_added, table_subnetwork in zip(