Improve Node representation

python/ribasim/playground.py

@@ -1,21 +1,5 @@
 # %%
-from pydantic import BaseModel
-
-
-class A(BaseModel):
-    x: float
-    y: float
-
-
-class B(BaseModel):
-    a: A
-
-
-b = B(a={"x": 1.2, "y": 1.4})
-
-print(b.model_dump())
-
-# %%
+import geopandas
 import pandas as pd
 from pandas import DataFrame
 from ribasim.add import Basins

python/ribasim/ribasim/add.py

@@ -5,6 +5,7 @@
 import pandas as pd
 from pandas import DataFrame, Timestamp
 from pydantic import ConfigDict, Field
+from shapely.geometry import Point
 
 from ribasim.config import Basin
 from ribasim.geometry.node import NodeSchema
@@ -37,7 +38,7 @@
 
 
 class _BaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
 
 
 class Node(_BaseModel):
@@ -68,27 +69,20 @@ class Basins(_BaseModel):
         default_factory=TableModel[BasinSubgridSchema],
         json_schema_extra={"sort_keys": ["subgrid_id", "basin_level"]},
     )
-    node: SpatialTableModel[NodeSchema] = Field(
-        default_factory=SpatialTableModel[NodeSchema],
-    )
+    node: list[Node] = []
 
     def __init__(self, basins: list[Basin]):
         super().__init__()
         for basin in basins:
             id = basin.node["id"]
 
             basin_dict = basin.model_dump()
+            node = basin_dict.pop("node")
+            self.node.append(Node(**node))
             for basin_key, basin_value in basin_dict.items():
-                if basin_key == "node":
-                    node_dict = {
-                        node_key: [value] for node_key, value in basin_value.items()
-                    }
-                    table_to_append = pd.DataFrame(node_dict)
-                else:
-                    table_to_append = basin_value.assign(node_id=id)
-
-                original_table = getattr(self, basin_key)
-                original_table.df = pd.concat([original_table.df, table_to_append])
+                member = getattr(self, basin_key)
+                table_to_append = basin_value.assign(node_id=id)
+                member.df = pd.concat([member.df, table_to_append])
 
 
 class Model(_BaseModel):