Show allocation flows on the map in QGIS

[SouthEndMusic]

What In addition to https://github.com/Deltares/Ribasim/issues/969, we also want to be able to visualize the allocation flows who have their own output file introduced in https://github.com/Deltares/Ribasim/pull/1012. There are some notable distinctions between allocation flow and physical flow:

• Allocation flow is given per priority;
• If the model has a main network and subnetworks, for the subnetworks there is both a demand collecting flow and an actual allocating flow per edge (distinguished using the boolean column collect_demands);
• Edges connecting the main network to a subnetwork are considered to be both part of the main network and the subnetwork;
• It can seem that the water balance isn't closed around the main network to subnetwork connection. This is because all flow allocated to a subnetwork is redistributed over priorities within the subnetwork.

Best done after #1399.

[visr]

This should be added to to_xugrid from #1369. One possible issue is that the time dimension is different, since saveat and allocation_timestep is not the same. Not sure what is better, to write a separate UGrid file, or to have two time dimensions in there.

[Huite]

This might be constrained by what QGIS and MDAL accept. It might be that "time" is special cased as a literal label.

In that case, the only option is a separate file.

[visr]

It looks like a separate file is needed. Allocation flow has 4 dimensions:

1. time
2. edge
3. optimization_type (one of the 3 given in https://deltares.github.io/Ribasim/core/allocation.html#the-high-level-algorithm)
4. priority

As discussed with @JoerivanEngelen the 3rd and 4th dimension needs to be added in the variable name of variables with dimension 1 and 2 only. To allow users to select one. So I used varname = f"{optimization_type}_priority_{priority}". Also discussed with @SouthEndMusic to add a variable I now call flow_rate_allocated, which is optimization_type = allocate and the sum over all priorities.

I'm not sure what's the best API. We need model.to_xugrid(add_results=False) to avoid the other time dimension in the same file. I suppose we could do model.to_xugrid(add_results=False, add_allocation_results=True), but open to other ideas.

import ribasim

model = ribasim.Model.read(toml_path)
uds = model.to_xugrid(add_results=False)

results_path = toml_path.parent / model.results_dir
alloc_flow_path = results_path / "allocation_flow.arrow"
alloc_flow_df = pd.read_feather(
    alloc_flow_path,
    columns=["time", "edge_id", "flow_rate", "optimization_type", "priority"],
)

edge_dim = uds.grid.edge_dimension
node_dim = uds.grid.node_dimension

# from edge_id to the edge_dim index
edge_lookup = pd.Series(
    index=uds["edge_id"],
    data=uds[edge_dim],
    name="edge_index",
)

alloc_flow_df[edge_dim] = edge_lookup[alloc_flow_df["edge_id"]].to_numpy()
alloc_flow_df["time"] = alloc_flow_df["time"].astype("datetime64[ns]")

# "flow_rate_allocated" is the sum of all allocated flow rates over the priorities
allocate_df = alloc_flow_df.loc[alloc_flow_df["optimization_type"] == "allocate"]
uds["flow_rate_allocated"] = (
    allocate_df.groupby(["time", edge_dim])["flow_rate"].sum().to_xarray()
)

# also add the individual priorities and optimization types
# added as separate variables to ensure QGIS / MDAL compatibility
for (optimization_type, priority), group in alloc_flow_df.groupby(
    ["optimization_type", "priority"]
):
    varname = f"{optimization_type}_priority_{priority}"
    da = group.set_index(["time", edge_dim])["flow_rate"].to_xarray()
    uds[varname] = da

uds.ugrid.to_netcdf(toml_path.parent / "results" / "ugrid_allocation.nc")