Add model.to_xugrid() (#1314)

This adds an optional dependency on `xugrid`, to facilitate creating
UGRID netCDFs, to support #969.

For now I did not yet add any results to the dataset, but this could be
done as a follow-up. This is just the network.

---------

Co-authored-by: Hofer-Julian <[email protected]>

python/ribasim/ribasim/model.py

@@ -47,6 +47,12 @@
     FileModel,
     context_file_loading,
 )
+from ribasim.utils import MissingOptionalModule
+
+try:
+    import xugrid
+except ImportError:
+    xugrid = MissingOptionalModule("xugrid")
 
 
 class Model(FileModel):
@@ -342,3 +348,59 @@ def plot(self, ax=None, indicate_subnetworks: bool = True) -> Any:
         ax.legend(handles, labels, loc="lower left", bbox_to_anchor=(1, 0.5))
 
         return ax
+
+    def to_xugrid(self) -> xugrid.UgridDataset:
+        """Convert the network to a xugrid.UgridDataset."""
+        node_df = self.node_table().df
+
+        # This will need to be adopted for locally unique node IDs,
+        # otherwise the `node_lookup` with `argsort` is not correct.
+        if not node_df.node_id.is_unique:
+            raise ValueError("node_id must be unique")
+        node_df.sort_values("node_id", inplace=True)
+
+        edge_df = self.edge.df.copy()
+        # We assume only the flow network is of interest.
+        edge_df = edge_df[edge_df.edge_type == "flow"]
+
+        node_id = node_df.node_id.to_numpy()
+        from_node_id = edge_df.from_node_id.to_numpy()
+        to_node_id = edge_df.to_node_id.to_numpy()
+
+        # from node_id to the node_dim index
+        node_lookup = pd.Series(
+            index=node_id,
+            data=node_id.argsort().astype(np.int32),
+            name="node_index",
+        )
+
+        if node_df.crs is None:
+            # TODO: can be removed when CRS is required, #1254
+            projected = False
+        else:
+            projected = node_df.crs.is_projected
+
+        grid = xugrid.Ugrid1d(
+            node_x=node_df.geometry.x,
+            node_y=node_df.geometry.y,
+            fill_value=-1,
+            edge_node_connectivity=np.column_stack(
+                (
+                    node_lookup[from_node_id],
+                    node_lookup[to_node_id],
+                )
+            ),
+            name="ribasim",
+            projected=projected,
+            crs=node_df.crs,
+        )
+
+        edge_dim = grid.edge_dimension
+        node_dim = grid.node_dimension
+
+        uds = xugrid.UgridDataset(None, grid)
+        uds = uds.assign_coords(node_id=(node_dim, node_id))
+        uds = uds.assign_coords(from_node_id=(edge_dim, from_node_id))
+        uds = uds.assign_coords(to_node_id=(edge_dim, to_node_id))
+
+        return uds

python/ribasim/ribasim/utils.py

@@ -5,3 +5,13 @@ def _pascal_to_snake(pascal_str):
     # Insert a '_' before all uppercase letters that are not at the start of the string
     # and convert the string to lowercase
     return re.sub(r"(?<!^)(?=[A-Z])", "_", pascal_str).lower()
+
+
+class MissingOptionalModule:
+    """Presents a clear error for optional modules."""
+
+    def __init__(self, name):
+        self.name = name
+
+    def __getattr__(self, name):
+        raise ImportError(f"{self.name} is required for this functionality")

python/ribasim/tests/test_model.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import xugrid
 from pydantic import ValidationError
 from ribasim.config import Solver
 from ribasim.geometry.edge import NodeData
@@ -216,3 +217,15 @@ def test_indexing(basic):
         match=re.escape("Cannot index into Basin / time: it contains no data."),
     ):
         model.basin.time[1]
+
+
+def test_xugrid(basic, tmp_path):
+    uds = basic.to_xugrid()
+    assert isinstance(uds, xugrid.UgridDataset)
+    assert uds.grid.edge_dimension == "ribasim_nEdges"
+    assert uds.grid.node_dimension == "ribasim_nNodes"
+    assert uds.grid.crs is None
+    assert uds.node_id.dtype == np.int32
+    uds.ugrid.to_netcdf(tmp_path / "ribasim.nc")
+    uds = xugrid.open_dataset(tmp_path / "ribasim.nc")
+    assert uds.attrs["Conventions"] == "CF-1.9 UGRID-1.0"