diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py index 34c808f73..7a395ab5a 100644 --- a/python/ribasim/ribasim/model.py +++ b/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[edge_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 diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py index 476cf8485..bd6fcc7a8 100644 --- a/python/ribasim/tests/test_io.py +++ b/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) diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py index bdc60368c..3e2ed89d6 100644 --- a/python/ribasim/tests/test_model.py +++ b/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