diff --git a/.gitignore b/.gitignore index 7b5586e8c..8ee64eb35 100644 --- a/.gitignore +++ b/.gitignore @@ -154,3 +154,4 @@ report.xml # Designated working dir for working on Ribasim models models/ +playground/output diff --git a/playground/playground.py b/playground/playground.py new file mode 100644 index 000000000..99085e2ff --- /dev/null +++ b/playground/playground.py @@ -0,0 +1,50 @@ +# %% +import shutil +from pathlib import Path + +from pandas import DataFrame +from ribasim.add import Basins, Model +from ribasim.config import Basin + +# %% +output_dir = Path(__file__).parent / "output" +shutil.rmtree(output_dir, ignore_errors=True) +# %% + +basins = Basins( + [ + Basin( + profile=DataFrame({"area": [1.0, 3.0], "level": [1.1, 2.2]}), + # static=DataFrame( + # {"precipitation": [1.0, 3.0], "control_state": [1.1, 2.2]} + # ), + node={ + "node_id": 2, + "geometry": (2.0, 3.6), + "name": "IJsselmeer", + "subnetwork_id": 5, + }, + ), + Basin( + profile=DataFrame({"area": [2.0, 4.0], "level": [6.1, 7.2]}), + # static=DataFrame( + # {"precipitation": [1.0, 3.0], "control_state": [1.1, 2.2]} + # ), + node={ + "node_id": 1, + "geometry": (5.0, 3.6), + "name": "Zoetermeer", + "subnetwork_id": 6, + }, + ), + ] +) +basins +# %% +model = Model( + starttime="2020-01-01 00:00:00", endtime="2021-01-01 00:00:00", basins=basins +) +model +# %% +toml_path = output_dir / "ribasim.toml" +model.write(toml_path) diff --git a/python/ribasim/playground.py b/python/ribasim/playground.py new file mode 100644 index 000000000..9797182ce --- /dev/null +++ b/python/ribasim/playground.py @@ -0,0 +1,38 @@ +# %% +import tempfile +from pathlib import Path + +import geopandas +import pandas as pd +from pandas import DataFrame +from ribasim.add import Basins, Model +from ribasim.config import Basin + +basins = Basins( + [ + Basin( + profile=DataFrame({"area": [1.0, 3.0], "level": [1.1, 2.2]}), + # static=DataFrame( + # {"precipitation": [1.0, 3.0], "control_state": [1.1, 2.2]} + # ), + node={ + "id": 2, + # allow tuples, call gpd.points_from_xy ourselves + "geometry": (2.0, 3.6), + "name": "IJsselmeer", + "allocation_network_id": 5, + }, + ) + ] +) +basins +# %% +model = Model( + starttime="2020-01-01 00:00:00", endtime="2021-01-01 00:00:00", basins=basins +) +model +# %% +with tempfile.TemporaryDirectory() as tmpdirname: + toml_path = Path(tmpdirname) / "ribasim.toml" + model.write(toml_path) + print(toml_path.read_text()) diff --git a/python/ribasim/ribasim/add.py b/python/ribasim/ribasim/add.py new file mode 100644 index 000000000..d97570aa2 --- /dev/null +++ b/python/ribasim/ribasim/add.py @@ -0,0 +1,311 @@ +import datetime +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Self + +import pandas as pd +import tomli +import tomli_w +from geopandas import GeoDataFrame +from pydantic import ( + DirectoryPath, + Field, + FilePath, + field_serializer, + model_serializer, + model_validator, +) +from shapely.geometry import Point + +from ribasim import Node +from ribasim.config import SingleNodeModel +from ribasim.geometry.edge import Edge +from ribasim.input_base import ( + BaseModel, + ChildModel, + FileModel, + NodeModel, + TableModel, + context_file_loading, +) +from ribasim.schemas import ( + BasinProfileSchema, + BasinStateSchema, + BasinStaticSchema, + BasinSubgridSchema, + BasinTimeSchema, +) + + +class NodeData(BaseModel): + node_id: int + node_type: str + geometry: tuple[float, float] + name: str = "" + subnetwork_id: int | None = None + + def into_geodataframe(self) -> GeoDataFrame: + return GeoDataFrame( + data={ + "node_id": [self.node_id], + "node_type": [self.node_type], + "name": [self.name], + "subnetwork_id": [self.subnetwork_id], + }, + geometry=[Point(*self.geometry)], + ) + + +class MultiNodeModel(NodeModel): + node: Node = Field(default_factory=Node) + + def __init__(self, nodes: Sequence[SingleNodeModel] | None = None) -> None: + if nodes is None: + super().__init__() + return + + data = {} + node_ids = set() + for node in nodes: + node_id = node.node["node_id"] + if node_id in node_ids: + raise ValueError( + f"Node IDs have to be unique, but {node_id=} exists multiple times." + ) + else: + node_ids.add(node_id) + for field in node.fields(): + existing_table = data[field] if field in data else pd.DataFrame() + if field == "node": + table_to_append = NodeData( + node_type=node.__class__.__name__, **node.node + ).into_geodataframe() + else: + table_to_append = getattr(node, field).df + if table_to_append is None: + continue + table_to_append = table_to_append.assign(node_id=node_id) + data[field] = pd.concat([existing_table, table_to_append]) + super().__init__(**data) + + +class Basins(MultiNodeModel): + profile: TableModel[BasinProfileSchema] = Field( + default_factory=TableModel[BasinProfileSchema], + json_schema_extra={"sort_keys": ["node_id", "level"]}, + ) + state: TableModel[BasinStateSchema] = Field( + default_factory=TableModel[BasinStateSchema], + json_schema_extra={"sort_keys": ["node_id"]}, + ) + static: TableModel[BasinStaticSchema] = Field( + default_factory=TableModel[BasinStaticSchema], + json_schema_extra={"sort_keys": ["node_id"]}, + ) + time: TableModel[BasinTimeSchema] = Field( + default_factory=TableModel[BasinTimeSchema], + json_schema_extra={"sort_keys": ["node_id", "time"]}, + ) + subgrid: TableModel[BasinSubgridSchema] = Field( + default_factory=TableModel[BasinSubgridSchema], + json_schema_extra={"sort_keys": ["subgrid_id", "basin_level"]}, + ) + + def __init__(self, nodes: Sequence[SingleNodeModel] | None = None) -> None: + super().__init__(nodes) + + +# TODO: This only deals with edges so should probably be refactored +class Network(FileModel, NodeModel): + filepath: Path | None = Field( + default=Path("database.gpkg"), exclude=True, repr=False + ) + + edge: Edge = Field(default_factory=Edge) + + @classmethod + def _load(cls, filepath: Path | None) -> dict[str, Any]: + directory = context_file_loading.get().get("directory", None) + if directory is not None: + context_file_loading.get()["database"] = directory / "database.gpkg" + return {} + + @classmethod + def _layername(cls, field: str) -> str: + return field.capitalize() + + def _save(self, directory, input_dir=Path(".")): + # We write all tables to a temporary database with a dot prefix, + # and at the end move this over the target file. + # This does not throw a PermissionError if the file is open in QGIS. + directory = Path(directory) + db_path = directory / input_dir / "database.gpkg" + db_path = db_path.resolve() + db_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = db_path.with_stem(".database") + + # avoid adding tables to existing model + temp_path.unlink(missing_ok=True) + context_file_loading.get()["database"] = temp_path + + # TODO: save edges + # self.edge._save(directory, input_dir) + # shutil.move(temp_path, db_path) + context_file_loading.get()["database"] = db_path + + @model_serializer + def set_modelname(self) -> str: + if self.filepath is not None: + return str(self.filepath.name) + else: + return str(self.model_fields["filepath"].default) + + +class Model(FileModel): + starttime: datetime.datetime + endtime: datetime.datetime + + input_dir: Path = Field(default_factory=lambda: Path(".")) + results_dir: Path = Field(default_factory=lambda: Path("results")) + + basins: Basins = Field(default_factory=Basins) + # TODO: refactor Network + network: Network = Field(default_factory=Network, alias="database", exclude=True) + + @model_validator(mode="after") + def set_node_parent(self) -> Self: + for ( + k, + v, + ) in self.children().items(): + setattr(v, "_parent", self) + setattr(v, "_parent_field", k) + return self + + @field_serializer("input_dir", "results_dir") + def serialize_path(self, path: Path) -> str: + return str(path) + + def model_post_init(self, __context: Any) -> None: + # Always write dir fields + self.model_fields_set.update({"input_dir", "results_dir"}) + + def __repr__(self) -> str: + """Generate a succinct overview of the Model content. + + Skip "empty" NodeModel instances: when all dataframes are None. + """ + content = ["ribasim.Model("] + INDENT = " " + for field in self.fields(): + attr = getattr(self, field) + content.append(f"{INDENT}{field}={repr(attr)},") + + content.append(")") + return "\n".join(content) + + def _write_toml(self, fn: FilePath): + fn = Path(fn) + + content = self.model_dump(exclude_unset=True, exclude_none=True, by_alias=True) + # Filter empty dicts (default Nodes) + content = dict(filter(lambda x: x[1], content.items())) + with open(fn, "wb") as f: + tomli_w.dump(content, f) + return fn + + def _save(self, directory: DirectoryPath, input_dir: DirectoryPath): + self.network._save(directory, input_dir) + for sub in self.nodes().values(): + sub._save(directory, input_dir) + + def nodes(self): + return { + k: getattr(self, k) + for k in self.model_fields.keys() + if isinstance(getattr(self, k), NodeModel) + } + + def children(self): + return { + k: getattr(self, k) + for k in self.model_fields.keys() + if isinstance(getattr(self, k), ChildModel) + } + + def validate_model_node_field_ids(self): + raise NotImplementedError() + + def validate_model_node_ids(self): + raise NotImplementedError() + + def validate_model(self): + """Validate the model. + + Checks: + - Whether the node IDs of the node_type fields are valid + - Whether the node IDs in the node field correspond to the node IDs on the node type fields + """ + + self.validate_model_node_field_ids() + self.validate_model_node_ids() + + @classmethod + def read(cls, filepath: FilePath) -> "Model": + """Read model from TOML file.""" + return cls(filepath=filepath) # type: ignore + + def write(self, filepath: Path | str) -> Path: + """ + Write the contents of the model to disk and save it as a TOML configuration file. + + If ``filepath.parent`` does not exist, it is created before writing. + + Parameters + ---------- + filepath: FilePath ending in .toml + """ + # TODO + # self.validate_model() + filepath = Path(filepath) + if not filepath.suffix == ".toml": + raise ValueError(f"Filepath '{filepath}' is not a .toml file.") + context_file_loading.set({}) + filepath = Path(filepath) + directory = filepath.parent + directory.mkdir(parents=True, exist_ok=True) + self._save(directory, self.input_dir) + fn = self._write_toml(filepath) + + context_file_loading.set({}) + return fn + + @classmethod + def _load(cls, filepath: Path | None) -> dict[str, Any]: + context_file_loading.set({}) + + if filepath is not None: + with open(filepath, "rb") as f: + config = tomli.load(f) + + context_file_loading.get()["directory"] = filepath.parent / config.get( + "input_dir", "." + ) + return config + else: + return {} + + @model_validator(mode="after") + def reset_contextvar(self) -> Self: + # Drop database info + context_file_loading.set({}) + return self + + def plot_control_listen(self, ax): + raise NotImplementedError() + + def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any: + raise NotImplementedError() + + def print_discrete_control_record(self, path: FilePath) -> None: + raise NotImplementedError() diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py index d6c3005a0..f3569e18b 100644 --- a/python/ribasim/ribasim/config.py +++ b/python/ribasim/ribasim/config.py @@ -1,4 +1,5 @@ from enum import Enum +from typing import Any from pydantic import Field @@ -73,14 +74,18 @@ class Logging(ChildModel): timing: bool = False -class Terminal(NodeModel): +class SingleNodeModel(NodeModel): + node: dict[str, Any] + + +class Terminal(SingleNodeModel): static: TableModel[TerminalStaticSchema] = Field( default_factory=TableModel[TerminalStaticSchema], json_schema_extra={"sort_keys": ["node_id"]}, ) -class PidControl(NodeModel): +class PidControl(SingleNodeModel): static: TableModel[PidControlStaticSchema] = Field( default_factory=TableModel[PidControlStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, @@ -91,7 +96,7 @@ class PidControl(NodeModel): ) -class LevelBoundary(NodeModel): +class LevelBoundary(SingleNodeModel): static: TableModel[LevelBoundaryStaticSchema] = Field( default_factory=TableModel[LevelBoundaryStaticSchema], json_schema_extra={"sort_keys": ["node_id"]}, @@ -102,14 +107,14 @@ class LevelBoundary(NodeModel): ) -class Pump(NodeModel): +class Pump(SingleNodeModel): static: TableModel[PumpStaticSchema] = Field( default_factory=TableModel[PumpStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, ) -class TabulatedRatingCurve(NodeModel): +class TabulatedRatingCurve(SingleNodeModel): static: TableModel[TabulatedRatingCurveStaticSchema] = Field( default_factory=TableModel[TabulatedRatingCurveStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state", "level"]}, @@ -120,7 +125,7 @@ class TabulatedRatingCurve(NodeModel): ) -class UserDemand(NodeModel): +class UserDemand(SingleNodeModel): static: TableModel[UserDemandStaticSchema] = Field( default_factory=TableModel[UserDemandStaticSchema], json_schema_extra={"sort_keys": ["node_id", "priority"]}, @@ -131,7 +136,7 @@ class UserDemand(NodeModel): ) -class LevelDemand(NodeModel): +class LevelDemand(SingleNodeModel): static: TableModel[LevelDemandStaticSchema] = Field( default_factory=TableModel[LevelDemandStaticSchema], json_schema_extra={"sort_keys": ["node_id", "priority"]}, @@ -142,7 +147,7 @@ class LevelDemand(NodeModel): ) -class FlowBoundary(NodeModel): +class FlowBoundary(SingleNodeModel): static: TableModel[FlowBoundaryStaticSchema] = Field( default_factory=TableModel[FlowBoundaryStaticSchema], json_schema_extra={"sort_keys": ["node_id"]}, @@ -153,7 +158,7 @@ class FlowBoundary(NodeModel): ) -class Basin(NodeModel): +class Basin(SingleNodeModel): profile: TableModel[BasinProfileSchema] = Field( default_factory=TableModel[BasinProfileSchema], json_schema_extra={"sort_keys": ["node_id", "level"]}, @@ -176,14 +181,14 @@ class Basin(NodeModel): ) -class ManningResistance(NodeModel): +class ManningResistance(SingleNodeModel): static: TableModel[ManningResistanceStaticSchema] = Field( default_factory=TableModel[ManningResistanceStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, ) -class DiscreteControl(NodeModel): +class DiscreteControl(SingleNodeModel): condition: TableModel[DiscreteControlConditionSchema] = Field( default_factory=TableModel[DiscreteControlConditionSchema], json_schema_extra={ @@ -196,21 +201,21 @@ class DiscreteControl(NodeModel): ) -class Outlet(NodeModel): +class Outlet(SingleNodeModel): static: TableModel[OutletStaticSchema] = Field( default_factory=TableModel[OutletStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, ) -class LinearResistance(NodeModel): +class LinearResistance(SingleNodeModel): static: TableModel[LinearResistanceStaticSchema] = Field( default_factory=TableModel[LinearResistanceStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, ) -class FractionalFlow(NodeModel): +class FractionalFlow(SingleNodeModel): static: TableModel[FractionalFlowStaticSchema] = Field( default_factory=TableModel[FractionalFlowStaticSchema], json_schema_extra={"sort_keys": ["node_id", "control_state"]}, diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py index 03a48e723..2e97894ba 100644 --- a/python/ribasim/ribasim/input_base.py +++ b/python/ribasim/ribasim/input_base.py @@ -354,9 +354,7 @@ def _write_table(self, path: FilePath) -> None: gdf = gpd.GeoDataFrame(data=self.df) gdf = gdf.set_geometry("geometry") - gdf.index.name = "fid" - - gdf.to_file(path, layer=self.tablename(), driver="GPKG", index=True) + gdf.to_file(path, layer=self.tablename(), driver="GPKG") def sort(self): self.df.sort_index(inplace=True) @@ -425,11 +423,15 @@ def node_ids_and_types(self) -> tuple[list[int], list[str]]: return list(ids), len(ids) * [self.get_input_type()] def _save(self, directory: DirectoryPath, input_dir: DirectoryPath, **kwargs): - for field in self.fields(): - getattr(self, field)._save( - directory, - input_dir, - ) + # TODO: stop sorting loop so that "node" comes first + for field in sorted(self.fields(), key=lambda x: x != "node"): + attr = getattr(self, field) + # TODO + if hasattr(attr, "_save"): + attr._save( + directory, + input_dir, + ) def _repr_content(self) -> str: """Generate a succinct overview of the content. diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py index 52120e5d7..6785bb3c7 100644 --- a/python/ribasim/ribasim/schemas.py +++ b/python/ribasim/ribasim/schemas.py @@ -12,18 +12,18 @@ class Config: class BasinProfileSchema(_BaseSchema): - node_id: Series[int] = pa.Field(nullable=False) + node_id: Series[int] = pa.Field(nullable=False, default=0) area: Series[float] = pa.Field(nullable=False) level: Series[float] = pa.Field(nullable=False) class BasinStateSchema(_BaseSchema): - node_id: Series[int] = pa.Field(nullable=False) + node_id: Series[int] = pa.Field(nullable=False, default=0) level: Series[float] = pa.Field(nullable=False) class BasinStaticSchema(_BaseSchema): - node_id: Series[int] = pa.Field(nullable=False) + node_id: Series[int] = pa.Field(nullable=False, default=0) drainage: Series[float] = pa.Field(nullable=True) potential_evaporation: Series[float] = pa.Field(nullable=True) infiltration: Series[float] = pa.Field(nullable=True) @@ -33,13 +33,13 @@ class BasinStaticSchema(_BaseSchema): class BasinSubgridSchema(_BaseSchema): subgrid_id: Series[int] = pa.Field(nullable=False) - node_id: Series[int] = pa.Field(nullable=False) + node_id: Series[int] = pa.Field(nullable=False, default=0) basin_level: Series[float] = pa.Field(nullable=False) subgrid_level: Series[float] = pa.Field(nullable=False) class BasinTimeSchema(_BaseSchema): - node_id: Series[int] = pa.Field(nullable=False) + node_id: Series[int] = pa.Field(nullable=False, default=0) time: Series[Timestamp] = pa.Field(nullable=False) drainage: Series[float] = pa.Field(nullable=True) potential_evaporation: Series[float] = pa.Field(nullable=True)