Pilots: derive secondary and tertiary basins based on MODFLOW6 input

[Huite]

What
Derive 2nd and 3rd basins based on MODFLOW6 input

Why
From yesterday's discussion: our goal with testing in the pilots is to also test the explicit coupling of secondary basins, which are "stacked" on top of the primary basins. The challenge here is that the in the pilot area (Hooge Raam), only the primary water is described in the hydraulic model. Hence, the result of the lumping tool will only provide parametrizations of these basins.

How
To properly test the functionality, we need to:

[ ] Add a secondary basin for each primary basin:

• The parametrization may be done in a very pragmatic manner. My thinking currently:
  • We define the profiles on the basis of the MODFLOW 6 conductance.
  • We use the conductance to compute a wetted area.
  • We assume either a rectangular profile, or a (45 degree) trapezoidal profile.
  • We derive an average water level; since the primary water level is generally located behind the weir, I suggest we compute both MODFLOW6 primary and secondary average water levels, and compute the difference. We use this difference as an bed level offset for the secondary basin (which should be located somewhat higher).

[ ] The second step consists of computing either an "effective" Manning element, or a TabulatedRatingCurve.

• An easy way to estimate the TabulatedRatingCurve is by assuming a linear Q-h relationship and defining two water depth + two flow rates.
  • A linear resistance is also a possibility.
  • The choice depends on what we want to see in terms of infiltration drawing from the primary water (which likely will not happen in this pilot area for most lumping levels).

[ ] Define a subgrid entry for each secondary river element. We can easily do this by assuming a constant water depth.
[ ] Next, we update primod to take the "stacked" basin definitions. This requires a bit of thinking in terms of what's nicest, but inevitably boils down to creating an additional mapping (which basin-definition or which set of nodes for which set of MODFLOW 6 packages). Perhaps the easiest way is to add the basin_definition entry to the driver coupling.

[ ] The next step is to repeat this for the drainage package in the MODFLOW model.

Additional context
The final form of this is explicitly not an official part of the tooling: it's an ad hoc solution to get testing input.

[Huite]

This is far easier to do with via an "add" API. #1110

[Huite]

Implemented somewhat pragmatically, sharing it here for posterity:

# %%

import ribasim
import numpy as np
import pandas as pd
import geopandas as gpd
import imod

from ribasim import Node
from ribasim.nodes import basin, linear_resistance

from scipy.spatial import KDTree

# %%

ribasim_model = ribasim.Model.read("ribasim-test/ribasim.toml")
basins = gpd.read_file("basins.gpkg", layer="basin_areas")
simulation = imod.mf6.Modflow6Simulation.from_file("mf6-test-10/GWF_1.toml")
gwf = simulation["GWF_1"]
secondary = gwf["riv-4"]

# %%

conductance = secondary["conductance"].isel(time=0, drop=True).max("layer")
bottom = secondary["elevation"].isel(time=0, drop=True).max("layer")
entry_resistance = 1.0
wetted_area = conductance / entry_resistance

gridded_basins = imod.prepare.rasterize(basins, column="node_id", like=wetted_area)
wetted_area_per_basin = wetted_area.groupby(gridded_basins).sum()
wetted_area_per_basin = wetted_area_per_basin.rename({"group": "node_id"})
wetted_area_per_basin.name = "area"
wetted_area_per_basin["node_id"] = wetted_area_per_basin["node_id"].astype(int)
wetted_area_df = wetted_area_per_basin.to_dataframe()

# %%

assert ribasim_model.basin.node.df is not None
primary_basin_node_ids = basins["node_id"].to_numpy()
secondary_basin_geometries = ribasim_model.basin.node.df[
    ["node_id", "geometry"]
].set_index("node_id")
secondary_basin_geometries.geometry = (
    secondary_basin_geometries.geometry.affine_transform([1, 0, 0, 1, -100, -100])
)
linear_resistance_geometries = ribasim_model.basin.node.df[
    ["node_id", "geometry"]
].set_index("node_id")
linear_resistance_geometries.geometry = (
    linear_resistance_geometries.geometry.affine_transform([1, 0, 0, 1, -50, -50])
)
# %%

node_offset = 1000
basin_area = basins[["node_id", "geometry"]].set_index("node_id")

subgrid_df = ribasim_model.basin.subgrid.df
n_subgrid = subgrid_df["subgrid_id"].iloc[-1].item() + 1
subgrid_xy = subgrid_df.groupby("subgrid_id").first()[["meta_x", "meta_y"]]
subgrid_gdf = gpd.GeoDataFrame(
    index=subgrid_xy.index,
    geometry=gpd.points_from_xy(x=subgrid_xy["meta_x"], y=subgrid_xy["meta_y"])
)
# %%

tmp = subgrid_df.set_index("subgrid_id")
nrow = subgrid_df.groupby("subgrid_id").count()["node_id"]

# %%
subgrid_dfs = []
for node_id in primary_basin_node_ids:
    basin_polygon = basin_area.loc[node_id].geometry
    basin_point = secondary_basin_geometries.loc[node_id].item()
    resistance_point = linear_resistance_geometries.loc[node_id].item()

    # Compute properties based on primary profile and modflow input.
    df = ribasim_model.basin.profile.df
    df = df.loc[df["node_id"] == node_id, :]
    basin_level = df["level"].to_numpy()
    is_basin = gridded_basins == node_id
    area = wetted_area.where(is_basin).sum()
    # Compute resistance assuming 1 mm/d of drainage leads to 1 cm of level change.
    Q = basin_polygon.area * 0.001 / 86_400.0
    dH = 0.01
    resistance = Q / dH
    
    # Setup subgrid data
    # Find xy data of modflow boundaries
    is_basin = gridded_basins == node_id
    xy = (
        wetted_area.where(is_basin)
        .stack(stacked_yx=["y", "x"])
        .dropna("stacked_yx")
        .to_dataframe(name="meta")
        .loc[:, ["x" ,"y"]]
        .to_numpy()
    )
    inside = subgrid_gdf.loc[subgrid_gdf.intersects(basin_polygon)]
    tree = KDTree(data=np.column_stack(
        (inside.geometry.x, inside.geometry.y)
    ))
    _, indices = tree.query(xy)
    subgrid_id = inside.index[indices]
    new_subgrid_entries = tmp.loc[subgrid_id].reset_index()
    nrow_subgrid = nrow.loc[subgrid_id]
    new_subgrid_entries["subgrid_id"] = np.repeat(np.arange(n_subgrid, n_subgrid + len(indices)), nrow_subgrid)
    new_subgrid_entries["meta_x"] = np.repeat(xy[:, 0], nrow_subgrid)
    new_subgrid_entries["meta_y"] = np.repeat(xy[:, 1], nrow_subgrid)
    new_subgrid_entries["meta_original_subgrid_id"] = np.repeat(subgrid_id, nrow_subgrid)
    n_subgrid += len(indices)
    subgrid_dfs.append(new_subgrid_entries)


    # Add new nodes
    new_id = node_id + node_offset
    ribasim_model.basin.add(
        node=Node(node_id=new_id, geometry=basin_point),
        tables=[
            basin.Profile(area=[area, area], level=[basin_level[0], basin_level[1]]),
            basin.Static(precipitation=[0.001 / 86_400.0]),
            basin.State(node_id=new_id, level=[basin_level[0] + 0.1]),
        ],
    )
    ribasim_model.linear_resistance.add(
        node=Node(node_id=new_id, geometry=resistance_point),
        tables=[linear_resistance.Static(resistance=[resistance])],
    )

    # Create the edges and connect.
    ribasim_model.edge.add(
        ribasim_model.basin[new_id], ribasim_model.linear_resistance[new_id]
    )
    ribasim_model.edge.add(
        ribasim_model.linear_resistance[new_id], ribasim_model.basin[node_id]
    )

# %%
ribasim_model.basin.subgrid.df["meta_original_subgrid_id"] = ribasim_model.basin.subgrid.df["subgrid_id"].copy()
ribasim_model.basin.subgrid.df = pd.concat(
    [ribasim_model.basin.subgrid.df] + subgrid_dfs
)

# %%

check = ribasim_model.basin.subgrid.df.groupby("subgrid_id").first()
gdf = gpd.GeoDataFrame(
    index=check.index,
    geometry=gpd.points_from_xy(x=check["meta_x"], y=check["meta_y"]),
    data=check,
)
gdf.to_file("all_subgrid.shp")

# %%

secondary_basins = basins.copy()
secondary_basins["node_id"] += node_offset
ribasim_model.write("test-secondary-basins-subgrid/ribasim.toml")
basins.to_file("test-secondary-basins-subgrid/primary-basins.gpkg")
secondary_basins.to_file("test-secondary-basins-subgrid/secondary-basins.gpkg")

# %%
#import pandas as pd
#df = pd.read_feather(r"c:\projects\TKI-Ribasim\joachim\test-secondary-basins-subgrid\results\basin.arrow")
#df.loc[df["node_id"] == 98, ["time", "level"]].set_index("time").plot()
#df.loc[df["node_id"] == 1098, ["time", "level"]].set_index("time").plot()