From 2bec7008fa270f73ec07c419a9058c6cf895341c Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Sat, 18 Jul 2026 21:01:03 +0200 Subject: [PATCH 1/4] Add nhflotools.lakes and polder DRN exclude for Bergen pond/lake stages Extract the lake-cell carve into nhflotools.lakes with a single aggregator (_aggregate_lake_cells) shared by carve_lake_cells, riv_from_lakes_pwn (per-lake RIV stage boundary) and recharge_pond_mask, so the carved-cell set, the stage reach set and the recharge-exclusion mask are equal by construction. Add an optional exclude mask to polder.drn_from_waterboard_data that nulls drn_cond and drn_elev at excluded cells so a reach is dropped whether its stage came from HHNK peilgebied data or the maaiveld fallback. Add unit tests on a synthetic disv grid. --- src/nhflotools/lakes.py | 198 +++++++++++++++++++++++++++++++++++++++ src/nhflotools/polder.py | 18 +++- tests/conftest.py | 97 +++++++++++++++++++ tests/test_lakes.py | 157 +++++++++++++++++++++++++++++++ tests/test_polder.py | 59 ++++++++++++ 5 files changed, 528 insertions(+), 1 deletion(-) create mode 100644 src/nhflotools/lakes.py create mode 100644 tests/conftest.py create mode 100644 tests/test_lakes.py create mode 100644 tests/test_polder.py diff --git a/src/nhflotools/lakes.py b/src/nhflotools/lakes.py new file mode 100644 index 0000000..5feeb97 --- /dev/null +++ b/src/nhflotools/lakes.py @@ -0,0 +1,198 @@ +"""Helpers for the Bergen pond/lake (``lakes_pwn``) features. + +The PWN model carves the model top down to the lake bottom in every grid cell that a +managed pond/lake sufficiently covers, and then holds those carved cells at their +prescribed lake stage with a per-lake RIV boundary. Both the carved cell set and the RIV +reach set are derived from the single aggregator :func:`_aggregate_lake_cells`, so the two +sets are equal by construction: no carved cell is ever left without a stage, and no stage +reach ever lands on an un-carved cell. + +The logic lives here (rather than inline in the model script) so it can be imported and +unit-tested without running the full REGIS/AHN/MF6 build. +""" + +import logging + +import flopy +import nlmod +import numpy as np +import xarray as xr + +logger = logging.getLogger(__name__) + + +def _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): + """Aggregate lake pieces to one carved/stage record per grid cell. + + This is the single source of truth for the carved-cell set. It (1) keeps only lake + pieces that carry both a stage (``strt``) and a bottom (``botm``), (2) computes each + cell's lake coverage over exactly those pieces, (3) keeps cells whose coverage is + strictly greater than ``min_area_fraction``, and (4) collapses the surviving pieces to + one record per cell. + + Parameters + ---------- + ds : xarray.Dataset + Model dataset. Only ``ds['area']`` (cell area indexed by ``icell2d``) is used. + gdf_lake_grid : geopandas.GeoDataFrame + Lake polygons already intersected with the model grid (see + :func:`nlmod.dims.gdf_to_grid`). Must carry a ``cellid`` column and the + ``strt``, ``botm``, ``clake`` and ``identificatie`` columns of ``lakes_pwn``. + min_area_fraction : float, optional + A cell is kept when the lake covers strictly more than this fraction of its area. + The default is 0.5, so a cell exactly half covered is not carved. + + Returns + ------- + pandas.DataFrame + Indexed by ``cellid`` with columns ``strt`` (area-weighted), ``botm`` (minimum), + ``cond`` (summed ``piece_area / clake``, units m2/d) and ``identificatie`` + (first). Empty-input semantics follow :func:`nlmod.grid.aggregate_vector_per_cell`. + """ + lake = gdf_lake_grid.dropna(subset=["strt", "botm"]).copy() + n_drop = len(gdf_lake_grid) - len(lake) + if n_drop: + logger.warning("Dropping %d lake piece(s) missing strt or botm", n_drop) + + # Use the true clipped-piece area throughout (coverage, conductance and the + # area-weighted stage), so the three are mutually consistent regardless of any stale + # 'area' column carried over from gdf_to_grid. + lake["area"] = lake.geometry.area + overlap = lake["area"].groupby(lake["cellid"]).transform("sum").to_numpy() + cell_area = ds["area"].sel(icell2d=lake["cellid"].to_numpy()).to_numpy() + # Strict '>' so a cell exactly at min_area_fraction is excluded. + lake = lake[overlap / cell_area > min_area_fraction] + + lake["cond"] = lake["area"] / lake["clake"] + return nlmod.grid.aggregate_vector_per_cell( + lake, + fields_methods={ + "strt": "area_weighted", + "botm": "min", + "cond": "sum", + "identificatie": "first", + }, + ) + + +def carve_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): + """Lower the model top to the lake bottom in sufficiently covered cells. + + The model top is set to the aggregated lake bottom in every carved cell and a boolean + ``ds['lake_cell']`` marker (True exactly at the carved cells) is added for the recharge + mask, the polder-drain exclusion and the lake RIV to reuse. + + Parameters + ---------- + ds : xarray.Dataset + Model dataset with ``top``, ``botm`` and ``area``. + gdf_lake_grid : geopandas.GeoDataFrame + Lake polygons intersected with the model grid (see :func:`_aggregate_lake_cells`). + min_area_fraction : float, optional + Minimum lake coverage for a cell to be carved, by default 0.5 (strict ``>``). + + Returns + ------- + ds : xarray.Dataset + Copy of the input dataset with the lowered top, ``thickness`` dropped (so it is + recomputed from the new top/botm), and the boolean ``lake_cell`` marker. + lake_cellids : numpy.ndarray + The ``icell2d`` values of the carved cells. + + Notes + ----- + ``ds['ahn']`` is intentionally not refreshed to the carved top. The only post-carve + reader of ``ds['ahn']`` is the polder DRN, which now excludes these cells, so a refresh + would be inert; refreshing ``ahn`` to the lake bottom would also corrupt its meaning as + the measured maaiveld. + """ + agg = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) + lake_cellids = agg.index.to_numpy() + + top = ds["top"].copy() + top.loc[{"icell2d": lake_cellids}] = agg["botm"].to_numpy() + ds = nlmod.layers.set_model_top(ds, top) + ds = ds.drop_vars("thickness", errors="ignore") + + ds["lake_cell"] = xr.zeros_like(ds["top"], dtype=bool) + ds["lake_cell"].loc[{"icell2d": lake_cellids}] = True + return ds, lake_cellids + + +def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): + """Hold the carved lake cells at their prescribed stage with a per-lake RIV. + + A RIV reach is placed in every carved cell with ``stage = strt``, ``rbot = botm`` (the + carved top) and ``cond = sum(piece_area / clake)``. Because ``rbot`` equals the carved + top, the reach caps bed infiltration once the head drops below the lakebed + (perched-pond behaviour) while draining freely when the head rises above the stage. + + The reach set is derived from the same :func:`_aggregate_lake_cells` call as + :func:`carve_lake_cells`, so it equals the carved-cell set by construction. + + Parameters + ---------- + ds : xarray.Dataset + Model dataset (post-carve), with ``top``, ``botm``, ``kh`` and idomain. + gwf : flopy.mf6.ModflowGwf + Groundwater flow model the RIV package is added to. + gdf_lake_grid : geopandas.GeoDataFrame + Lake polygons intersected with the model grid (see :func:`_aggregate_lake_cells`). + min_area_fraction : float, optional + Minimum lake coverage for a cell to be carved/bounded, by default 0.5. + + Returns + ------- + flopy.mf6.ModflowGwfriv + The lake RIV package. When ``ds.transport`` is set, its package name is appended to + ``ds.attrs['ssm_sources']`` (if absent) so its CONCENTRATION aux is used by SSM. + """ + agg = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction).rename( + columns={"strt": "stage", "botm": "rbot", "identificatie": "boundname"} + ) + agg["aux"] = 0.0 + + riv_spd = nlmod.gwf.build_spd(agg, "RIV", ds, layer_method="lay_of_rbot") + + riv = flopy.mf6.ModflowGwfriv( + gwf, + auxiliary="CONCENTRATION", + boundnames=True, + stress_period_data={0: riv_spd}, + save_flows=True, + pname="riv_lake", + ) + if ds.transport: + ssm_sources = list(ds.attrs.get("ssm_sources", [])) + if riv.package_name not in ssm_sources: + ds.attrs["ssm_sources"] = [*ssm_sources, riv.package_name] + return riv + + +def recharge_pond_mask(ds, panden_riv=None): + """Mark cells whose meteoric input is carried by a stage boundary, not by RCH. + + Areal recharge is applied to the top active cell of every non-sea cell, which + double-counts precipitation on the managed panden (already implicit in the prescribed + RIV stage) and applies land P-E to the carved open-water lake cells. This mask flags + those cells so the RCH package can exclude them. + + Parameters + ---------- + ds : xarray.Dataset + Model dataset carrying the boolean ``ds['lake_cell']`` marker (see + :func:`carve_lake_cells`). + panden_riv : flopy.mf6.ModflowGwfriv or None, optional + The infiltration-panden RIV package. When None (the default, e.g. the Bergen extent + where the panden lie outside the grid) the mask is exactly ``ds['lake_cell']``. + + Returns + ------- + xarray.DataArray + Boolean mask over ``icell2d``: True at carved lake cells and at panden RIV cells. + """ + mask = ds["lake_cell"].copy() + if panden_riv is not None: + cells = np.unique([cid[-1] for cid in panden_riv.stress_period_data.data[0]["cellid"]]) + mask.loc[{"icell2d": cells}] = True + return mask diff --git a/src/nhflotools/polder.py b/src/nhflotools/polder.py index 8b32a4d..e163305 100644 --- a/src/nhflotools/polder.py +++ b/src/nhflotools/polder.py @@ -1,3 +1,5 @@ +"""Regional polder drainage (HHNK peilgebieden) as a MODFLOW 6 DRN package.""" + import itertools from collections import Counter, defaultdict @@ -7,7 +9,7 @@ import xarray as xr -def drn_from_waterboard_data(ds, gwf, wb="Hollands Noorderkwartier", cbot=1.0): +def drn_from_waterboard_data(ds, gwf, wb="Hollands Noorderkwartier", cbot=1.0, exclude=None): """Create DRN package from waterboard data. Het oppervlaktewater in de polders is vlakdekkend geschematiseerd op basis van @@ -31,6 +33,13 @@ def drn_from_waterboard_data(ds, gwf, wb="Hollands Noorderkwartier", cbot=1.0): cbot : float, optional Bottom resistance of the drains [days], by default 1.0. The per-cell conductance is ``cell_area / cbot``. + exclude : xarray.DataArray or None, optional + Boolean mask over ``icell2d`` marking cells at which no DRN reach should be + emitted. Both the conductance and the elevation are nulled at these cells, so a + reach is dropped whether its stage comes from HHNK peilgebied data or from the + maaiveld fallback (``nlmod.gwf.drn`` masks on ``cond > 0``). The default None + applies no exclusion and is fully backward-compatible. Used to hand carved lake + cells over to a dedicated stage boundary. Returns ------- @@ -81,4 +90,11 @@ def drn_from_waterboard_data(ds, gwf, wb="Hollands Noorderkwartier", cbot=1.0): ds["drn_elev"] = drn_elev ds["drn_cond"] = drn_cond + if exclude is not None: + # Null both fields at excluded cells. nlmod.gwf.drn builds a reach wherever + # cond > 0, so nulling drn_cond drops the reach whether its stage came from the + # celldata assignment or the maaiveld fallback above. + ds["drn_cond"] = ds["drn_cond"].where(~exclude) + ds["drn_elev"] = ds["drn_elev"].where(~exclude) + return nlmod.gwf.drn(ds, gwf, elev="drn_elev", cond="drn_cond") diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e92f4f9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,97 @@ +"""Shared synthetic disv fixtures for the lake/polder tests. + +These build a tiny vertex (disv) MODFLOW 6 model plus a matching xarray dataset, so the +lake carve/stage/mask helpers can be exercised without a solver, real PWN data, or a +network connection. +""" + +import flopy +import numpy as np +import pytest +import xarray as xr +from shapely.geometry import box + + +def _build_disv(nrow=3, ncol=3, top=5.0, dz=10.0, transport=0): + """Build an independent (ds, gwf, cell_geometries) synthetic disv model. + + Parameters + ---------- + nrow, ncol : int, optional + Number of rows/columns of 200 m square cells. The default is a 3x3 grid. + top : float, optional + Uniform model top, by default 5.0 m NAP. + dz : float, optional + Uniform layer thickness, by default 10.0 m (two layers). + transport : int, optional + Value stored under ``ds.attrs['transport']``, by default 0 (no transport). + + Returns + ------- + ds : xarray.Dataset + Dataset with ``top``, ``botm``, ``kh``, ``idomain`` and ``area``. + gwf : flopy.mf6.ModflowGwf + Groundwater flow model with a matching disv grid and a single stress period. + geoms : list of shapely.geometry.Polygon + Cell footprints, indexed by ``icell2d``. + """ + delr = delc = 200.0 + nlay = 2 + ncpl = nrow * ncol + + verts = [] + vid = {} + k = 0 + for j in range(nrow + 1): + for i in range(ncol + 1): + vid[(j, i)] = k + verts.append([k, i * delr, (nrow - j) * delc]) + k += 1 + + cell2d = [] + xc, yc, geoms = [], [], [] + for r in range(nrow): + for c in range(ncol): + icpl = r * ncol + c + v = [vid[(r, c)], vid[(r, c + 1)], vid[(r + 1, c + 1)], vid[(r + 1, c)]] + cx, cy = (c + 0.5) * delr, (nrow - r - 0.5) * delc + xc.append(cx) + yc.append(cy) + cell2d.append([icpl, cx, cy, 4, *v]) + geoms.append(box(c * delr, (nrow - r - 1) * delc, (c + 1) * delr, (nrow - r) * delc)) + + botm = np.array([[top - dz * (lay + 1)] * ncpl for lay in range(nlay)], dtype=float) + + sim = flopy.mf6.MFSimulation(sim_name="t", exe_name="mf6") + flopy.mf6.ModflowTdis(sim, nper=1, perioddata=[(1.0, 1, 1.0)]) + gwf = flopy.mf6.ModflowGwf(sim, modelname="t") + flopy.mf6.ModflowGwfdisv( + gwf, nlay=nlay, ncpl=ncpl, nvert=len(verts), top=top, botm=botm, vertices=verts, cell2d=cell2d + ) + + ds = xr.Dataset( + data_vars={ + "top": ("icell2d", np.full(ncpl, top, dtype=float)), + "botm": (("layer", "icell2d"), botm), + "kh": (("layer", "icell2d"), np.full((nlay, ncpl), 10.0)), + "idomain": (("layer", "icell2d"), np.ones((nlay, ncpl), dtype=int)), + "area": ("icell2d", np.full(ncpl, delr * delc, dtype=float)), + }, + coords={ + "layer": np.arange(nlay), + "icell2d": np.arange(ncpl), + "x": ("icell2d", np.array(xc)), + "y": ("icell2d", np.array(yc)), + }, + ) + ds.attrs["gridtype"] = "vertex" + ds.attrs["transport"] = transport + ds.attrs["extent"] = [0.0, ncol * delr, 0.0, nrow * delc] + ds.attrs["ssm_sources"] = [] + return ds, gwf, geoms + + +@pytest.fixture +def disv_grid(): + """Return the :func:`_build_disv` factory (each call builds a fresh, independent model).""" + return _build_disv diff --git a/tests/test_lakes.py b/tests/test_lakes.py new file mode 100644 index 0000000..f0c8bdb --- /dev/null +++ b/tests/test_lakes.py @@ -0,0 +1,157 @@ +"""Tests for nhflotools.lakes. + +The Bergen pond/lake carve, the per-lake stage RIV and the recharge pond-mask are all +derived from the single aggregator ``_aggregate_lake_cells``, so the carved-cell set and the +stage-reach set are equal by construction. These tests pin that identity, the strict 50% +coverage threshold, the per-cell collapse of overlapping lake pieces, and the recharge mask. +""" + +import logging +import types + +import geopandas as gpd +import numpy as np +import pytest +from shapely.geometry import box + +from nhflotools import lakes + +# Cell ids of the shared 3x3 disv grid used throughout (icell2d, row-major from the top-left). +CELL_40 = 0 # 40% lake coverage -> not carved +CELL_60 = 1 # 60% lake coverage -> carved +CELL_50 = 2 # exactly 50% coverage -> not carved (strict '>') +CELL_FULL = 4 # centre cell, fully covered -> carved +PANDEN_A = 3 # stub panden RIV cell +PANDEN_B = 5 # stub panden RIV cell + +DEEPEST_BOTM = 0.2 # minimum botm across the overlapping pieces of CELL_FULL + + +def _piece(cellid, geom, strt, botm, clake=10.0, ident="lake"): + """Build one lake-piece record already assigned to a grid cell.""" + return { + "cellid": cellid, + "geometry": geom, + "strt": strt, + "botm": botm, + "clake": clake, + "identificatie": ident, + } + + +def _lake_gdf(rows): + """Build a lake gdf (as gdf_to_grid would return) from piece dicts.""" + return gpd.GeoDataFrame(rows, geometry="geometry") + + +def test_carve_selection_min_area_fraction(disv_grid): + """Coverage uses a strict '>' threshold and the carved bottom is the per-cell minimum.""" + ds, _gwf, _geoms = disv_grid() + # Cell 0 at 40%, cell 1 at 60%, cell 2 at exactly 50%, cell 4 fully covered by two pieces. + rows = [ + _piece(CELL_40, box(0, 400, 80, 600), 2.0, 1.0), # 16000 / 40000 = 0.40 -> excluded + _piece(CELL_60, box(200, 400, 320, 600), 2.0, 1.0), # 24000 / 40000 = 0.60 -> carved + _piece(CELL_50, box(400, 400, 500, 600), 2.0, 1.0), # 20000 / 40000 = 0.50 -> excluded + _piece(CELL_FULL, box(200, 200, 400, 300), 2.0, 1.0, ident="deep"), # 20000, botm 1.0 + _piece(CELL_FULL, box(200, 300, 400, 400), 2.0, DEEPEST_BOTM, ident="deeper"), # 20000, botm 0.2 + ] + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, _lake_gdf(rows), min_area_fraction=0.5) + + assert set(lake_cellids.tolist()) == {CELL_60, CELL_FULL} + assert CELL_40 not in lake_cellids # 40% not carved + assert CELL_50 not in lake_cellids # exactly 50% not carved -> pins strict '>' + # the carved top is the minimum botm across the overlapping pieces of the cell + assert ds_carved["top"].sel(icell2d=CELL_FULL).item() == pytest.approx(DEEPEST_BOTM) + + +def test_carve_and_stage_sets_identical(disv_grid): + """A >50%-covered cell whose only piece lacks a stage is dropped from BOTH sets.""" + ds, gwf, geoms = disv_grid(transport=0) + rows = [ + _piece(CELL_FULL, geoms[CELL_FULL], 2.5, 0.0, ident="lake_full"), # has stage -> both sets + _piece(CELL_40, geoms[CELL_40], np.nan, 1.0, ident="botm_only"), # NO stage -> neither set + ] + gdf = _lake_gdf(rows) + + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + riv = lakes.riv_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + stage_cells = {rec["cellid"][-1] for rec in riv.stress_period_data.data[0]} + + # The unified dropna(['strt', 'botm']) filter drops the botm-only cell from both sets, so the + # script's `stage_cells == lake_cellids` consistency assertion cannot fire on a valid build. + assert stage_cells == set(lake_cellids.tolist()) + assert CELL_40 not in stage_cells + assert CELL_40 not in set(lake_cellids.tolist()) + assert set(lake_cellids.tolist()) == {CELL_FULL} + + +def test_lake_riv_stage_boundary_on_synthetic_grid(disv_grid, caplog): + """Single- and multi-piece cells yield exactly one reach each with the expected values.""" + ds, gwf, geoms = disv_grid(transport=1) + # cell 1: two overlapping pieces, areas 24000 (strt 2.0, botm 0.5) and 12000 (strt 3.0, botm 0.0) + area_a, area_b = 24000.0, 12000.0 + strt_a, strt_b = 2.0, 3.0 + clake = 10.0 + rows = [ + _piece(CELL_FULL, geoms[CELL_FULL], 2.5, 0.0, clake=clake, ident="lake_single"), + _piece(CELL_60, box(200, 400, 320, 600), strt_a, 0.5, clake=clake, ident="piece_a"), + _piece(CELL_60, box(320, 400, 400, 550), strt_b, 0.0, clake=clake, ident="piece_b"), + ] + gdf = _lake_gdf(rows) + ds_carved, _lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + + with caplog.at_level(logging.WARNING): + riv = lakes.riv_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + + recs = {rec["cellid"][-1]: rec for rec in riv.stress_period_data.data[0]} + assert set(recs) == {CELL_60, CELL_FULL} # one reach per cell, no duplicate cellids + + # single-piece cell: stage/rbot/cond and top-active-layer placement + single = recs[CELL_FULL] + assert single["cellid"][0] == 0 # lands in the top active layer -> pins lay_of_rbot + assert single["stage"] == pytest.approx(2.5) + assert single["rbot"] == pytest.approx(0.0) + assert single["cond"] == pytest.approx(40000.0 / clake) # piece_area / clake + + # two-piece cell: min rbot, summed cond, area-weighted stage + two = recs[CELL_60] + assert two["rbot"] == pytest.approx(0.0) # min(0.5, 0.0) + assert two["cond"] == pytest.approx(area_a / clake + area_b / clake) # summed + assert two["stage"] == pytest.approx((area_a * strt_a + area_b * strt_b) / (area_a + area_b)) + + # no "stage below bottom elevation" / "records without a stage" warnings on valid data + messages = " ".join(rec.getMessage().lower() for rec in caplog.records) + assert "stage below bottom" not in messages + assert "without a stage" not in messages + + # transport build registers the RIV as an SSM source exactly once + assert riv.package_name in ds_carved.attrs["ssm_sources"] + assert ds_carved.attrs["ssm_sources"].count(riv.package_name) == 1 + + +def test_pond_mask_excludes_recharge_cells(disv_grid): + """The mask equals lake_cell in the None branch and adds panden cells via cid[-1].""" + ds, _gwf, geoms = disv_grid() + gdf = _lake_gdf([_piece(CELL_60, geoms[CELL_60], 2.0, 1.0), _piece(CELL_FULL, geoms[CELL_FULL], 2.0, 1.0)]) + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + assert set(lake_cellids.tolist()) == {CELL_60, CELL_FULL} + + # None branch (the default Bergen build): mask is exactly lake_cell + mask_none = lakes.recharge_pond_mask(ds_carved, None) + assert bool((mask_none == ds_carved["lake_cell"]).all()) + + # stub panden RIV whose cellid entries are (layer, icell2d) tuples (with a duplicate cell) + stub = types.SimpleNamespace( + stress_period_data=types.SimpleNamespace(data={0: {"cellid": [(0, PANDEN_A), (0, PANDEN_B), (0, PANDEN_A)]}}) + ) + mask = lakes.recharge_pond_mask(ds_carved, stub) + pond_cells = set(ds_carved["icell2d"].values[mask.values].tolist()) + expected = {CELL_60, CELL_FULL, PANDEN_A, PANDEN_B} + assert pond_cells == expected # lake cells + panden cells via cid[-1] + + northsea = ds_carved["top"].astype(int) * 0 # northsea == 0 everywhere -> all land + masked_in = (northsea == 0) & ~mask + # no lake or panden cell survives the recharge mask + assert not bool(masked_in.sel(icell2d=sorted(expected)).any()) + # masked-in count drops by exactly |lake union panden| relative to the northsea-only mask + assert int((northsea == 0).sum()) - int(masked_in.sum()) == len(expected) diff --git a/tests/test_polder.py b/tests/test_polder.py new file mode 100644 index 0000000..ba4876f --- /dev/null +++ b/tests/test_polder.py @@ -0,0 +1,59 @@ +"""Tests for nhflotools.polder. + +The polder DRN gains an optional ``exclude`` mask so carved lake cells can be handed over +to a dedicated stage boundary. The reach must be dropped whether its stage comes from real +HHNK peilgebied data or from the maaiveld fallback, which is why the exclusion nulls the +conductance after the celldata assignment rather than only editing the fallback mask. +""" + +import geopandas as gpd +import nlmod +import xarray as xr +from shapely.geometry import box + +from nhflotools import polder + +# Cell ids of the shared 3x3 disv grid. +CELL_REAL_STAGE = 4 # centre cell, covered by the mocked HHNK level area (real peilgebied stage) +CELL_FALLBACK = 0 # a cell with no HHNK stage -> maaiveld fallback drain + + +def _reach_cells(drn): + """Return the set of icell2d values that carry a DRN reach.""" + if drn is None: + return set() + return {rec["cellid"][-1] for rec in drn.stress_period_data.data[0]} + + +def _run_drn(monkeypatch, ds, gwf, exclude): + """Run drn_from_waterboard_data with a real HHNK level area over the centre cell mocked in.""" + ds["ahn"] = xr.full_like(ds["top"], 1.0) + ds["northsea"] = xr.zeros_like(ds["top"]).astype(int) + level_areas = gpd.GeoDataFrame( + {"summer_stage": [0.5], "winter_stage": [0.3]}, + geometry=[box(200, 200, 400, 400)], # exactly the centre cell of the 3x3 grid + index=["area_a"], + ) + monkeypatch.setattr(nlmod.read.waterboard, "download_data", lambda **_kw: level_areas.copy()) + return polder.drn_from_waterboard_data(ds=ds, gwf=gwf, cbot=1.0, exclude=exclude) + + +def test_drn_fallback_excludes_lake_cells(disv_grid, monkeypatch): + """Exclude drops the reach at a lake cell carrying real HHNK stage; baseline keeps it.""" + # Baseline (no exclude): the centre cell carries real HHNK peilgebied stage, so it gets a reach. + ds, gwf, _geoms = disv_grid() + drn_base = _run_drn(monkeypatch, ds, gwf, exclude=None) + base_cells = _reach_cells(drn_base) + assert CELL_REAL_STAGE in base_cells # pin: without exclude the real-stage lake cell IS drained + + # With exclude: the reach at the centre cell is dropped even though its stage is real (not + # fallback), while a non-excluded fallback cell still gets a reach. + ds2, gwf2, _geoms2 = disv_grid() + exclude = xr.zeros_like(ds2["top"]).astype(bool) + exclude.loc[{"icell2d": CELL_REAL_STAGE}] = True + drn_excl = _run_drn(monkeypatch, ds2, gwf2, exclude=exclude) + excl_cells = _reach_cells(drn_excl) + + assert CELL_REAL_STAGE not in excl_cells + assert CELL_FALLBACK in excl_cells # a fallback (maaiveld) cell is unaffected + assert excl_cells == base_cells - {CELL_REAL_STAGE} From 7881dab3d1def2d728647ed91659d1277c879a60 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Sun, 19 Jul 2026 07:26:10 +0200 Subject: [PATCH 2/4] Guard empty lake set in nhflotools.lakes _aggregate_lake_cells operated on an empty geometry column when no cell cleared the coverage threshold (empty gdf, every piece missing strt/botm, or all below min_area_fraction), raising 'can only use area methods with polygon geometries'. Return an empty cellid-indexed frame instead, so carve_lake_cells carves nothing and riv_from_lakes_pwn returns None. Add regression tests for both empty paths. --- src/nhflotools/lakes.py | 38 ++++++++++++++++++++++++++------------ tests/test_lakes.py | 25 +++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 12 deletions(-) diff --git a/src/nhflotools/lakes.py b/src/nhflotools/lakes.py index 5feeb97..81c02bf 100644 --- a/src/nhflotools/lakes.py +++ b/src/nhflotools/lakes.py @@ -16,6 +16,7 @@ import flopy import nlmod import numpy as np +import pandas as pd import xarray as xr logger = logging.getLogger(__name__) @@ -47,21 +48,31 @@ def _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): pandas.DataFrame Indexed by ``cellid`` with columns ``strt`` (area-weighted), ``botm`` (minimum), ``cond`` (summed ``piece_area / clake``, units m2/d) and ``identificatie`` - (first). Empty-input semantics follow :func:`nlmod.grid.aggregate_vector_per_cell`. + (first). When no cell clears the coverage threshold (empty input, every piece + missing ``strt``/``botm``, or all cells below ``min_area_fraction``), an empty + ``cellid``-indexed frame with those columns is returned so the callers carve + nothing and emit no RIV rather than raising on an empty geometry column. """ lake = gdf_lake_grid.dropna(subset=["strt", "botm"]).copy() n_drop = len(gdf_lake_grid) - len(lake) if n_drop: logger.warning("Dropping %d lake piece(s) missing strt or botm", n_drop) - # Use the true clipped-piece area throughout (coverage, conductance and the - # area-weighted stage), so the three are mutually consistent regardless of any stale - # 'area' column carried over from gdf_to_grid. - lake["area"] = lake.geometry.area - overlap = lake["area"].groupby(lake["cellid"]).transform("sum").to_numpy() - cell_area = ds["area"].sel(icell2d=lake["cellid"].to_numpy()).to_numpy() - # Strict '>' so a cell exactly at min_area_fraction is excluded. - lake = lake[overlap / cell_area > min_area_fraction] + if not lake.empty: + # Use the true clipped-piece area throughout (coverage, conductance and the + # area-weighted stage), so the three are mutually consistent regardless of any stale + # 'area' column carried over from gdf_to_grid. + lake["area"] = lake.geometry.area + overlap = lake["area"].groupby(lake["cellid"]).transform("sum").to_numpy() + cell_area = ds["area"].sel(icell2d=lake["cellid"].to_numpy()).to_numpy() + # Strict '>' so a cell exactly at min_area_fraction is excluded. + lake = lake[overlap / cell_area > min_area_fraction] + + if lake.empty: + return pd.DataFrame( + {"strt": [], "botm": [], "cond": [], "identificatie": []}, + index=pd.Index([], name="cellid", dtype="int64"), + ) lake["cond"] = lake["area"] / lake["clake"] return nlmod.grid.aggregate_vector_per_cell( @@ -143,13 +154,16 @@ def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): Returns ------- - flopy.mf6.ModflowGwfriv - The lake RIV package. When ``ds.transport`` is set, its package name is appended to - ``ds.attrs['ssm_sources']`` (if absent) so its CONCENTRATION aux is used by SSM. + flopy.mf6.ModflowGwfriv or None + The lake RIV package, or ``None`` when no cell clears the coverage threshold (so the + caller adds no package). When ``ds.transport`` is set, the package name is appended + to ``ds.attrs['ssm_sources']`` (if absent) so its CONCENTRATION aux is used by SSM. """ agg = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction).rename( columns={"strt": "stage", "botm": "rbot", "identificatie": "boundname"} ) + if agg.empty: + return None agg["aux"] = 0.0 riv_spd = nlmod.gwf.build_spd(agg, "RIV", ds, layer_method="lay_of_rbot") diff --git a/tests/test_lakes.py b/tests/test_lakes.py index f0c8bdb..c554dc3 100644 --- a/tests/test_lakes.py +++ b/tests/test_lakes.py @@ -155,3 +155,28 @@ def test_pond_mask_excludes_recharge_cells(disv_grid): assert not bool(masked_in.sel(icell2d=sorted(expected)).any()) # masked-in count drops by exactly |lake union panden| relative to the northsea-only mask assert int((northsea == 0).sum()) - int(masked_in.sum()) == len(expected) + + +def test_carve_lake_cells_empty_when_no_cell_clears_threshold(disv_grid): + """An empty / all-below-threshold lake gdf carves nothing instead of crashing. + + Regression for the empty-geometry ``.area`` crash: ``carve_lake_cells`` previously + raised ``TypeError: can only use area methods with polygon geometries`` whenever the + coverage filter (or an empty input) left no surviving lake piece. + """ + ds, _gwf, _geoms = disv_grid() + sliver = _lake_gdf([_piece(CELL_40, box(0, 400, 20, 420), 2.0, 1.0)]) # 400/40000 = 1% < 50% + empty = _lake_gdf([_piece(CELL_40, box(0, 400, 80, 600), 2.0, 1.0)]).iloc[0:0] + for gdf in (sliver, empty): + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + assert lake_cellids.tolist() == [] + assert not bool(ds_carved["lake_cell"].any()) + assert bool((ds_carved["top"] == ds["top"]).all()) # nothing carved + + +def test_riv_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): + """The per-lake RIV is ``None`` (no package added) when no cell clears the threshold.""" + ds, gwf, _geoms = disv_grid() + sliver = _lake_gdf([_piece(CELL_40, box(0, 400, 20, 420), 2.0, 1.0)]) + ds_carved, _ = lakes.carve_lake_cells(ds, sliver, min_area_fraction=0.5) + assert lakes.riv_from_lakes_pwn(ds_carved, gwf, sliver, min_area_fraction=0.5) is None From 44a9b1ac96910681a22c769c3d3306df7e02e5e3 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Tue, 21 Jul 2026 18:18:34 +0200 Subject: [PATCH 3/4] Aggregate lake cells per lake, add fractional recharge mask and LAK variant Responds to the review of #58: - _aggregate_lake_cells now returns one record per (cellid, identificatie) plus a per-cell coverage Series. Lakes sharing a grid cell are no longer merged: each keeps its own RIV reach with its own stage, bed (rbot), conductance and boundname, so per-lake budgets stay attributable and a future MVR mover/weir or LAK configuration can address lakes individually. The cell is still carved once, to the deepest bed among its lakes. - carve_lake_cells stores ds['lake_coverage'] (combined stage-carrying lake coverage per cell, also below the carve threshold) next to ds['lake_cell'], and documents its lifecycle contract: called after the layer model is finalized, lowering-only (see the consolidated adjust-top note on NHFLO/models#126). - recharge_pond_mask gains a keyword-only fractional=True mode: 1.0 at carved lake and panden cells, the open-water coverage fraction at cells below the carve threshold, so recharge can be scaled by 1 - fraction instead of the all-or-nothing boolean exclusion. - lak_from_lakes_pwn builds the LAK alternative from the same aggregator, with an effective bed resistance that reproduces the RIV summed-piece conductance and from_ds meteorology. Verified: recharge double-counting is impossible in the current build -- KNMI recharge (method='linear') nets Makkink evaporation into the single recharge variable and no EVT package is built, so masking a cell removes its meteoric term exactly once. Tests: regression for the cross-lake merge (fails on the previous aggregation, including the per-lake rbot in a shared cell), fractional mask, LAK connection/conductance equivalence; suite mutation-tested (13 kernel mutations, all caught). --- src/nhflotools/lakes.py | 295 +++++++++++++++++++++++++++++++--------- tests/test_lakes.py | 222 ++++++++++++++++++++++++++++-- 2 files changed, 439 insertions(+), 78 deletions(-) diff --git a/src/nhflotools/lakes.py b/src/nhflotools/lakes.py index 81c02bf..53bfd2f 100644 --- a/src/nhflotools/lakes.py +++ b/src/nhflotools/lakes.py @@ -1,11 +1,16 @@ """Helpers for the Bergen pond/lake (``lakes_pwn``) features. -The PWN model carves the model top down to the lake bottom in every grid cell that a -managed pond/lake sufficiently covers, and then holds those carved cells at their -prescribed lake stage with a per-lake RIV boundary. Both the carved cell set and the RIV -reach set are derived from the single aggregator :func:`_aggregate_lake_cells`, so the two -sets are equal by construction: no carved cell is ever left without a stage, and no stage -reach ever lands on an un-carved cell. +The PWN model carves the model top down to the lake bottom in every grid cell that the +managed ponds/lakes together sufficiently cover, and then holds each carved cell at the +prescribed stage of its lake(s) with a per-lake RIV boundary: every lake keeps its own +reach, stage, conductance and boundname, also where two lakes share a cell, so per-lake +budgets stay attributable and a future MVR/weir or LAK configuration can address lakes +individually. Both the carved cell set and the RIV reach set are derived from the single +aggregator :func:`_aggregate_lake_cells`, so the two sets are equal by construction: no +carved cell is ever left without a stage, and no stage reach ever lands on an un-carved +cell. For detailed studies :func:`lak_from_lakes_pwn` builds the LAK package from the +same aggregator instead, solving the stage from the lake water balance with the same +lakebed leakance as the RIV. The logic lives here (rather than inline in the model script) so it can be imported and unit-tested without running the full REGIS/AHN/MF6 build. @@ -23,13 +28,14 @@ def _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): - """Aggregate lake pieces to one carved/stage record per grid cell. + """Aggregate lake pieces to one carved/stage record per lake per grid cell. This is the single source of truth for the carved-cell set. It (1) keeps only lake pieces that carry both a stage (``strt``) and a bottom (``botm``), (2) computes each - cell's lake coverage over exactly those pieces, (3) keeps cells whose coverage is - strictly greater than ``min_area_fraction``, and (4) collapses the surviving pieces to - one record per cell. + cell's combined lake coverage over exactly those pieces, (3) keeps cells whose combined + coverage is strictly greater than ``min_area_fraction``, and (4) collapses the surviving + pieces to one record per lake (``identificatie``) per cell — lakes are never merged + across ``identificatie``, so each keeps its own stage and boundname. Parameters ---------- @@ -40,58 +46,79 @@ def _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): :func:`nlmod.dims.gdf_to_grid`). Must carry a ``cellid`` column and the ``strt``, ``botm``, ``clake`` and ``identificatie`` columns of ``lakes_pwn``. min_area_fraction : float, optional - A cell is kept when the lake covers strictly more than this fraction of its area. - The default is 0.5, so a cell exactly half covered is not carved. + A cell is kept when the lakes combined cover strictly more than this fraction of + its area. The default is 0.5, so a cell exactly half covered is not carved. Returns ------- - pandas.DataFrame - Indexed by ``cellid`` with columns ``strt`` (area-weighted), ``botm`` (minimum), - ``cond`` (summed ``piece_area / clake``, units m2/d) and ``identificatie`` - (first). When no cell clears the coverage threshold (empty input, every piece - missing ``strt``/``botm``, or all cells below ``min_area_fraction``), an empty - ``cellid``-indexed frame with those columns is returned so the callers carve + agg : pandas.DataFrame + One row per (cell, lake), indexed by ``cellid`` (repeated when lakes share a + cell), with columns ``strt`` (area-weighted within the lake), ``botm`` (minimum + within the lake), ``cond`` (summed ``piece_area / clake``, units m2/d) and + ``identificatie``. When no cell clears the coverage threshold (empty input, every + piece missing ``strt``/``botm``, or all cells below ``min_area_fraction``), an + empty ``cellid``-indexed frame with those columns is returned so the callers carve nothing and emit no RIV rather than raising on an empty geometry column. + coverage : pandas.Series + Combined stage-carrying lake coverage fraction per ``cellid``, *before* the + ``min_area_fraction`` threshold, for every cell any surviving piece touches. Used + for the fractional recharge mask. """ lake = gdf_lake_grid.dropna(subset=["strt", "botm"]).copy() n_drop = len(gdf_lake_grid) - len(lake) if n_drop: logger.warning("Dropping %d lake piece(s) missing strt or botm", n_drop) - if not lake.empty: - # Use the true clipped-piece area throughout (coverage, conductance and the - # area-weighted stage), so the three are mutually consistent regardless of any stale - # 'area' column carried over from gdf_to_grid. - lake["area"] = lake.geometry.area - overlap = lake["area"].groupby(lake["cellid"]).transform("sum").to_numpy() - cell_area = ds["area"].sel(icell2d=lake["cellid"].to_numpy()).to_numpy() - # Strict '>' so a cell exactly at min_area_fraction is excluded. - lake = lake[overlap / cell_area > min_area_fraction] + empty_agg = pd.DataFrame( + {"strt": [], "botm": [], "cond": [], "identificatie": []}, + index=pd.Index([], name="cellid", dtype="int64"), + ) + empty_coverage = pd.Series([], index=pd.Index([], name="cellid", dtype="int64"), dtype="float64") + if lake.empty: + return empty_agg, empty_coverage + + # Use the true clipped-piece area throughout (coverage, conductance and the + # area-weighted stage), so the three are mutually consistent regardless of any stale + # 'area' column carried over from gdf_to_grid. + lake["area"] = lake.geometry.area + overlap = lake["area"].groupby(lake["cellid"]).sum() + cell_area = ds["area"].sel(icell2d=overlap.index.to_numpy()).to_numpy() + coverage = overlap / cell_area + # Strict '>' so a cell exactly at min_area_fraction is excluded. + lake = lake[lake["cellid"].map(coverage) > min_area_fraction] if lake.empty: - return pd.DataFrame( - {"strt": [], "botm": [], "cond": [], "identificatie": []}, - index=pd.Index([], name="cellid", dtype="int64"), - ) + return empty_agg, coverage lake["cond"] = lake["area"] / lake["clake"] - return nlmod.grid.aggregate_vector_per_cell( - lake, - fields_methods={ - "strt": "area_weighted", - "botm": "min", - "cond": "sum", - "identificatie": "first", - }, - ) + parts = [] + for ident, group in lake.groupby("identificatie", sort=False, dropna=False): + agg = nlmod.grid.aggregate_vector_per_cell( + group, + fields_methods={"strt": "area_weighted", "botm": "min", "cond": "sum"}, + ) + agg["identificatie"] = ident + parts.append(agg) + agg = pd.concat(parts).sort_index() + agg.index.name = "cellid" + return agg, coverage def carve_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): """Lower the model top to the lake bottom in sufficiently covered cells. - The model top is set to the aggregated lake bottom in every carved cell and a boolean - ``ds['lake_cell']`` marker (True exactly at the carved cells) is added for the recharge - mask, the polder-drain exclusion and the lake RIV to reuse. + The model top is set to the deepest aggregated lake bottom in every carved cell and two + markers are added for the recharge mask, the polder-drain exclusion and the lake RIV to + reuse: a boolean ``ds['lake_cell']`` (True exactly at the carved cells) and a float + ``ds['lake_coverage']`` (combined stage-carrying lake coverage fraction per cell, + including cells below the carve threshold). + + This is called from the model script directly after the layer model is finalized (the + nlmod top from AHN/REGIS) and before ``starting_head`` and any package build. It only + lowers the top: :func:`nlmod.layers.set_model_top` is a one-way ratchet, so when more + carve sources join (bathymetry, panden opt-in) their bed elevations must be merged into + one top with a single ``set_model_top`` call rather than carved incrementally — see the + consolidated adjust-top note on NHFLO/models#126. Parameters ---------- @@ -100,13 +127,15 @@ def carve_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): gdf_lake_grid : geopandas.GeoDataFrame Lake polygons intersected with the model grid (see :func:`_aggregate_lake_cells`). min_area_fraction : float, optional - Minimum lake coverage for a cell to be carved, by default 0.5 (strict ``>``). + Minimum combined lake coverage for a cell to be carved, by default 0.5 (strict + ``>``). Returns ------- ds : xarray.Dataset Copy of the input dataset with the lowered top, ``thickness`` dropped (so it is - recomputed from the new top/botm), and the boolean ``lake_cell`` marker. + recomputed from the new top/botm), and the ``lake_cell`` / ``lake_coverage`` + markers. lake_cellids : numpy.ndarray The ``icell2d`` values of the carved cells. @@ -117,29 +146,41 @@ def carve_lake_cells(ds, gdf_lake_grid, min_area_fraction=0.5): would be inert; refreshing ``ahn`` to the lake bottom would also corrupt its meaning as the measured maaiveld. """ - agg = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) - lake_cellids = agg.index.to_numpy() + agg, coverage = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) + # Deepest lake bed per cell: lakes sharing a cell carve it to the lower of their beds. + cell_botm = agg["botm"].groupby(level="cellid").min() + lake_cellids = cell_botm.index.to_numpy() top = ds["top"].copy() - top.loc[{"icell2d": lake_cellids}] = agg["botm"].to_numpy() + top.loc[{"icell2d": lake_cellids}] = cell_botm.to_numpy() ds = nlmod.layers.set_model_top(ds, top) ds = ds.drop_vars("thickness", errors="ignore") ds["lake_cell"] = xr.zeros_like(ds["top"], dtype=bool) ds["lake_cell"].loc[{"icell2d": lake_cellids}] = True + # Coverage can nominally exceed 1 when lake polygons overlap each other; clip so the + # fractional recharge mask never removes more than a cell's full meteoric term. + ds["lake_coverage"] = xr.zeros_like(ds["top"], dtype=float) + ds["lake_coverage"].loc[{"icell2d": coverage.index.to_numpy()}] = coverage.clip(upper=1.0).to_numpy() return ds, lake_cellids def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): """Hold the carved lake cells at their prescribed stage with a per-lake RIV. - A RIV reach is placed in every carved cell with ``stage = strt``, ``rbot = botm`` (the - carved top) and ``cond = sum(piece_area / clake)``. Because ``rbot`` equals the carved - top, the reach caps bed infiltration once the head drops below the lakebed - (perched-pond behaviour) while draining freely when the head rises above the stage. + A RIV reach is placed in every carved cell for every lake that covers it, with + ``stage = strt``, ``rbot = botm`` (that lake's bed) and ``cond = sum(piece_area / + clake)``, and the lake's ``identificatie`` as boundname. Because ``rbot`` equals the + lake bed, the reach caps bed infiltration once the head drops below the lakebed + (perched-pond behaviour) while draining freely when the head rises above the stage. In + a cell carved by a single lake (the usual case) ``rbot`` therefore equals the carved + top; where lakes share a cell, the cell is carved to the deepest bed and the shallower + lake's reach keeps its own, higher ``rbot``. Keeping one reach per lake preserves + per-lake budgets and boundnames, so MVR movers and weir outlets between individual + lakes (and a future LAK swap) remain configurable. The reach set is derived from the same :func:`_aggregate_lake_cells` call as - :func:`carve_lake_cells`, so it equals the carved-cell set by construction. + :func:`carve_lake_cells`, so its cell set equals the carved-cell set by construction. Parameters ---------- @@ -150,7 +191,7 @@ def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): gdf_lake_grid : geopandas.GeoDataFrame Lake polygons intersected with the model grid (see :func:`_aggregate_lake_cells`). min_area_fraction : float, optional - Minimum lake coverage for a cell to be carved/bounded, by default 0.5. + Minimum combined lake coverage for a cell to be carved/bounded, by default 0.5. Returns ------- @@ -159,9 +200,8 @@ def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): caller adds no package). When ``ds.transport`` is set, the package name is appended to ``ds.attrs['ssm_sources']`` (if absent) so its CONCENTRATION aux is used by SSM. """ - agg = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction).rename( - columns={"strt": "stage", "botm": "rbot", "identificatie": "boundname"} - ) + agg, _ = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) + agg = agg.rename(columns={"strt": "stage", "botm": "rbot", "identificatie": "boundname"}) if agg.empty: return None agg["aux"] = 0.0 @@ -183,30 +223,157 @@ def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): return riv -def recharge_pond_mask(ds, panden_riv=None): +def lak_from_lakes_pwn( + ds, + gwf, + gdf_lake_grid, + min_area_fraction=0.5, + gwt=None, + rainfall="from_ds", + evaporation="from_ds", + **kwargs, +): + """Model the carved lake cells with the LAK package instead of a stage RIV. + + Alternative to :func:`riv_from_lakes_pwn` for detailed studies: the lake stage + follows from the simulated lake water balance instead of being prescribed, which + resolves stage drawdown, lake storage and (via ``lakeout`` columns) outlets, at the + cost of a harder nonlinear solve. The connection set is derived from the same + :func:`_aggregate_lake_cells` call as :func:`carve_lake_cells`, so it equals the + carved-cell set by construction, and the lakebed leakance matches the RIV variant: + each (cell, lake) row becomes one VERTICAL connection with an effective bed + resistance ``clake_eff = cell_area / sum(piece_area / clake)``, so ``bedleak * + cell_area`` equals the summed piece conductance of the corresponding RIV reach. The + total exchange is nevertheless somewhat weaker than the RIV's, because MF6 places + the connected cell's half-cell vertical resistance (``0.5 * thickness / k33``) in + series with the lakebed for VERTICAL connections, which the RIV formulation applies + directly to the cell node (12-16% lower on the PWN layer model). Like the RIV (and + unlike a GHB), the exchange is capped at the lakebed once the aquifer head drops + below it. + + Parameters + ---------- + ds : xarray.Dataset + Model dataset (post-carve) with ``top``, ``botm``, ``area``, ``idomain`` and a + time discretisation in days. ``recharge`` is required when + ``rainfall='from_ds'`` and ``chloride`` when ``gwt`` is given. + gwf : flopy.mf6.ModflowGwf + Groundwater flow model the LAK package is added to. + gdf_lake_grid : geopandas.GeoDataFrame + Lake polygons intersected with the model grid (see + :func:`_aggregate_lake_cells`). Per-lake outlet columns (``lakeout``, + ``couttype``, ``outlet_invert``, ``outlet_width``, ``outlet_rough``, + ``outlet_slope``) are carried through when present. + min_area_fraction : float, optional + Minimum combined lake coverage for a cell to be connected, by default 0.5 + (strict ``>``), identical to the carve threshold. + gwt : flopy.mf6.ModflowGwt, optional + When given, a matching LKT transport package is created as well and + ``(lak, lkt)`` is returned. The default is None. + rainfall, evaporation : str, float, pandas.DataFrame or None, optional + Passed to :func:`nlmod.gwf.lake.lake_from_gdf`. The default ``'from_ds'`` + derives per-lake area-weighted series from the meteorology in ``ds`` with + :func:`nlmod.gwf.lake.copy_meteorological_data_from_ds`. Exclude the lake cells + from the RCH package (see :func:`recharge_pond_mask`) so the meteoric term is + not counted both on the lake and on the aquifer. + **kwargs + Passed to :func:`nlmod.gwf.lake.lake_from_gdf` and on to + :class:`flopy.mf6.ModflowGwflak` (e.g. ``save_flows``, + ``package_convergence_filerecord``). + + Returns + ------- + flopy.mf6.ModflowGwflak, tuple or None + The LAK package, ``(lak, lkt)`` when ``gwt`` is given, or ``None`` when no cell + clears the coverage threshold (so the caller adds no package). + + Notes + ----- + When two lakes share a carved cell, MF6 applies each lake's RAINFALL over that + lake's full connection (cell) area while RCH excluded the cell only once, so the + meteoric term on such a cell is over-applied. No current ``lakes_pwn`` cell is + shared between lakes. + """ + agg, _ = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) + if agg.empty: + return None + + gdf = agg.drop(columns=["botm"]) + gdf["clake"] = ds["area"].sel(icell2d=gdf.index.to_numpy()).to_numpy() / gdf.pop("cond") + # LAK requires one exact strt per lake; the per-cell area-weighted values can differ + # at floating-point level. + gdf["strt"] = gdf.groupby("identificatie")["strt"].transform("mean") + + outlet_columns = [ + column + for column in ("lakeout", "couttype", "outlet_invert", "outlet_width", "outlet_rough", "outlet_slope") + if column in gdf_lake_grid.columns + ] + if outlet_columns: + per_lake = gdf_lake_grid.drop_duplicates("identificatie").set_index("identificatie")[outlet_columns] + gdf = gdf.join(per_lake, on="identificatie") + + rainfall_from_ds = isinstance(rainfall, str) and rainfall == "from_ds" + evaporation_from_ds = isinstance(evaporation, str) and evaporation == "from_ds" + if rainfall_from_ds or evaporation_from_ds: + ds_rainfall, ds_evaporation = nlmod.gwf.lake.copy_meteorological_data_from_ds( + gdf, ds, boundname_column="identificatie" + ) + if rainfall_from_ds: + rainfall = ds_rainfall + if evaporation_from_ds: + evaporation = ds_evaporation + + return nlmod.gwf.lake.lake_from_gdf( + gwf, + gdf, + ds, + rainfall=rainfall, + evaporation=evaporation, + boundname_column="identificatie", + gwt=gwt, + **kwargs, + ) + + +def recharge_pond_mask(ds, panden_riv=None, *, fractional=False): """Mark cells whose meteoric input is carried by a stage boundary, not by RCH. Areal recharge is applied to the top active cell of every non-sea cell, which double-counts precipitation on the managed panden (already implicit in the prescribed RIV stage) and applies land P-E to the carved open-water lake cells. This mask flags - those cells so the RCH package can exclude them. + those cells so the RCH package can exclude them. Note that ``ds['recharge']`` (KNMI, + ``method='linear'``) already nets Makkink evaporation against precipitation, and no EVT + package is built, so excluding a cell removes its entire meteoric term exactly once. Parameters ---------- ds : xarray.Dataset - Model dataset carrying the boolean ``ds['lake_cell']`` marker (see - :func:`carve_lake_cells`). + Model dataset carrying the ``ds['lake_cell']`` and ``ds['lake_coverage']`` markers + (see :func:`carve_lake_cells`). panden_riv : flopy.mf6.ModflowGwfriv or None, optional The infiltration-panden RIV package. When None (the default, e.g. the Bergen extent - where the panden lie outside the grid) the mask is exactly ``ds['lake_cell']``. + where the panden lie outside the grid) the mask is derived from the lake markers + alone. + fractional : bool, optional + When False (default) return the boolean mask: True at carved lake cells and panden + RIV cells. When True return a float fraction per cell instead: 1.0 at carved lake + cells (the whole cell is modeled as lake bed held by the RIV) and at panden RIV + cells, and the open-water coverage fraction at cells below the carve threshold — + precipitation on such a partial lake sliver feeds the lake (whose balance the + prescribed stage absorbs), not the aquifer, so the caller can scale recharge by + ``1 - fraction`` there instead of keeping the full land P-E. Returns ------- xarray.DataArray - Boolean mask over ``icell2d``: True at carved lake cells and at panden RIV cells. + Over ``icell2d``: boolean exclusion mask (default), or the fraction of each cell's + meteoric term carried by a stage boundary (``fractional=True``). """ mask = ds["lake_cell"].copy() if panden_riv is not None: cells = np.unique([cid[-1] for cid in panden_riv.stress_period_data.data[0]["cellid"]]) mask.loc[{"icell2d": cells}] = True - return mask + if not fractional: + return mask + return xr.where(mask, 1.0, ds["lake_coverage"], keep_attrs=True) diff --git a/tests/test_lakes.py b/tests/test_lakes.py index c554dc3..3e229ef 100644 --- a/tests/test_lakes.py +++ b/tests/test_lakes.py @@ -3,7 +3,8 @@ The Bergen pond/lake carve, the per-lake stage RIV and the recharge pond-mask are all derived from the single aggregator ``_aggregate_lake_cells``, so the carved-cell set and the stage-reach set are equal by construction. These tests pin that identity, the strict 50% -coverage threshold, the per-cell collapse of overlapping lake pieces, and the recharge mask. +coverage threshold, the per-lake-per-cell aggregation (pieces merge within a lake, never +across lakes), and the boolean and fractional recharge masks. """ import logging @@ -11,7 +12,9 @@ import geopandas as gpd import numpy as np +import pandas as pd import pytest +import xarray as xr from shapely.geometry import box from nhflotools import lakes @@ -26,6 +29,13 @@ DEEPEST_BOTM = 0.2 # minimum botm across the overlapping pieces of CELL_FULL +# Shared by the RIV and LAK equivalence tests: one single-piece lake covering the centre +# cell (strt 2.5, botm 0.0) plus one two-piece lake in CELL_60. Both tests assert on the +# same aggregated values, so they must consume the identical input. +CLAKE = 10.0 +AREA_A, AREA_B = 24000.0, 12000.0 # piece areas of the two-piece lake +STRT_A, STRT_B = 2.0, 3.0 # stages of the two pieces (area-weighted in aggregation) + def _piece(cellid, geom, strt, botm, clake=10.0, ident="lake"): """Build one lake-piece record already assigned to a grid cell.""" @@ -44,6 +54,15 @@ def _lake_gdf(rows): return gpd.GeoDataFrame(rows, geometry="geometry") +def _single_and_two_piece_lakes(geoms): + """Build the lake gdf shared by the RIV and LAK equivalence tests.""" + return _lake_gdf([ + _piece(CELL_FULL, geoms[CELL_FULL], 2.5, 0.0, clake=CLAKE, ident="lake_single"), + _piece(CELL_60, box(200, 400, 320, 600), STRT_A, 0.5, clake=CLAKE, ident="lake_two_pieces"), + _piece(CELL_60, box(320, 400, 400, 550), STRT_B, 0.0, clake=CLAKE, ident="lake_two_pieces"), + ]) + + def test_carve_selection_min_area_fraction(disv_grid): """Coverage uses a strict '>' threshold and the carved bottom is the per-cell minimum.""" ds, _gwf, _geoms = disv_grid() @@ -88,16 +107,7 @@ def test_carve_and_stage_sets_identical(disv_grid): def test_lake_riv_stage_boundary_on_synthetic_grid(disv_grid, caplog): """Single- and multi-piece cells yield exactly one reach each with the expected values.""" ds, gwf, geoms = disv_grid(transport=1) - # cell 1: two overlapping pieces, areas 24000 (strt 2.0, botm 0.5) and 12000 (strt 3.0, botm 0.0) - area_a, area_b = 24000.0, 12000.0 - strt_a, strt_b = 2.0, 3.0 - clake = 10.0 - rows = [ - _piece(CELL_FULL, geoms[CELL_FULL], 2.5, 0.0, clake=clake, ident="lake_single"), - _piece(CELL_60, box(200, 400, 320, 600), strt_a, 0.5, clake=clake, ident="piece_a"), - _piece(CELL_60, box(320, 400, 400, 550), strt_b, 0.0, clake=clake, ident="piece_b"), - ] - gdf = _lake_gdf(rows) + gdf = _single_and_two_piece_lakes(geoms) ds_carved, _lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) with caplog.at_level(logging.WARNING): @@ -111,13 +121,13 @@ def test_lake_riv_stage_boundary_on_synthetic_grid(disv_grid, caplog): assert single["cellid"][0] == 0 # lands in the top active layer -> pins lay_of_rbot assert single["stage"] == pytest.approx(2.5) assert single["rbot"] == pytest.approx(0.0) - assert single["cond"] == pytest.approx(40000.0 / clake) # piece_area / clake + assert single["cond"] == pytest.approx(40000.0 / CLAKE) # piece_area / clake # two-piece cell: min rbot, summed cond, area-weighted stage two = recs[CELL_60] assert two["rbot"] == pytest.approx(0.0) # min(0.5, 0.0) - assert two["cond"] == pytest.approx(area_a / clake + area_b / clake) # summed - assert two["stage"] == pytest.approx((area_a * strt_a + area_b * strt_b) / (area_a + area_b)) + assert two["cond"] == pytest.approx(AREA_A / CLAKE + AREA_B / CLAKE) # summed + assert two["stage"] == pytest.approx((AREA_A * STRT_A + AREA_B * STRT_B) / (AREA_A + AREA_B)) # no "stage below bottom elevation" / "records without a stage" warnings on valid data messages = " ".join(rec.getMessage().lower() for rec in caplog.records) @@ -129,6 +139,51 @@ def test_lake_riv_stage_boundary_on_synthetic_grid(disv_grid, caplog): assert ds_carved.attrs["ssm_sources"].count(riv.package_name) == 1 +def test_two_lakes_sharing_a_cell_get_separate_reaches(disv_grid): + """Two different lakes overlapping one cell each keep their own reach, stage and boundname. + + Regression for the per-cell collapse: merging lakes within a cell blends their stages + (each lake holds its own prescribed stage) and destroys the per-lake boundname needed to + configure MVR movers and weir outlets between individual lakes. + """ + ds, gwf, _geoms = disv_grid(transport=0) + clake = 10.0 + strt_a, strt_b = 2.0, 3.0 + # Lake A covers the west half of the centre cell, lake B the east half; both also cover a + # cell of their own so each lake exists independently of the shared cell. + rows = [ + _piece(CELL_FULL, box(200, 200, 300, 400), strt_a, 1.0, clake=clake, ident="lake_a"), + _piece(CELL_FULL, box(300, 200, 400, 400), strt_b, 0.5, clake=clake, ident="lake_b"), + _piece(CELL_60, box(200, 400, 400, 600), strt_a, 1.0, clake=clake, ident="lake_a"), + _piece(PANDEN_B, box(400, 200, 600, 400), strt_b, 0.5, clake=clake, ident="lake_b"), + ] + gdf = _lake_gdf(rows) + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + + # the shared cell is carved once (combined coverage 100%), to the deepest lake bed + assert CELL_FULL in lake_cellids + assert ds_carved["top"].sel(icell2d=CELL_FULL).item() == pytest.approx(0.5) + + riv = lakes.riv_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + recs = riv.stress_period_data.data[0] + shared = [rec for rec in recs if rec["cellid"][-1] == CELL_FULL] + + # one reach per lake in the shared cell, not one blended reach per cell + assert sorted(rec["boundname"] for rec in shared) == ["lake_a", "lake_b"] + by_name = {rec["boundname"]: rec for rec in shared} + # each reach keeps its own lake's stage and conductance (half a 40000 m2 cell each) + assert by_name["lake_a"]["stage"] == pytest.approx(strt_a) + assert by_name["lake_b"]["stage"] == pytest.approx(strt_b) + assert by_name["lake_a"]["cond"] == pytest.approx(20000.0 / clake) + assert by_name["lake_b"]["cond"] == pytest.approx(20000.0 / clake) + # each reach also keeps its own lake's bed: the cell is carved to the deepest bed (0.5), + # but lake_a's perched-infiltration cap stays at its own higher bed (rbot 1.0) + assert by_name["lake_a"]["rbot"] == pytest.approx(1.0) + assert by_name["lake_b"]["rbot"] == pytest.approx(0.5) + # the stage-cell set still equals the carved-cell set (identity by construction) + assert {rec["cellid"][-1] for rec in recs} == set(lake_cellids.tolist()) + + def test_pond_mask_excludes_recharge_cells(disv_grid): """The mask equals lake_cell in the None branch and adds panden cells via cid[-1].""" ds, _gwf, geoms = disv_grid() @@ -157,6 +212,35 @@ def test_pond_mask_excludes_recharge_cells(disv_grid): assert int((northsea == 0).sum()) - int(masked_in.sum()) == len(expected) +def test_pond_mask_fractional(disv_grid): + """``fractional=True`` returns 1.0 at carved/panden cells and the coverage fraction below threshold.""" + ds, _gwf, geoms = disv_grid() + rows = [ + _piece(CELL_40, box(0, 400, 80, 600), 2.0, 1.0), # 40% -> not carved, fraction 0.4 + _piece(CELL_60, box(200, 400, 320, 600), 2.0, 1.0), # 60% -> carved, fraction 1.0 + _piece(CELL_FULL, geoms[CELL_FULL], 2.0, 1.0), # 100% -> carved, fraction 1.0 + ] + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, _lake_gdf(rows), min_area_fraction=0.5) + assert set(lake_cellids.tolist()) == {CELL_60, CELL_FULL} + + frac = lakes.recharge_pond_mask(ds_carved, None, fractional=True) + # carved cells are fully stage-carried (whole cell modeled as lake bed), even at 60% coverage + assert frac.sel(icell2d=CELL_60).item() == pytest.approx(1.0) + assert frac.sel(icell2d=CELL_FULL).item() == pytest.approx(1.0) + # below-threshold cell keeps its open-water fraction; untouched cells keep 0.0 + assert frac.sel(icell2d=CELL_40).item() == pytest.approx(0.4) + assert frac.sel(icell2d=PANDEN_A).item() == pytest.approx(0.0) + + # the boolean default is unchanged: the 40% cell is not excluded + mask = lakes.recharge_pond_mask(ds_carved, None) + assert not bool(mask.sel(icell2d=CELL_40).item()) + + # a panden RIV cell is fully stage-carried in both modes + stub = types.SimpleNamespace(stress_period_data=types.SimpleNamespace(data={0: {"cellid": [(0, PANDEN_A)]}})) + frac_panden = lakes.recharge_pond_mask(ds_carved, stub, fractional=True) + assert frac_panden.sel(icell2d=PANDEN_A).item() == pytest.approx(1.0) + + def test_carve_lake_cells_empty_when_no_cell_clears_threshold(disv_grid): """An empty / all-below-threshold lake gdf carves nothing instead of crashing. @@ -180,3 +264,113 @@ def test_riv_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): sliver = _lake_gdf([_piece(CELL_40, box(0, 400, 20, 420), 2.0, 1.0)]) ds_carved, _ = lakes.carve_lake_cells(ds, sliver, min_area_fraction=0.5) assert lakes.riv_from_lakes_pwn(ds_carved, gwf, sliver, min_area_fraction=0.5) is None + + +def _add_time_and_recharge(ds, recharge=0.0007): + """Extend the synthetic ds with the time axis and recharge that LAK needs.""" + ds = ds.assign_coords(time=pd.to_datetime(["2023-01-01"])) + ds["time"].attrs["time_units"] = "days" + ds["recharge"] = xr.DataArray( + np.full((1, ds.sizes["icell2d"]), recharge), dims=("time", "icell2d"), attrs={"units": "m/d"} + ) + return ds + + +def test_lak_connections_match_carved_cells_with_equivalent_conductance(disv_grid): + """LAK connects exactly the carved cells, with the same lakebed conductance as the RIV. + + The effective bed resistance is chosen such that ``bedleak * cell_area`` equals the + summed piece conductance (piece_area / clake) of the corresponding RIV reach. The total + MF6 exchange is still somewhat weaker than the RIV's, because VERTICAL connections add + the half-cell vertical resistance in series with the lakebed. + """ + ds, gwf, geoms = disv_grid() + ds = _add_time_and_recharge(ds) + gdf = _single_and_two_piece_lakes(geoms) + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + + lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + + conns = lak.connectiondata.array + assert {int(cid[-1]) for cid in conns["cellid"]} == set(lake_cellids.tolist()) + assert all(cid[0] == 0 for cid in conns["cellid"]) # first active layer + + cell_area = 40000.0 + bedleak = {int(cid[-1]): bl for cid, bl in zip(conns["cellid"], conns["bedleak"], strict=True)} + assert bedleak[CELL_FULL] * cell_area == pytest.approx(cell_area / CLAKE) + assert bedleak[CELL_60] * cell_area == pytest.approx((AREA_A + AREA_B) / CLAKE) + + strt_by_name = {rec[-1]: rec[1] for rec in lak.packagedata.array} + assert strt_by_name["lake_single"] == pytest.approx(2.5) + assert strt_by_name["lake_two_pieces"] == pytest.approx((AREA_A * STRT_A + AREA_B * STRT_B) / (AREA_A + AREA_B)) + + # the default rainfall='from_ds' feeds ds['recharge'] to the lakes as RAINFALL + settings = {(int(rec[0]), rec[1]): rec[2] for rec in lak.perioddata.data[0]} + lakeno_by_name = {rec[-1]: int(rec[0]) for rec in lak.packagedata.array} + assert settings[lakeno_by_name["lake_single"], "RAINFALL"] == pytest.approx(0.0007) + + +def test_lak_two_lakes_sharing_a_cell_stay_separate(disv_grid): + """Two lakes overlapping one cell each keep their own lake, connection and strt. + + Pins the strt collapse's grouping key: grouping by cell instead of by lake would blend + the two lakes' stages in the shared cell and crash nlmod's single-strt-per-lake check. + """ + ds, gwf, _geoms = disv_grid() + ds = _add_time_and_recharge(ds) + clake = 10.0 + strt_a, strt_b = 2.0, 3.0 + rows = [ + _piece(CELL_FULL, box(200, 200, 300, 400), strt_a, 1.0, clake=clake, ident="lake_a"), + _piece(CELL_FULL, box(300, 200, 400, 400), strt_b, 0.5, clake=clake, ident="lake_b"), + _piece(CELL_60, box(200, 400, 400, 600), strt_a, 1.0, clake=clake, ident="lake_a"), + _piece(PANDEN_B, box(400, 200, 600, 400), strt_b, 0.5, clake=clake, ident="lake_b"), + ] + gdf = _lake_gdf(rows) + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + + lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + + strt_by_name = {rec[-1]: rec[1] for rec in lak.packagedata.array} + assert strt_by_name["lake_a"] == pytest.approx(strt_a) + assert strt_by_name["lake_b"] == pytest.approx(strt_b) + assert {int(rec[0]) for rec in lak.packagedata.array} == {0, 1} # two distinct lakes + + conns = lak.connectiondata.array + shared = [conn for conn in conns if int(conn["cellid"][-1]) == CELL_FULL] + assert len(shared) == len(strt_by_name) # one connection per lake in the shared cell + # each lake's lakebed conductance in the shared cell equals its RIV cond (half cell each) + for conn in shared: + assert conn["bedleak"] * 40000.0 == pytest.approx(20000.0 / clake) + assert {int(conn["cellid"][-1]) for conn in conns} == set(lake_cellids.tolist()) + + +def test_lak_single_strt_for_lake_spanning_multiple_cells(disv_grid): + """A lake over several cells with differing per-cell stages gets one exact strt. + + Without the per-lake strt collapse, nlmod's single-strt-per-lake check raises on any + per-cell difference (deterministic proxy for float-level aggregation noise). + """ + ds, gwf, geoms = disv_grid() + ds = _add_time_and_recharge(ds) + rows = [ + _piece(CELL_FULL, geoms[CELL_FULL], 2.0, 0.0, ident="one_lake"), + _piece(CELL_60, geoms[CELL_60], 3.0, 0.5, ident="one_lake"), + ] + gdf = _lake_gdf(rows) + ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) + + lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + + assert len(lak.packagedata.array) == 1 + assert lak.packagedata.array[0][1] == pytest.approx(2.5) # mean over equal-area cells + assert {int(conn["cellid"][-1]) for conn in lak.connectiondata.array} == set(lake_cellids.tolist()) + + +def test_lak_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): + """LAK is ``None`` (no package added) when no cell clears the threshold.""" + ds, gwf, _geoms = disv_grid() + ds = _add_time_and_recharge(ds) + sliver = _lake_gdf([_piece(CELL_40, box(0, 400, 20, 420), 2.0, 1.0)]) + ds_carved, _ = lakes.carve_lake_cells(ds, sliver, min_area_fraction=0.5) + assert lakes.lak_from_lakes_pwn(ds_carved, gwf, sliver, min_area_fraction=0.5) is None From a220ae2d3243db1e2f50e7279c7cff9d8187f547 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Wed, 22 Jul 2026 09:27:32 +0200 Subject: [PATCH 4/4] Prepare LAK input as a frame; build the package via nlmod.gwf.lake_from_gdf lak_from_lakes_pwn wrapped nlmod's lake builder; replace it with lak_gdf_from_lakes_pwn, which only aggregates the lakes_pwn pieces to the per-(cell, lake) frame that nlmod.gwf.lake_from_gdf consumes (effective clake = cell_area / sum(piece_area/clake), one exact strt per lake, outlet columns carried through, None when no cell clears the coverage threshold). The model script now calls nlmod.gwf.lake_from_gdf and copy_meteorological_data_from_ds directly. --- src/nhflotools/lakes.py | 114 ++++++++++++---------------------------- tests/test_lakes.py | 26 ++++++--- 2 files changed, 53 insertions(+), 87 deletions(-) diff --git a/src/nhflotools/lakes.py b/src/nhflotools/lakes.py index 53bfd2f..33584d1 100644 --- a/src/nhflotools/lakes.py +++ b/src/nhflotools/lakes.py @@ -8,9 +8,10 @@ individually. Both the carved cell set and the RIV reach set are derived from the single aggregator :func:`_aggregate_lake_cells`, so the two sets are equal by construction: no carved cell is ever left without a stage, and no stage reach ever lands on an un-carved -cell. For detailed studies :func:`lak_from_lakes_pwn` builds the LAK package from the -same aggregator instead, solving the stage from the lake water balance with the same -lakebed leakance as the RIV. +cell. For detailed studies :func:`lak_gdf_from_lakes_pwn` prepares a LAK input frame +from the same aggregator instead, which the model script feeds directly to +:func:`nlmod.gwf.lake_from_gdf` to solve the stage from the lake water balance with +the same lakebed leakance as the RIV. The logic lives here (rather than inline in the model script) so it can be imported and unit-tested without running the full REGIS/AHN/MF6 build. @@ -223,42 +224,29 @@ def riv_from_lakes_pwn(ds, gwf, gdf_lake_grid, min_area_fraction=0.5): return riv -def lak_from_lakes_pwn( - ds, - gwf, - gdf_lake_grid, - min_area_fraction=0.5, - gwt=None, - rainfall="from_ds", - evaporation="from_ds", - **kwargs, -): - """Model the carved lake cells with the LAK package instead of a stage RIV. - - Alternative to :func:`riv_from_lakes_pwn` for detailed studies: the lake stage - follows from the simulated lake water balance instead of being prescribed, which - resolves stage drawdown, lake storage and (via ``lakeout`` columns) outlets, at the - cost of a harder nonlinear solve. The connection set is derived from the same - :func:`_aggregate_lake_cells` call as :func:`carve_lake_cells`, so it equals the - carved-cell set by construction, and the lakebed leakance matches the RIV variant: - each (cell, lake) row becomes one VERTICAL connection with an effective bed - resistance ``clake_eff = cell_area / sum(piece_area / clake)``, so ``bedleak * - cell_area`` equals the summed piece conductance of the corresponding RIV reach. The - total exchange is nevertheless somewhat weaker than the RIV's, because MF6 places - the connected cell's half-cell vertical resistance (``0.5 * thickness / k33``) in - series with the lakebed for VERTICAL connections, which the RIV formulation applies - directly to the cell node (12-16% lower on the PWN layer model). Like the RIV (and - unlike a GHB), the exchange is capped at the lakebed once the aquifer head drops - below it. +def lak_gdf_from_lakes_pwn(ds, gdf_lake_grid, min_area_fraction=0.5): + """Prepare the per-cell lake frame that :func:`nlmod.gwf.lake_from_gdf` consumes. + + Alternative input to :func:`riv_from_lakes_pwn` for detailed studies: the caller + builds the LAK package directly with :func:`nlmod.gwf.lake_from_gdf` (and typically + :func:`nlmod.gwf.copy_meteorological_data_from_ds` for per-lake rainfall and + evaporation), so the lake stage follows from the simulated water balance instead of + being prescribed. The frame is derived from the same :func:`_aggregate_lake_cells` + call as :func:`carve_lake_cells`, so the LAK connection set equals the carved-cell + set by construction, and the lakebed leakance matches the RIV variant: each + (cell, lake) row yields one VERTICAL connection with an effective bed resistance + ``clake = cell_area / sum(piece_area / clake)``, so ``bedleak * cell_area`` equals + the summed piece conductance of the corresponding RIV reach. The total exchange is + nevertheless somewhat weaker than the RIV's, because MF6 places the connected + cell's half-cell vertical resistance (``0.5 * thickness / k33``) in series with the + lakebed for VERTICAL connections, which the RIV formulation applies directly to the + cell node (12-16% lower on the PWN layer model). Like the RIV (and unlike a GHB), + the exchange is capped at the lakebed once the aquifer head drops below it. Parameters ---------- ds : xarray.Dataset - Model dataset (post-carve) with ``top``, ``botm``, ``area``, ``idomain`` and a - time discretisation in days. ``recharge`` is required when - ``rainfall='from_ds'`` and ``chloride`` when ``gwt`` is given. - gwf : flopy.mf6.ModflowGwf - Groundwater flow model the LAK package is added to. + Model dataset (post-carve). Only ``ds['area']`` is used. gdf_lake_grid : geopandas.GeoDataFrame Lake polygons intersected with the model grid (see :func:`_aggregate_lake_cells`). Per-lake outlet columns (``lakeout``, @@ -267,32 +255,23 @@ def lak_from_lakes_pwn( min_area_fraction : float, optional Minimum combined lake coverage for a cell to be connected, by default 0.5 (strict ``>``), identical to the carve threshold. - gwt : flopy.mf6.ModflowGwt, optional - When given, a matching LKT transport package is created as well and - ``(lak, lkt)`` is returned. The default is None. - rainfall, evaporation : str, float, pandas.DataFrame or None, optional - Passed to :func:`nlmod.gwf.lake.lake_from_gdf`. The default ``'from_ds'`` - derives per-lake area-weighted series from the meteorology in ``ds`` with - :func:`nlmod.gwf.lake.copy_meteorological_data_from_ds`. Exclude the lake cells - from the RCH package (see :func:`recharge_pond_mask`) so the meteoric term is - not counted both on the lake and on the aquifer. - **kwargs - Passed to :func:`nlmod.gwf.lake.lake_from_gdf` and on to - :class:`flopy.mf6.ModflowGwflak` (e.g. ``save_flows``, - ``package_convergence_filerecord``). Returns ------- - flopy.mf6.ModflowGwflak, tuple or None - The LAK package, ``(lak, lkt)`` when ``gwt`` is given, or ``None`` when no cell - clears the coverage threshold (so the caller adds no package). + pandas.DataFrame or None + One row per (cell, lake), indexed by ``icell2d`` as ``lake_from_gdf`` expects, + with ``strt`` (one exact value per lake — the per-cell area-weighted values can + differ at floating-point level and LAK requires a single strt), the effective + ``clake``, ``identificatie`` and any outlet columns. ``None`` when no cell + clears the coverage threshold, so the caller adds no package. Notes ----- - When two lakes share a carved cell, MF6 applies each lake's RAINFALL over that - lake's full connection (cell) area while RCH excluded the cell only once, so the - meteoric term on such a cell is over-applied. No current ``lakes_pwn`` cell is - shared between lakes. + Exclude the lake cells from the RCH package (see :func:`recharge_pond_mask`) so the + meteoric term is not counted both on the lake and on the aquifer. When two lakes + share a carved cell, MF6 applies each lake's RAINFALL over that lake's full + connection (cell) area while RCH excluded the cell only once, so the meteoric term + on such a cell is over-applied; no current ``lakes_pwn`` cell is shared. """ agg, _ = _aggregate_lake_cells(ds, gdf_lake_grid, min_area_fraction=min_area_fraction) if agg.empty: @@ -300,8 +279,6 @@ def lak_from_lakes_pwn( gdf = agg.drop(columns=["botm"]) gdf["clake"] = ds["area"].sel(icell2d=gdf.index.to_numpy()).to_numpy() / gdf.pop("cond") - # LAK requires one exact strt per lake; the per-cell area-weighted values can differ - # at floating-point level. gdf["strt"] = gdf.groupby("identificatie")["strt"].transform("mean") outlet_columns = [ @@ -312,28 +289,7 @@ def lak_from_lakes_pwn( if outlet_columns: per_lake = gdf_lake_grid.drop_duplicates("identificatie").set_index("identificatie")[outlet_columns] gdf = gdf.join(per_lake, on="identificatie") - - rainfall_from_ds = isinstance(rainfall, str) and rainfall == "from_ds" - evaporation_from_ds = isinstance(evaporation, str) and evaporation == "from_ds" - if rainfall_from_ds or evaporation_from_ds: - ds_rainfall, ds_evaporation = nlmod.gwf.lake.copy_meteorological_data_from_ds( - gdf, ds, boundname_column="identificatie" - ) - if rainfall_from_ds: - rainfall = ds_rainfall - if evaporation_from_ds: - evaporation = ds_evaporation - - return nlmod.gwf.lake.lake_from_gdf( - gwf, - gdf, - ds, - rainfall=rainfall, - evaporation=evaporation, - boundname_column="identificatie", - gwt=gwt, - **kwargs, - ) + return gdf def recharge_pond_mask(ds, panden_riv=None, *, fractional=False): diff --git a/tests/test_lakes.py b/tests/test_lakes.py index 3e229ef..996aa8f 100644 --- a/tests/test_lakes.py +++ b/tests/test_lakes.py @@ -11,6 +11,7 @@ import types import geopandas as gpd +import nlmod import numpy as np import pandas as pd import pytest @@ -266,6 +267,15 @@ def test_riv_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): assert lakes.riv_from_lakes_pwn(ds_carved, gwf, sliver, min_area_fraction=0.5) is None +def _build_lak(ds, gwf, gdf_lake_grid): + """Mirror the model script's LAK branch: prepare the frame, then build via nlmod.""" + gdf = lakes.lak_gdf_from_lakes_pwn(ds, gdf_lake_grid, min_area_fraction=0.5) + rainfall, evaporation = nlmod.gwf.copy_meteorological_data_from_ds(gdf, ds, boundname_column="identificatie") + return nlmod.gwf.lake_from_gdf( + gwf, gdf, ds, rainfall=rainfall, evaporation=evaporation, boundname_column="identificatie" + ) + + def _add_time_and_recharge(ds, recharge=0.0007): """Extend the synthetic ds with the time axis and recharge that LAK needs.""" ds = ds.assign_coords(time=pd.to_datetime(["2023-01-01"])) @@ -289,7 +299,7 @@ def test_lak_connections_match_carved_cells_with_equivalent_conductance(disv_gri gdf = _single_and_two_piece_lakes(geoms) ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) - lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + lak = _build_lak(ds_carved, gwf, gdf) conns = lak.connectiondata.array assert {int(cid[-1]) for cid in conns["cellid"]} == set(lake_cellids.tolist()) @@ -304,7 +314,7 @@ def test_lak_connections_match_carved_cells_with_equivalent_conductance(disv_gri assert strt_by_name["lake_single"] == pytest.approx(2.5) assert strt_by_name["lake_two_pieces"] == pytest.approx((AREA_A * STRT_A + AREA_B * STRT_B) / (AREA_A + AREA_B)) - # the default rainfall='from_ds' feeds ds['recharge'] to the lakes as RAINFALL + # copy_meteorological_data_from_ds feeds ds['recharge'] to the lakes as RAINFALL settings = {(int(rec[0]), rec[1]): rec[2] for rec in lak.perioddata.data[0]} lakeno_by_name = {rec[-1]: int(rec[0]) for rec in lak.packagedata.array} assert settings[lakeno_by_name["lake_single"], "RAINFALL"] == pytest.approx(0.0007) @@ -329,7 +339,7 @@ def test_lak_two_lakes_sharing_a_cell_stay_separate(disv_grid): gdf = _lake_gdf(rows) ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) - lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + lak = _build_lak(ds_carved, gwf, gdf) strt_by_name = {rec[-1]: rec[1] for rec in lak.packagedata.array} assert strt_by_name["lake_a"] == pytest.approx(strt_a) @@ -360,17 +370,17 @@ def test_lak_single_strt_for_lake_spanning_multiple_cells(disv_grid): gdf = _lake_gdf(rows) ds_carved, lake_cellids = lakes.carve_lake_cells(ds, gdf, min_area_fraction=0.5) - lak = lakes.lak_from_lakes_pwn(ds_carved, gwf, gdf, min_area_fraction=0.5) + lak = _build_lak(ds_carved, gwf, gdf) assert len(lak.packagedata.array) == 1 assert lak.packagedata.array[0][1] == pytest.approx(2.5) # mean over equal-area cells assert {int(conn["cellid"][-1]) for conn in lak.connectiondata.array} == set(lake_cellids.tolist()) -def test_lak_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): - """LAK is ``None`` (no package added) when no cell clears the threshold.""" - ds, gwf, _geoms = disv_grid() +def test_lak_gdf_from_lakes_pwn_returns_none_when_no_lake_cells(disv_grid): + """The LAK input frame is ``None`` when no cell clears the threshold (no package built).""" + ds, _gwf, _geoms = disv_grid() ds = _add_time_and_recharge(ds) sliver = _lake_gdf([_piece(CELL_40, box(0, 400, 20, 420), 2.0, 1.0)]) ds_carved, _ = lakes.carve_lake_cells(ds, sliver, min_area_fraction=0.5) - assert lakes.lak_from_lakes_pwn(ds_carved, gwf, sliver, min_area_fraction=0.5) is None + assert lakes.lak_gdf_from_lakes_pwn(ds_carved, sliver, min_area_fraction=0.5) is None