From 04b23d005a78e8ab044e66e3eeede60a08542651 Mon Sep 17 00:00:00 2001
From: Maarten Pronk <git@evetion.nl>
Date: Sat, 20 Apr 2024 22:07:14 +0200
Subject: [PATCH] Merge main.

---
 coupling/delwaq/delwaq_util.py | 54 ----------------------------------
 coupling/delwaq/gen_delwaq.py  | 29 ++++++++++--------
 2 files changed, 17 insertions(+), 66 deletions(-)

diff --git a/coupling/delwaq/delwaq_util.py b/coupling/delwaq/delwaq_util.py
index 0cb413713..1c0efa36d 100644
--- a/coupling/delwaq/delwaq_util.py
+++ b/coupling/delwaq/delwaq_util.py
@@ -6,8 +6,6 @@
 
 import numpy as np
 import pandas as pd
-import ribasim
-import xugrid as xu
 
 
 def strfdelta(tdelta):
@@ -89,55 +87,3 @@ def write_flows(fn: Path | str, data: pd.DataFrame, timestep: timedelta):
         # Delwaq needs an extra timestep after the end
         f.write(struct.pack("<i", int(time + timestep.total_seconds())))
         f.write(group.flow_rate.to_numpy().astype("float32").tobytes())
-
-
-def ugridify(model: ribasim.Model):
-    node_df = model.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 = model.edge.df.copy()
-    # edge_df = edge_df[edge_df.edge_type == "flow"]
-    edge_df.set_crs(epsg=28992, inplace=True, allow_override=True)
-    node_df.set_crs(epsg=28992, inplace=True, allow_override=True)
-
-    node_id = node_df.node_id.to_numpy(dtype="int32")
-    edge_id = edge_df.index.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",
-    )
-
-    grid = xu.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_network",
-        projected=node_df.crs.is_projected,
-        crs=node_df.crs,
-    )
-
-    edge_dim = grid.edge_dimension
-    node_dim = grid.node_dimension
-
-    uds = xu.UgridDataset(None, grid)
-    uds = uds.assign_coords(node_id=(node_dim, node_id))
-    uds = uds.assign_coords(edge_id=(edge_dim, edge_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
diff --git a/coupling/delwaq/gen_delwaq.py b/coupling/delwaq/gen_delwaq.py
index ef153bdc2..63a9fada7 100644
--- a/coupling/delwaq/gen_delwaq.py
+++ b/coupling/delwaq/gen_delwaq.py
@@ -11,7 +11,6 @@
 import ribasim
 from delwaq_util import (
     strfdelta,
-    ugridify,
     write_flows,
     write_pointer,
     write_volumes,
@@ -52,10 +51,17 @@
 # Flows on non-existing edges indicate where the boundaries are
 tg = flows.groupby("time")
 g = tg.get_group(flows.time[0])
-nboundary = g.edge_id.isna().sum()
-boundary_nodes = g.to_node_id[g.edge_id.isna()]
-new_boundary_ids = np.tile(-np.arange(1, nboundary + 1), tg.ngroups)
-flows.loc[flows.edge_id.isna(), "from_node_id"] = new_boundary_ids
+
+boundary_types = ["LevelBoundary", "FlowBoundary", "Terminal"]
+node = model.node_table().df
+bids = node.node_id[node.node_type.isin(boundary_types)]
+nboundary = len(bids)
+
+m = flows.from_node_type.isin(boundary_types)
+flows.from_node_id[m] = flows.from_node_id[m] * -1
+m = flows.to_node_type.isin(boundary_types)
+flows.to_node_id[m] = flows.to_node_id[m] * -1
+
 # flows.to_csv(output_folder / "flows.csv", index=False)  # not needed
 
 
@@ -67,8 +73,8 @@
 )
 write_pointer(output_folder / "ribasim.poi", pointer)
 
-total_segments = len(node.index)
-total_exchanges = len(edge.index) + nboundary
+total_segments = len(node) - nboundary
+total_exchanges = len(edge.index)
 
 
 # Write attributes template
@@ -81,7 +87,7 @@
     )
 
 # Generate mesh and write to NetCDF
-uds = ugridify(model)
+uds = model.to_xugrid()
 uds.ugrid.to_netcdf(output_folder / "ribasim.nc")
 
 # Generate area and flows
@@ -99,11 +105,11 @@
 ntime = basins.time.unique()
 basins.drop(columns=["level"], inplace=True)
 
-non_basins = set(node.index) - set(basins.node_id)
+non_basins = set(node.index) - set(basins.node_id) - set(bids)
 rtime = ntime - ntime[0]
 volumes_nbasin = pd.DataFrame(
     {
-        "time": np.repeat(rtime, len(non_basins)),
+        "time": np.repeat(basins.time.unique(), len(non_basins)),
         "node_id": np.tile(list(non_basins), len(rtime)),
         "storage": fillvolume,
     }
@@ -129,7 +135,6 @@
 write_flows(output_folder / "ribasim.len", lengths, timestep)
 
 # Find our boundaries
-
 boundaries = []
 substances = set()
 
@@ -175,7 +180,7 @@ def make_boundary(id, type):
         )
     )
 
-bnd = node[node.node_id.isin(boundary_nodes)].reset_index(drop=True)
+bnd = node[node.node_id.isin(bids)].reset_index(drop=True)
 bnd["fid"] = bnd["node_type"].str[:10] + "_" + bnd["node_id"].astype(str)
 bnd["comment"] = ""
 bnd = bnd[["fid", "comment", "node_type"]]