Merge branch 'main' into equal_ratio_allocation_objective

.github/workflows/python_lint.yml → .github/workflows/mypy.yml

@@ -1,4 +1,4 @@
-name: Python Lint
+name: Mypy Type Check
 on:
   push:
     branches: [main, update/pixi-lock]

.vscode/extensions.json

@@ -7,6 +7,7 @@
         "njpwerner.autodocstring",
         "quarto.quarto",
         "tamasfe.even-better-toml",
-        "samuelcolvin.jinjahtml"
+        "samuelcolvin.jinjahtml",
+        "yy0931.vscode-sqlite3-editor"
     ]
 }

docs/python/examples.ipynb

@@ -854,7 +854,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Setup flow boundary:\n"
+    "Setup level boundary:\n"
    ]
   },
   {
@@ -865,7 +865,7 @@
    "source": [
     "model.level_boundary.add(\n",
     "    Node(4, Point(3.0, 0.0)),\n",
-    "    [level_boundary.Static(level=[1])],\n",
+    "    [level_boundary.Static(level=[5.0])],\n",
     ")"
    ]
   },
@@ -882,24 +882,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pid_control_data = [\n",
-    "    pid_control.Time(\n",
-    "        time=[\n",
-    "            \"2020-01-01\",\n",
-    "            \"2020-05-01\",\n",
-    "            \"2020-07-01\",\n",
-    "            \"2020-12-01\",\n",
-    "        ],\n",
-    "        listen_node_id=[2, 2, 2, 2],\n",
-    "        listen_node_type=[\"Basin\", \"Basin\", \"Basin\", \"Basin\"],\n",
-    "        target=[5.0, 5.0, 7.5, 7.5],\n",
-    "        proportional=[-1e-3, 1e-3, 1e-3, 1e-3],\n",
-    "        integral=[-1e-7, 1e-7, -1e-7, 1e-7],\n",
-    "        derivative=[0.0, 0.0, 0.0, 0.0],\n",
-    "    )\n",
-    "]\n",
-    "model.pid_control.add(Node(5, Point(1.5, 1.0)), pid_control_data)\n",
-    "model.pid_control.add(Node(7, Point(1.5, -1.0)), pid_control_data)"
+    "for node, proportional, integral in [\n",
+    "    (Node(5, Point(1.5, 1.0)), -1e-3, -1e-7),\n",
+    "    (Node(7, Point(1.5, -1.0)), 1e-3, 1e-7),\n",
+    "]:\n",
+    "    pid_control_data = [\n",
+    "        pid_control.Time(\n",
+    "            time=[\n",
+    "                \"2020-01-01\",\n",
+    "                \"2020-05-01\",\n",
+    "                \"2020-07-01\",\n",
+    "                \"2020-12-01\",\n",
+    "            ],\n",
+    "            listen_node_id=2,\n",
+    "            listen_node_type=\"Basin\",\n",
+    "            target=[5.0, 5.0, 7.5, 7.5],\n",
+    "            proportional=proportional,\n",
+    "            integral=integral,\n",
+    "            derivative=0.0,\n",
+    "        )\n",
+    "    ]\n",
+    "    model.pid_control.add(node, pid_control_data)"
    ]
   },
   {

python/ribasim/ribasim/model.py

@@ -264,10 +264,10 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         # TODO
         # self.validate_model()
         filepath = Path(filepath)
+        self.filepath = 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)
@@ -280,7 +280,7 @@ def write(self, filepath: str | PathLike[str]) -> Path:
     def _load(cls, filepath: Path | None) -> dict[str, Any]:
         context_file_loading.set({})
 
-        if filepath is not None:
+        if filepath is not None and filepath.is_file():
             with open(filepath, "rb") as f:
                 config = tomli.load(f)
 
@@ -395,9 +395,10 @@ def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
 
         return ax
 
-    def to_xugrid(self):
+    def to_xugrid(self, add_results: bool = True):
         """
-        Convert the network to a `xugrid.UgridDataset`.
+        Convert the network and results to a `xugrid.UgridDataset`.
+        To get the network only, set `add_results=False`.
         This method will throw `ImportError`,
         if the optional dependency `xugrid` isn't installed.
         """
@@ -449,4 +450,63 @@ def to_xugrid(self):
         uds = uds.assign_coords(from_node_id=(edge_dim, from_node_id))
         uds = uds.assign_coords(to_node_id=(edge_dim, to_node_id))
 
+        if add_results:
+            uds = self._add_results(uds)
+
+        return uds
+
+    def _add_results(self, uds):
+        toml_path = self.filepath
+        if toml_path is None:
+            raise FileNotFoundError("Model must be written to disk to add results.")
+
+        results_path = toml_path.parent / self.results_dir
+        basin_path = results_path / "basin.arrow"
+        flow_path = results_path / "flow.arrow"
+
+        if not basin_path.is_file() or not flow_path.is_file():
+            raise FileNotFoundError(
+                f"Cannot find results in '{results_path}', "
+                "perhaps the model needs to be run first."
+            )
+
+        basin_df = pd.read_feather(basin_path)
+        flow_df = pd.read_feather(flow_path)
+
+        edge_dim = uds.grid.edge_dimension
+        node_dim = uds.grid.node_dimension
+
+        # from node_id to the node_dim index
+        node_lookup = pd.Series(
+            index=uds["node_id"],
+            data=uds[node_dim],
+            name="node_index",
+        )
+        # from edge_id to the edge_dim index
+        edge_lookup = pd.Series(
+            index=uds["edge_id"],
+            data=uds[edge_dim],
+            name="edge_index",
+        )
+
+        basin_df = pd.read_feather(basin_path)
+        flow_df = pd.read_feather(flow_path)
+
+        # datetime64[ms] gives trouble; https://github.com/pydata/xarray/issues/6318
+        flow_df["time"] = flow_df["time"].astype("datetime64[ns]")
+        basin_df["time"] = basin_df["time"].astype("datetime64[ns]")
+
+        # add flow results to the UgridDataset
+        flow_df[edge_dim] = edge_lookup[flow_df["edge_id"]].to_numpy()
+        flow_da = flow_df.set_index(["time", edge_dim])["flow_rate"].to_xarray()
+        uds[flow_da.name] = flow_da
+
+        # add basin results to the UgridDataset
+        basin_df[node_dim] = node_lookup[basin_df["node_id"]].to_numpy()
+        basin_df.drop(columns=["node_id"], inplace=True)
+        basin_ds = basin_df.set_index(["time", node_dim]).to_xarray()
+
+        for var_name, da in basin_ds.data_vars.items():
+            uds[var_name] = da
+
         return uds

python/ribasim/tests/test_io.py

@@ -29,8 +29,11 @@ def __assert_equal(a: DataFrame, b: DataFrame) -> None:
 def test_basic(basic, tmp_path):
     model_orig = basic
     toml_path = tmp_path / "basic/ribasim.toml"
+    assert model_orig.filepath is None
     model_orig.write(toml_path)
+    assert model_orig.filepath == toml_path
     model_loaded = Model.read(toml_path)
+    assert model_loaded.filepath == toml_path
 
     with open(toml_path, "rb") as f:
         toml_dict = tomli.load(f)

python/ribasim/tests/test_model.py

@@ -233,7 +233,7 @@ def test_indexing(basic):
 
 
 def test_xugrid(basic, tmp_path):
-    uds = basic.to_xugrid()
+    uds = basic.to_xugrid(add_results=False)
     assert isinstance(uds, xugrid.UgridDataset)
     assert uds.grid.edge_dimension == "ribasim_nEdges"
     assert uds.grid.node_dimension == "ribasim_nNodes"
@@ -243,6 +243,13 @@ def test_xugrid(basic, tmp_path):
     uds = xugrid.open_dataset(tmp_path / "ribasim.nc")
     assert uds.attrs["Conventions"] == "CF-1.9 UGRID-1.0"
 
+    with pytest.raises(FileNotFoundError, match="Model must be written to disk"):
+        basic.to_xugrid(add_results=True)
+
+    basic.write(tmp_path / "ribasim.toml")
+    with pytest.raises(FileNotFoundError, match="Cannot find results"):
+        basic.to_xugrid(add_results=True)
+
 
 def test_to_crs(bucket: Model):
     model = bucket