diff --git a/src/nhflotools/lakes.py b/src/nhflotools/lakes.py new file mode 100644 index 0000000..33584d1 --- /dev/null +++ b/src/nhflotools/lakes.py @@ -0,0 +1,335 @@ +"""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 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_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. +""" + +import logging + +import flopy +import nlmod +import numpy as np +import pandas as pd +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 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 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 + ---------- + 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 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 + ------- + 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) + + 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 empty_agg, coverage + + lake["cond"] = lake["area"] / lake["clake"] + 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 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 + ---------- + 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 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 ``lake_cell`` / ``lake_coverage`` + markers. + 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, 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}] = 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 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 its cell set 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 combined lake coverage for a cell to be carved/bounded, by default 0.5. + + Returns + ------- + 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) + agg = agg.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") + + 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 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). 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``, + ``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. + + Returns + ------- + 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 + ----- + 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: + return None + + gdf = agg.drop(columns=["botm"]) + gdf["clake"] = ds["area"].sel(icell2d=gdf.index.to_numpy()).to_numpy() / gdf.pop("cond") + 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") + return gdf + + +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. 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 ``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 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 + 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 + if not fractional: + return mask + return xr.where(mask, 1.0, ds["lake_coverage"], keep_attrs=True) 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..996aa8f --- /dev/null +++ b/tests/test_lakes.py @@ -0,0 +1,386 @@ +"""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-lake-per-cell aggregation (pieces merge within a lake, never +across lakes), and the boolean and fractional recharge masks. +""" + +import logging +import types + +import geopandas as gpd +import nlmod +import numpy as np +import pandas as pd +import pytest +import xarray as xr +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 + +# 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.""" + 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 _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() + # 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) + 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): + 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_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() + 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) + + +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. + + 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 + + +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"])) + 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 = _build_lak(ds_carved, gwf, gdf) + + 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)) + + # 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) + + +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 = _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) + 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 = _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_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_gdf_from_lakes_pwn(ds_carved, sliver, min_area_fraction=0.5) is None 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}