Add model.to_xugrid()

python/ribasim/ribasim/model.py

@@ -17,6 +17,11 @@
     model_validator,
 )
 
+try:
+    import xugrid
+except ImportError:
+    xugrid = None
+
 import ribasim
 from ribasim.config import (
     Allocation,
@@ -342,3 +347,58 @@
         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.
+        assert node_df.node_id.is_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(dtype="int32")
+        from_node_id = edge_df.from_node_id.to_numpy(dtype="int32")
+        to_node_id = edge_df.to_node_id.to_numpy(dtype="int32")
+
+        # from node_id to the node_dim index
+        node_lookup = pd.Series(
+            index=node_id,
+            data=node_id.argsort().astype("int32"),
+            name="node_index",
+        )
+
+        if node_df.crs is None:
+            # 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/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"