From 467c95b62c8fe70a755926e5bc587eab7d81453c Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Tue, 20 Aug 2024 10:06:54 +0200 Subject: [PATCH] Fix Julia side of things. --- core/src/graph.jl | 6 ++--- core/src/validation.jl | 6 ++--- core/test/run_models_test.jl | 2 +- core/test/validation_test.jl | 4 +-- python/ribasim/ribasim/geometry/edge.py | 33 +++++++++++-------------- python/ribasim/ribasim/geometry/node.py | 7 ++---- python/ribasim/ribasim/input_base.py | 4 +-- utils/templates/schemas.py.jinja | 2 +- 8 files changed, 28 insertions(+), 36 deletions(-) diff --git a/core/src/graph.jl b/core/src/graph.jl index d374b516b..bc0984221 100644 --- a/core/src/graph.jl +++ b/core/src/graph.jl @@ -14,7 +14,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra db, """ SELECT - Edge.fid, + Edge.edge_id, FromNode.node_id AS from_node_id, FromNode.node_type AS from_node_type, ToNode.node_id AS to_node_id, @@ -59,7 +59,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra errors = false for (; - fid, + edge_id, from_node_type, from_node_id, to_node_type, @@ -79,7 +79,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra subnetwork_id = 0 end edge_metadata = EdgeMetadata(; - id = fid, + id = edge_id, flow_idx = edge_type == EdgeType.flow ? flow_counter + 1 : 0, type = edge_type, subnetwork_id_source = subnetwork_id, diff --git a/core/src/validation.jl b/core/src/validation.jl index c5a0e1a0e..25ac02c61 100644 --- a/core/src/validation.jl +++ b/core/src/validation.jl @@ -510,14 +510,14 @@ end function valid_edge_types(db::DB)::Bool edge_rows = execute( db, - "SELECT fid, from_node_id, to_node_id, edge_type FROM Edge ORDER BY fid", + "SELECT edge_id, from_node_id, to_node_id, edge_type FROM Edge ORDER BY edge_id", ) errors = false - for (; fid, from_node_id, to_node_id, edge_type) in edge_rows + for (; edge_id, from_node_id, to_node_id, edge_type) in edge_rows if edge_type ∉ ["flow", "control"] errors = true - @error "Invalid edge type '$edge_type' for edge #$fid from node #$from_node_id to node #$to_node_id." + @error "Invalid edge type '$edge_type' for edge #$edge_id from node #$from_node_id to node #$to_node_id." end end return !errors diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl index 6e911cf02..95d80e26e 100644 --- a/core/test/run_models_test.jl +++ b/core/test/run_models_test.jl @@ -94,7 +94,7 @@ @testset "Results values" begin @test flow.time[1] == DateTime(2020) - @test coalesce.(flow.edge_id[1:2], -1) == [0, 1] + @test coalesce.(flow.edge_id[1:2], -1) == [100, 101] @test flow.from_node_id[1:2] == [6, 0] @test flow.to_node_id[1:2] == [0, 2147483647] diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl index cda3b0a76..e94bbbe87 100644 --- a/core/test/validation_test.jl +++ b/core/test/validation_test.jl @@ -281,10 +281,10 @@ end @test length(logger.logs) == 2 @test logger.logs[1].level == Error @test logger.logs[1].message == - "Invalid edge type 'foo' for edge #0 from node #1 to node #2." + "Invalid edge type 'foo' for edge #1 from node #1 to node #2." @test logger.logs[2].level == Error @test logger.logs[2].message == - "Invalid edge type 'bar' for edge #1 from node #2 to node #3." + "Invalid edge type 'bar' for edge #2 from node #2 to node #3." end @testitem "Subgrid validation" begin diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py index c20e9e6a6..ca4468741 100644 --- a/python/ribasim/ribasim/geometry/edge.py +++ b/python/ribasim/ribasim/geometry/edge.py @@ -14,6 +14,7 @@ from shapely.geometry import LineString, MultiLineString, Point from ribasim.input_base import SpatialTableModel +from ribasim.schemas import _BaseSchema from ribasim.utils import UsedIDs __all__ = ("EdgeTable",) @@ -33,20 +34,15 @@ class NodeData(NamedTuple): geometry: Point -class EdgeSchema(pa.DataFrameModel): - edge_id: Index[Int32] = pa.Field(default=0, ge=0, check_name=True, coerce=True) +class EdgeSchema(_BaseSchema): + edge_id: Index[Int32] = pa.Field(default=0, ge=0, check_name=True) name: Series[str] = pa.Field(default="") - from_node_id: Series[Int32] = pa.Field(default=0, coerce=True) - to_node_id: Series[Int32] = pa.Field(default=0, coerce=True) - edge_type: Series[str] = pa.Field(default="flow", coerce=True) - subnetwork_id: Series[pd.Int32Dtype] = pa.Field( - default=pd.NA, nullable=True, coerce=True - ) + from_node_id: Series[Int32] = pa.Field(default=0) + to_node_id: Series[Int32] = pa.Field(default=0) + edge_type: Series[str] = pa.Field(default="flow") + subnetwork_id: Series[pd.Int32Dtype] = pa.Field(default=pd.NA, nullable=True) geometry: GeoSeries[Any] = pa.Field(default=None, nullable=True) - class Config: - add_missing_columns = True - @classmethod def _index_name(self) -> str: return "edge_id" @@ -101,20 +97,19 @@ def add( f"Edge IDs have to be unique, but {edge_id} already exists." ) - table_to_append = GeoDataFrame( + table_to_append = GeoDataFrame[EdgeSchema]( data={ - "edge_id": pd.Series([edge_id], dtype=np.int32), - "from_node_id": pd.Series([from_node.node_id], dtype=np.int32), - "to_node_id": pd.Series([to_node.node_id], dtype=np.int32), - "edge_type": pd.Series([edge_type], dtype=str), - "name": pd.Series([name], dtype=str), - "subnetwork_id": pd.Series([subnetwork_id], dtype=pd.Int32Dtype()), + "from_node_id": [from_node.node_id], + "to_node_id": [to_node.node_id], + "edge_type": [edge_type], + "name": [name], + "subnetwork_id": [subnetwork_id], **kwargs, }, geometry=geometry_to_append, crs=self.df.crs, + index=pd.Index([edge_id], name="edge_id"), ) - table_to_append.set_index("edge_id", inplace=True) self.df = GeoDataFrame[EdgeSchema](pd.concat([self.df, table_to_append])) if self.df.duplicated(subset=["from_node_id", "to_node_id"]).any(): diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py index faf02d749..04c7e1791 100644 --- a/python/ribasim/ribasim/geometry/node.py +++ b/python/ribasim/ribasim/geometry/node.py @@ -11,11 +11,12 @@ from pandera.typing.geopandas import GeoSeries from ribasim.input_base import SpatialTableModel +from ribasim.schemas import _BaseSchema __all__ = ("NodeTable",) -class NodeSchema(pa.DataFrameModel): +class NodeSchema(_BaseSchema): node_id: Index[Int32] = pa.Field(default=0, check_name=True) name: Series[str] = pa.Field(default="") node_type: Series[str] = pa.Field(default="") @@ -24,10 +25,6 @@ class NodeSchema(pa.DataFrameModel): ) geometry: GeoSeries[Any] = pa.Field(default=None, nullable=True) - class Config: - add_missing_columns = True - coerce = True - @classmethod def _index_name(self) -> str: return "node_id" diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py index 0cb8df99a..66afcf511 100644 --- a/python/ribasim/ribasim/input_base.py +++ b/python/ribasim/ribasim/input_base.py @@ -15,7 +15,6 @@ import geopandas as gpd import numpy as np import pandas as pd -import pandera as pa from pandera.typing import DataFrame from pandera.typing.geopandas import GeoDataFrame from pydantic import BaseModel as PydanticBaseModel @@ -32,6 +31,7 @@ ) import ribasim +from ribasim.schemas import _BaseSchema from .styles import _add_styles_to_geopackage @@ -64,7 +64,7 @@ "file_writing", default={} ) -TableT = TypeVar("TableT", bound=pa.DataFrameModel) +TableT = TypeVar("TableT", bound=_BaseSchema) def esc_id(identifier: str) -> str: diff --git a/utils/templates/schemas.py.jinja b/utils/templates/schemas.py.jinja index 8a9cac921..487d476e8 100644 --- a/utils/templates/schemas.py.jinja +++ b/utils/templates/schemas.py.jinja @@ -2,7 +2,7 @@ import pandera as pa from pandera.dtypes import Int32, Timestamp -from pandera.typing import Series, Index +from pandera.typing import Index, Series class _BaseSchema(pa.DataFrameModel):