Fix Julia side of things.

Deltares · Aug 20, 2024 · 467c95b · 467c95b
1 parent 552d7c5
commit 467c95b
Show file tree

Hide file tree

Showing 8 changed files with 28 additions and 36 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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):