From d6e471d57a1b22b99b047504a30d9862051d9cad Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Wed, 8 Jul 2026 09:47:24 +0200 Subject: [PATCH 1/7] Add drain_from_df helper Build DRN packages from vector and tabular drain data with deterministic MVR provider mapping. --- nlmod/gwf/__init__.py | 3 +- nlmod/gwf/drain.py | 421 ++++++++++++++++++++++++++++++++++++++++ tests/test_027_drain.py | 299 ++++++++++++++++++++++++++++ 3 files changed, 722 insertions(+), 1 deletion(-) create mode 100644 nlmod/gwf/drain.py create mode 100644 tests/test_027_drain.py diff --git a/nlmod/gwf/__init__.py b/nlmod/gwf/__init__.py index 9d5233f6..3cbffef3 100644 --- a/nlmod/gwf/__init__.py +++ b/nlmod/gwf/__init__.py @@ -1,5 +1,6 @@ # ruff: noqa: F401 F403 -from . import hfb, output, surface_water, wells +from . import drain, hfb, output, surface_water, wells +from .drain import * from .gwf import * from .hfb import * from .lake import * diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py new file mode 100644 index 00000000..15eac11e --- /dev/null +++ b/nlmod/gwf/drain.py @@ -0,0 +1,421 @@ +import logging + +import flopy +import geopandas as gpd +import numpy as np +import pandas as pd + +from ..dims.grid import gdf_to_grid +from ..util import tqdm +from .surface_water import build_spd + +logger = logging.getLogger(__name__) + +LINE_GEOM_TYPES = {"LineString", "MultiLineString"} +POLYGON_GEOM_TYPES = {"Polygon", "MultiPolygon"} +POINT_GEOM_TYPES = {"Point"} + + +def drain_from_df( + df, + gwf, + ds, + elev="elevation", + cond="cond", + conductance_per_length="conductance_per_meter", + conductance_per_area="conductance_per_squared_meter", + x="x", + y="y", + boundnames=None, + mover_destinations=None, + layer_method="lay_of_rbot", + pname="drn", + silent=False, + return_provider_mapping=False, + **kwargs, +): + """Add a Drain (DRN) package based on input from a (Geo)DataFrame. + + Parameters + ---------- + df : pd.DataFrame or gpd.GeoDataFrame + A (Geo)DataFrame containing the drain properties. Line and polygon + geometries are intersected with the model grid and converted to + conductance using ``conductance_per_length`` and ``conductance_per_area``. + Point geometries and non-geometric data require ``cond`` to contain the + integrated MF6 drain conductance. + gwf : flopy ModflowGwf + Groundwaterflow object to add the DRN package to. + ds : xarray.Dataset + Dataset with model data. Used for grid intersection and layer placement. + elev : str, optional + Column in ``df`` that contains the drain elevation. The default is + "elevation". + cond : str, optional + Column in ``df`` that contains the integrated drain conductance. Required + for point geometries and direct cellid input. The default is "cond". + conductance_per_length : str, optional + Column in ``df`` that contains conductance per metre for line geometries. + The default is "conductance_per_meter". + conductance_per_area : str, optional + Column in ``df`` that contains conductance per square metre for polygon + geometries. The default is "conductance_per_squared_meter". + x : str, optional + Column in ``df`` that contains the x-coordinate for point drains when + ``df`` is not a GeoDataFrame and no ``cellid`` column is present. The + default is "x". + y : str, optional + Column in ``df`` that contains the y-coordinate for point drains when + ``df`` is not a GeoDataFrame and no ``cellid`` column is present. The + default is "y". + boundnames : str, optional + Column in ``df`` that contains boundary names. These are written to the + DRN package and included in the provider mapping. The default is None. + mover_destinations : str, optional + Column in ``df`` that identifies the intended MVR receiver for each drain. + This function does not create MVR routes, but stores this column in the + provider mapping. The default is None. + layer_method : str, optional + Method used by ``nlmod.gwf.surface_water.build_spd`` for layer placement. + The default is "lay_of_rbot". + pname : str, optional + Package name. The default is "drn". + silent : bool, optional + Do not show progress bars when silent is True. The default is False. + return_provider_mapping : bool, optional + Return ``(drn, provider_mapping)`` when True. The mapping is always added + to the returned package as ``drn.mvr_provider_mapping``. Provider IDs are + zero-based FloPy numeric indices for use in MVR period data. The default + is False. + **kwargs : dict + Kwargs are passed to ``flopy.mf6.ModflowGwfdrn``. Use ``mover=True`` to + make the drain package available to an MVR package. + + Returns + ------- + drn : flopy.mf6.ModflowGwfdrn + DRN package. When ``return_provider_mapping`` is True, returns a tuple of + ``(drn, provider_mapping)``. + """ + logger.info("creating mf6 DRN from dataframe") + + celldata = _drain_celldata_from_df( + df, + gwf=gwf, + ds=ds, + elev=elev, + cond=cond, + conductance_per_length=conductance_per_length, + conductance_per_area=conductance_per_area, + x=x, + y=y, + boundnames=boundnames, + mover_destinations=mover_destinations, + silent=silent, + ) + spd, provider_mapping = _build_spd_with_provider_mapping( + celldata, + ds=ds, + layer_method=layer_method, + silent=silent, + ) + + if len(spd) == 0: + logger.warning("no drn pkg added") + if return_provider_mapping: + return None, provider_mapping + return None + + save_flows = kwargs.pop("save_flows", True) + drn = flopy.mf6.ModflowGwfdrn( + gwf, + maxbound=len(spd), + stress_period_data={0: spd}, + save_flows=save_flows, + boundnames=boundnames is not None, + pname=pname, + **kwargs, + ) + + provider_mapping["package"] = drn.package_name + drn.mvr_provider_mapping = provider_mapping + if return_provider_mapping: + return drn, provider_mapping + return drn + + +def mvr_perioddata_from_provider_mapping( + provider_mapping, + receiver_package, + receiver_id_map, + receiver_column="mover_destination", + provider_package=None, + mvrtype="FACTOR", + value=1.0, +): + """Build MVR perioddata for drain outflow using a provider mapping. + + Parameters + ---------- + provider_mapping : pd.DataFrame + Provider mapping returned by ``drain_from_df``. Must contain + ``mvr_provider_id`` and the receiver column. + receiver_package : str + Name of the receiving package in the MVR package. + receiver_id_map : dict or pd.Series + Mapping from values in ``receiver_column`` to zero-based receiver IDs. + receiver_column : str, optional + Column in ``provider_mapping`` that identifies the receiver. The default is + "mover_destination". + provider_package : str, optional + Name of the provider package. When None, the ``package`` column in + ``provider_mapping`` is used. The default is None. + mvrtype : str, optional + MVR rule type. The default is "FACTOR". + value : float, optional + MVR rule value. The default is 1.0. + + Returns + ------- + perioddata : list + Perioddata records that can be passed to ``flopy.mf6.ModflowGwfmvr``. + """ + _validate_columns( + provider_mapping, ["mvr_provider_id", receiver_column], "provider_mapping" + ) + if provider_package is None: + _validate_columns(provider_mapping, ["package"], "provider_mapping") + + perioddata = [] + for _, row in provider_mapping.dropna(subset=[receiver_column]).iterrows(): + receiver_name = row[receiver_column] + if receiver_name not in receiver_id_map: + raise KeyError(f"Receiver {receiver_name!r} not found in receiver_id_map") + perioddata.append( + ( + row["package"] if provider_package is None else provider_package, + row["mvr_provider_id"], + receiver_package, + receiver_id_map[receiver_name], + mvrtype, + value, + ) + ) + return perioddata + + +def _drain_celldata_from_df( + df, + gwf, + ds, + elev, + cond, + conductance_per_length, + conductance_per_area, + x, + y, + boundnames, + mover_destinations, + silent, +): + if not isinstance(df, (pd.DataFrame, gpd.GeoDataFrame)): + raise TypeError("df must be a pandas DataFrame or geopandas GeoDataFrame") + + _validate_columns(df, [elev], "df") + if boundnames is not None: + _validate_columns(df, [boundnames], "df") + if mover_destinations is not None: + _validate_columns(df, [mover_destinations], "df") + + if isinstance(df, gpd.GeoDataFrame): + gdf = df.copy() + elif "cellid" not in df.columns and {x, y}.issubset(df.columns): + gdf = gpd.GeoDataFrame(df.copy(), geometry=gpd.points_from_xy(df[x], df[y])) + else: + return _cellid_celldata_from_df( + df.copy(), + elev=elev, + cond=cond, + boundnames=boundnames, + mover_destinations=mover_destinations, + ) + + gdf["_source_index"] = gdf.index.to_numpy() + parts = [] + for geom_types, conductance_column, measure in ( + (LINE_GEOM_TYPES, conductance_per_length, "length"), + (POLYGON_GEOM_TYPES, conductance_per_area, "area"), + (POINT_GEOM_TYPES, cond, None), + ): + subset = gdf[gdf.geom_type.isin(geom_types)] + if subset.empty: + continue + if measure is None: + _validate_columns(subset, [cond], "df") + else: + _validate_columns(subset, [conductance_column], "df") + parts.append( + _geodataframe_celldata( + subset, + gwf=gwf, + ds=ds, + elev=elev, + cond=cond, + conductance_column=conductance_column, + measure=measure, + boundnames=boundnames, + mover_destinations=mover_destinations, + silent=silent, + ) + ) + + unsupported_geom_types = set(gdf.geom_type.unique()).difference( + LINE_GEOM_TYPES | POLYGON_GEOM_TYPES | POINT_GEOM_TYPES + ) + if unsupported_geom_types: + raise TypeError(f"Unsupported drain geometry types: {unsupported_geom_types}") + + if len(parts) == 0: + return _empty_celldata() + return pd.concat(parts, axis=0) + + +def _geodataframe_celldata( + gdf, + gwf, + ds, + elev, + cond, + conductance_column, + measure, + boundnames, + mover_destinations, + silent, +): + if "cellid" not in gdf.columns: + gdf = gdf_to_grid(gdf, ds if ds is not None else gwf, silent=silent) + + celldata = _base_celldata(gdf, elev, boundnames, mover_destinations) + if measure == "length": + celldata["cond"] = gdf.geometry.length.to_numpy() * gdf[conductance_column] + celldata["len_estimate"] = gdf.geometry.length.to_numpy() + celldata["area"] = np.nan + elif measure == "area": + celldata["cond"] = gdf.geometry.area.to_numpy() * gdf[conductance_column] + celldata["area"] = gdf.geometry.area.to_numpy() + celldata["len_estimate"] = np.nan + elif measure is None: + celldata["cond"] = gdf[cond].to_numpy() + celldata["area"] = np.nan + celldata["len_estimate"] = np.nan + else: + raise ValueError(f"Unknown measure: {measure}") + return celldata.set_index("cellid") + + +def _cellid_celldata_from_df(df, elev, cond, boundnames, mover_destinations): + _validate_columns(df, ["cellid", cond], "df") + celldata = _base_celldata(df, elev, boundnames, mover_destinations) + celldata["cond"] = df[cond].to_numpy() + celldata["area"] = np.nan + celldata["len_estimate"] = np.nan + return celldata.set_index("cellid") + + +def _base_celldata(df, elev, boundnames, mover_destinations): + celldata = pd.DataFrame(index=df.index) + celldata["cellid"] = df["cellid"].to_numpy() + celldata["stage"] = df[elev].to_numpy() + celldata["rbot"] = df[elev].to_numpy() + celldata["source_index"] = ( + df["_source_index"].to_numpy() + if "_source_index" in df.columns + else df.index.to_numpy() + ) + if boundnames is not None: + celldata["boundname"] = df[boundnames].to_numpy() + if mover_destinations is not None: + celldata["mover_destination"] = df[mover_destinations].to_numpy() + return celldata + + +def _build_spd_with_provider_mapping(celldata, ds, layer_method, silent): + spd = [] + provider_mapping = [] + for cellid, row in tqdm( + celldata.iterrows(), + total=celldata.index.size, + desc="Building stress period data DRN", + disable=silent, + ): + if _is_3d_cellid(cellid, ds): + row_spd = [_record_from_3d_cellid(cellid, row)] + else: + index = np.empty(1, dtype=object) + index[0] = cellid + row_df = pd.DataFrame( + [row], + index=pd.Index(index, name=celldata.index.name), + ) + row_spd = build_spd( + row_df, + "DRN", + ds, + layer_method=layer_method, + silent=True, + ) + for record in row_spd: + provider_mapping.append( + { + "mvr_provider_id": len(spd), + "cellid": record[0], + "elev": record[1], + "cond": record[2], + "source_index": row["source_index"], + "boundname": row.get("boundname"), + "mover_destination": row.get("mover_destination"), + } + ) + spd.append(record) + return spd, pd.DataFrame(provider_mapping) + + +def _is_3d_cellid(cellid, ds): + if not isinstance(cellid, tuple): + return False + if ds.gridtype == "vertex": + return len(cellid) == 2 + if ds.gridtype == "structured": + return len(cellid) == 3 + raise ValueError(f"Unsupported gridtype: {ds.gridtype}") + + +def _record_from_3d_cellid(cellid, row): + if np.isnan(row["cond"]): + raise ValueError(f"Conductance is NaN in cell {cellid}") + if row["cond"] < 0: + raise ValueError(f"Conductance is negative in cell {cellid}") + + auxlist = [] + if "boundname" in row: + auxlist.append(row["boundname"]) + return [cellid, row["stage"], row["cond"]] + auxlist + + +def _empty_celldata(): + return pd.DataFrame( + columns=[ + "stage", + "rbot", + "cond", + "area", + "len_estimate", + "source_index", + ] + ) + + +def _validate_columns(df, columns, name): + missing = set(columns).difference(df.columns) + if missing: + raise ValueError(f"Missing columns in {name}: {missing}") diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py new file mode 100644 index 00000000..c5125273 --- /dev/null +++ b/tests/test_027_drain.py @@ -0,0 +1,299 @@ +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +import test_010_wells +from shapely.geometry import LineString, MultiPoint, Polygon + +import nlmod + + +def test_drain_from_df_vector_keeps_drain_thresholds_and_mvr_mapping(): + """Test vector drain conductance and MVR provider metadata.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = gpd.GeoDataFrame( + { + "name": ["line-drain", "area-drain"], + "elevation": [-1.0, -3.0], + "conductance_per_meter": [2.0, np.nan], + "conductance_per_squared_meter": [np.nan, 3.0], + "mover_lake_name": ["lake-1", "lake-2"], + }, + geometry=[ + LineString([(-499.0, 499.0), (-497.0, 499.0)]), + Polygon( + [ + (-499.0, 498.0), + (-497.0, 498.0), + (-497.0, 497.0), + (-499.0, 497.0), + ] + ), + ], + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + boundnames="name", + mover_destinations="mover_lake_name", + mover=True, + pname="drn_test", + silent=True, + return_provider_mapping=True, + ) + + assert drn.package_name == "drn_test" + assert drn.mover.array is True + assert provider_mapping["mvr_provider_id"].tolist() == [0, 1] + assert provider_mapping["mover_destination"].tolist() == ["lake-1", "lake-2"] + assert provider_mapping["boundname"].tolist() == ["line-drain", "area-drain"] + assert provider_mapping["elev"].tolist() == [-1.0, -3.0] + assert provider_mapping["cond"].tolist() == pytest.approx([4.0, 6.0]) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + head = -2.0 + generated_flux = ( + provider_mapping["cond"] * np.maximum(head - provider_mapping["elev"], 0.0) + ).sum() + min_elevation_collapsed_flux = provider_mapping["cond"].sum() * max( + head - -3.0, 0.0 + ) + assert generated_flux == pytest.approx(6.0) + assert generated_flux < min_elevation_collapsed_flux + + +def test_drain_from_df_keeps_same_geometry_thresholds_separate(): + """Test same-geometry drains in one cell keep separate elevations.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = gpd.GeoDataFrame( + { + "elevation": [-1.0, -3.0], + "conductance_per_meter": [2.0, 2.0], + }, + geometry=[ + LineString([(-499.0, 499.0), (-497.0, 499.0)]), + LineString([(-499.0, 498.0), (-497.0, 498.0)]), + ], + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_same_geom", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping["elev"].tolist() == [-1.0, -3.0] + assert provider_mapping["cond"].tolist() == pytest.approx([4.0, 4.0]) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + head = -2.0 + generated_flux = ( + provider_mapping["cond"] * np.maximum(head - provider_mapping["elev"], 0.0) + ).sum() + min_elevation_collapsed_flux = provider_mapping["cond"].sum() * max( + head - -3.0, 0.0 + ) + assert generated_flux == pytest.approx(4.0) + assert generated_flux < min_elevation_collapsed_flux + + +def test_drain_from_df_uses_clipped_line_and_polygon_measures(): + """Test vector conductance uses clipped geometry measure per cell.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = gpd.GeoDataFrame( + { + "elevation": [-1.0, -2.0], + "conductance_per_meter": [2.0, np.nan], + "conductance_per_squared_meter": [np.nan, 2.0], + }, + geometry=[ + LineString([(-499.0, 499.0), (-481.0, 499.0)]), + Polygon( + [ + (-499.0, 498.0), + (-481.0, 498.0), + (-481.0, 495.0), + (-499.0, 495.0), + ] + ), + ], + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_clipped", + silent=True, + return_provider_mapping=True, + ) + + assert sorted(provider_mapping["cond"].tolist()) == pytest.approx( + [18.0, 18.0, 54.0, 54.0] + ) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_mvr_perioddata_from_provider_mapping(): + """Test conversion from DRN provider mapping to MVR perioddata.""" + provider_mapping = pd.DataFrame( + { + "package": ["drn_test", "drn_test"], + "mvr_provider_id": [0, 1], + "mover_destination": ["lake-1", "lake-2"], + } + ) + + perioddata = nlmod.gwf.drain.mvr_perioddata_from_provider_mapping( + provider_mapping, + receiver_package="lak", + receiver_id_map={"lake-1": 0, "lake-2": 1}, + ) + + assert perioddata == [ + ("drn_test", 0, "lak", 0, "FACTOR", 1.0), + ("drn_test", 1, "lak", 1, "FACTOR", 1.0), + ] + + provider_mapping = pd.DataFrame( + { + "package": ["drn_test", "drn_test", "drn_test"], + "mvr_provider_id": [0, 1, 2], + "mover_destination": ["lake-1", np.nan, "lake-2"], + }, + index=[10, 20, 30], + ) + perioddata = nlmod.gwf.drain.mvr_perioddata_from_provider_mapping( + provider_mapping, + receiver_package="lak", + receiver_id_map={"lake-1": 0, "lake-2": 1}, + ) + assert perioddata == [ + ("drn_test", 0, "lak", 0, "FACTOR", 1.0), + ("drn_test", 2, "lak", 1, "FACTOR", 1.0), + ] + + +def test_drain_from_df_preserves_direct_3d_cellids(): + """Test that explicit 3D cell IDs are not prefixed with another layer.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(1, 0, 0)], + "elevation": [-12.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_3d", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (1, 0, 0) + assert provider_mapping.loc[0, "elev"] == -12.0 + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_places_2d_cellids_in_layer_from_elevation(): + """Test that 2D cell IDs use drain elevation for layer placement.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0)], + "elevation": [-12.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_2d", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (1, 0, 0) + assert provider_mapping.loc[0, "elev"] == -12.0 + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_preserves_point_conductance(): + """Test that point drains use supplied integrated conductance unchanged.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "x": [-495.0], + "y": [495.0], + "elevation": [-1.0], + "cond": [7.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_point", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cond"] == 7.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_rejects_multipoint_conductance(): + """Test MultiPoint drains are rejected to avoid duplicating conductance.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = gpd.GeoDataFrame( + { + "elevation": [-1.0], + "cond": [7.0], + }, + geometry=[MultiPoint([(-495.0, 495.0), (-485.0, 495.0)])], + ) + + with pytest.raises(TypeError, match="Unsupported drain geometry types"): + nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_multipoint", + silent=True, + ) + + +def _assert_mapping_matches_stress_period_data(drn, provider_mapping): + spd = drn.stress_period_data.array[0] + assert len(spd) == len(provider_mapping) + assert sorted(provider_mapping["mvr_provider_id"].astype(int)) == list( + range(len(spd)) + ) + for _, row in provider_mapping.iterrows(): + record = spd[int(row["mvr_provider_id"])] + assert record["cellid"] == row["cellid"] + assert record["elev"] == pytest.approx(row["elev"]) + assert record["cond"] == pytest.approx(row["cond"]) + if "boundname" in record.dtype.names and pd.notna(row["boundname"]): + assert record["boundname"] == row["boundname"] From ac6c43bf4d6fb21ed5332348b6e0f6e9158da4dd Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Wed, 8 Jul 2026 16:22:17 +0200 Subject: [PATCH 2/7] Document drain overlap and IDOMAIN handling Clarify overlap with surface water seasonal DRN conversion and reject explicit 3D DRNs in inactive or pass-through cells. --- nlmod/gwf/drain.py | 31 +++++++- nlmod/gwf/surface_water.py | 18 ++++- tests/test_027_drain.py | 151 +++++++++++++++++++++++++++++++++++++ 3 files changed, 194 insertions(+), 6 deletions(-) diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index 15eac11e..fdf4f1f1 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -6,6 +6,7 @@ import pandas as pd from ..dims.grid import gdf_to_grid +from ..dims.layers import get_idomain from ..util import tqdm from .surface_water import build_spd @@ -96,6 +97,24 @@ def drain_from_df( drn : flopy.mf6.ModflowGwfdrn DRN package. When ``return_provider_mapping`` is True, returns a tuple of ``(drn, provider_mapping)``. + + Notes + ----- + This function overlaps with ``nlmod.gwf.surface_water.gdf_to_seasonal_pkg`` + for polygon-to-DRN conversion. Use ``gdf_to_seasonal_pkg`` for surface-water + polygons with winter and summer stages and seasonal conductance timeseries. + Use ``drain_from_df`` for fixed drain features such as pipes, basins, point + drains, or direct cellid input, and when deterministic MVR provider IDs are + needed. + + For vector geometries and 2D cellids, layer placement is delegated to + ``nlmod.gwf.surface_water.build_spd``. That helper uses + ``nlmod.dims.layers.get_idomain`` to skip columns without active cells and to + place the drain in a suitable active layer (``idomain > 0``), not in inactive + (``idomain == 0``) or vertical pass-through (``idomain < 0``) cells. FloPy + receives explicit 3D DRN cellids and does not relocate boundaries to the first + active layer. Therefore explicit 3D cellids passed to this function are + checked and rejected when they target ``idomain <= 0``. """ logger.info("creating mf6 DRN from dataframe") @@ -349,7 +368,7 @@ def _build_spd_with_provider_mapping(celldata, ds, layer_method, silent): disable=silent, ): if _is_3d_cellid(cellid, ds): - row_spd = [_record_from_3d_cellid(cellid, row)] + row_spd = [_record_from_3d_cellid(cellid, row, ds)] else: index = np.empty(1, dtype=object) index[0] = cellid @@ -390,7 +409,15 @@ def _is_3d_cellid(cellid, ds): raise ValueError(f"Unsupported gridtype: {ds.gridtype}") -def _record_from_3d_cellid(cellid, row): +def _record_from_3d_cellid(cellid, row, ds): + idomain_value = get_idomain(ds).data[cellid] + if idomain_value <= 0: + raise ValueError( + f"DRN cellid {cellid} is inactive or pass-through " + f"(idomain={idomain_value}). Use a 2D cellid or vector geometry to let " + "nlmod choose an active layer." + ) + if np.isnan(row["cond"]): raise ValueError(f"Conductance is NaN in cell {cellid}") if row["cond"] < 0: diff --git a/nlmod/gwf/surface_water.py b/nlmod/gwf/surface_water.py index a15c74ca..ec71bc51 100644 --- a/nlmod/gwf/surface_water.py +++ b/nlmod/gwf/surface_water.py @@ -16,7 +16,7 @@ ) from ..dims.layers import get_idomain from ..read import bgt, waterboard -from ..util import extent_to_polygon, gdf_intersection_join, zonal_statistics, tqdm +from ..util import extent_to_polygon, gdf_intersection_join, tqdm, zonal_statistics logger = logging.getLogger(__name__) @@ -912,9 +912,10 @@ def gdf_to_seasonal_pkg( season_filename="season.ts", **kwargs, ): - """Add a surface water package to a groundwater-model, based on input from a - GeoDataFrame. This method adds two boundary conditions for each record in the - GeoDataFrame: one for the winter_stage and one for the summer_stage. + """Add a surface water package to a groundwater-model from a GeoDataFrame. + + This method adds two boundary conditions for each record in the GeoDataFrame: + one for the winter_stage and one for the summer_stage. The conductance of each record is a time-series called 'winter' or 'summer' with values of either 0 or 1. These conductance values are multiplied by an auxiliary variable that contains the actual conductance. @@ -969,6 +970,15 @@ def gdf_to_seasonal_pkg( ------- package : ModflowGwfdrn, ModflowGwfriv or ModflowGwfghb The generated flopy-package + + Notes + ----- + This function overlaps with ``nlmod.gwf.drain.drain_from_df`` for + polygon-to-DRN conversion. Use this function for surface-water polygons with + winter and summer stages and seasonal conductance timeseries. Use + ``drain_from_df`` for fixed drain features such as pipes, basins, point + drains, or direct cellid input, and when deterministic MVR provider IDs are + needed. """ if gdf.index.name != "cellid": # if "cellid" not in gdf: diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py index c5125273..24b29d1a 100644 --- a/tests/test_027_drain.py +++ b/tests/test_027_drain.py @@ -236,6 +236,157 @@ def test_drain_from_df_places_2d_cellids_in_layer_from_elevation(): _assert_mapping_matches_stress_period_data(drn, provider_mapping) +def test_drain_from_df_places_2d_cellids_below_inactive_top_layer(): + """Test 2D cell IDs skip inactive layers during layer placement.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[0, 0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [0, 1, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0)], + "elevation": [-1.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_inactive_top", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (1, 0, 0) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_keeps_valid_2d_cellids_after_omitting_inactive_column(): + """Test omitted 2D rows do not suppress later valid rows.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["top"].notnull() + ds["active_domain"].data[0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [0, 0, 0] + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 1].tolist() == [1, 1, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0), (0, 1)], + "elevation": [-1.0, -1.0], + "cond": [5.0, 7.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_mixed_active", + silent=True, + return_provider_mapping=True, + ) + + assert drn is not None + assert provider_mapping["mvr_provider_id"].tolist() == [0] + assert provider_mapping.loc[0, "cellid"] == (0, 0, 1) + assert provider_mapping.loc[0, "cond"] == 7.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_places_2d_cellids_below_pass_through_layer(): + """Test 2D cell IDs skip pass-through layers during layer placement.""" + ds = test_010_wells.get_model_ds() + ds["botm"].data[1, 0, 0] = ds["botm"].data[0, 0, 0] + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [1, -1, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0)], + "elevation": [-12.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_pass_through", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (2, 0, 0) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_omits_2d_cellids_without_active_layers(): + """Test 2D cell IDs are omitted when the column has no active layers.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["top"].notnull() + ds["active_domain"].data[0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [0, 0, 0] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0)], + "elevation": [-1.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_no_active_layers", + silent=True, + return_provider_mapping=True, + ) + + assert drn is None + assert provider_mapping.empty + + +@pytest.mark.parametrize( + ("cellid", "setup_idomain", "expected_idomain"), + [ + ((0, 0, 0), "inactive_top", [0, 1, 1]), + ((1, 0, 0), "pass_through_middle", [1, -1, 1]), + ], +) +def test_drain_from_df_rejects_3d_inactive_or_pass_through_cellids( + cellid, setup_idomain, expected_idomain +): + """Test explicit 3D cell IDs must target active cells.""" + ds = test_010_wells.get_model_ds() + if setup_idomain == "inactive_top": + ds["botm"].data[0, 0, 0] = ds["top"].data[0, 0] + elif setup_idomain == "pass_through_middle": + ds["botm"].data[1, 0, 0] = ds["botm"].data[0, 0, 0] + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == expected_idomain + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [cellid], + "elevation": [-1.0], + "cond": [5.0], + } + ) + + with pytest.raises(ValueError, match="inactive or pass-through"): + nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_bad_3d_idomain", + silent=True, + ) + + def test_drain_from_df_preserves_point_conductance(): """Test that point drains use supplied integrated conductance unchanged.""" ds = test_010_wells.get_model_ds() From b2c51b1651a168095453bc48b98d2cdc176cf307 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Thu, 9 Jul 2026 10:59:21 +0200 Subject: [PATCH 3/7] Remap explicit drain cells to active layers Move explicit 3D DRNs in inactive or pass-through cells to the nearest active layer in the same column. --- nlmod/gwf/drain.py | 83 +++++++++++++++--- tests/test_027_drain.py | 182 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 245 insertions(+), 20 deletions(-) diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index fdf4f1f1..0dc5acaa 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -112,9 +112,11 @@ def drain_from_df( ``nlmod.dims.layers.get_idomain`` to skip columns without active cells and to place the drain in a suitable active layer (``idomain > 0``), not in inactive (``idomain == 0``) or vertical pass-through (``idomain < 0``) cells. FloPy - receives explicit 3D DRN cellids and does not relocate boundaries to the first - active layer. Therefore explicit 3D cellids passed to this function are - checked and rejected when they target ``idomain <= 0``. + receives explicit 3D DRN cellids and does not relocate boundaries. Therefore + explicit 3D cellids passed to this function that target ``idomain <= 0`` are + remapped to the nearest active layer in the same vertical column, based on the + drain elevation. If no active layer exists, or if remapping is required for a + nonnumeric drain elevation, a ``ValueError`` is raised. """ logger.info("creating mf6 DRN from dataframe") @@ -410,13 +412,7 @@ def _is_3d_cellid(cellid, ds): def _record_from_3d_cellid(cellid, row, ds): - idomain_value = get_idomain(ds).data[cellid] - if idomain_value <= 0: - raise ValueError( - f"DRN cellid {cellid} is inactive or pass-through " - f"(idomain={idomain_value}). Use a 2D cellid or vector geometry to let " - "nlmod choose an active layer." - ) + cellid = _get_active_cellid_for_3d_cellid(cellid, row, ds) if np.isnan(row["cond"]): raise ValueError(f"Conductance is NaN in cell {cellid}") @@ -429,6 +425,73 @@ def _record_from_3d_cellid(cellid, row, ds): return [cellid, row["stage"], row["cond"]] + auxlist +def _get_active_cellid_for_3d_cellid(cellid, row, ds): + idomain = get_idomain(ds).data + if idomain[cellid] > 0: + return cellid + + try: + drain_elevation = float(row["stage"]) + except (TypeError, ValueError) as err: + raise ValueError( + f"Cannot remap DRN cellid {cellid} to the nearest active layer " + "because the drain elevation is not numeric." + ) from err + if not np.isfinite(drain_elevation): + raise ValueError( + f"Cannot remap DRN cellid {cellid} to the nearest active layer " + "because the drain elevation is not finite." + ) + + layer = cellid[0] + if ds.gridtype == "vertex": + icell2d = cellid[1] + idomain_column = idomain[:, icell2d] + layer_tops = np.r_[ds["top"].data[icell2d], ds["botm"].data[:-1, icell2d]] + layer_botms = ds["botm"].data[:, icell2d] + column_cellid = (icell2d,) + elif ds.gridtype == "structured": + row_index, column_index = cellid[1:] + idomain_column = idomain[:, row_index, column_index] + layer_tops = np.r_[ + ds["top"].data[row_index, column_index], + ds["botm"].data[:-1, row_index, column_index], + ] + layer_botms = ds["botm"].data[:, row_index, column_index] + column_cellid = (row_index, column_index) + else: + raise ValueError(f"Unsupported gridtype: {ds.gridtype}") + + active_layers = np.where(idomain_column > 0)[0] + if len(active_layers) == 0: + raise ValueError( + f"Cannot remap DRN cellid {cellid}; the vertical column has no active " + "layers." + ) + + nearest_layer = min( + active_layers, + key=lambda active_layer: ( + _distance_to_layer_interval( + drain_elevation, + layer_top=layer_tops[active_layer], + layer_botm=layer_botms[active_layer], + ), + abs(active_layer - layer), + active_layer, + ), + ) + return (nearest_layer,) + column_cellid + + +def _distance_to_layer_interval(elevation, layer_top, layer_botm): + upper = max(layer_top, layer_botm) + lower = min(layer_top, layer_botm) + if lower <= elevation <= upper: + return 0.0 + return min(abs(elevation - lower), abs(elevation - upper)) + + def _empty_celldata(): return pd.DataFrame( columns=[ diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py index 24b29d1a..d2bd88e3 100644 --- a/tests/test_027_drain.py +++ b/tests/test_027_drain.py @@ -3,6 +3,7 @@ import pandas as pd import pytest import test_010_wells +import util from shapely.geometry import LineString, MultiPoint, Polygon import nlmod @@ -188,7 +189,7 @@ def test_drain_from_df_preserves_direct_3d_cellids(): _, gwf = test_010_wells.get_sim_and_gwf(ds) drains = pd.DataFrame( { - "cellid": [(1, 0, 0)], + "cellid": [(0, 0, 0)], "elevation": [-12.0], "cond": [5.0], } @@ -203,12 +204,39 @@ def test_drain_from_df_preserves_direct_3d_cellids(): return_provider_mapping=True, ) - assert provider_mapping.loc[0, "cellid"] == (1, 0, 0) + assert provider_mapping.loc[0, "cellid"] == (0, 0, 0) assert provider_mapping.loc[0, "elev"] == -12.0 assert provider_mapping.loc[0, "cond"] == 5.0 _assert_mapping_matches_stress_period_data(drn, provider_mapping) +def test_drain_from_df_preserves_active_3d_cellid_with_timeseries_elevation(): + """Test active explicit 3D cell IDs can keep timeseries elevations.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0, 0)], + "elevation": ["stage_ts"], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_3d_ts", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (0, 0, 0) + assert provider_mapping.loc[0, "elev"] == "stage_ts" + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + def test_drain_from_df_places_2d_cellids_in_layer_from_elevation(): """Test that 2D cell IDs use drain elevation for layer placement.""" ds = test_010_wells.get_model_ds() @@ -352,32 +380,67 @@ def test_drain_from_df_omits_2d_cellids_without_active_layers(): @pytest.mark.parametrize( - ("cellid", "setup_idomain", "expected_idomain"), + ("cellid", "setup_idomain", "elevation", "expected_idomain", "expected_cellid"), [ - ((0, 0, 0), "inactive_top", [0, 1, 1]), - ((1, 0, 0), "pass_through_middle", [1, -1, 1]), + ((0, 0, 0), "inactive_top", -1.0, [0, 1, 1], (1, 0, 0)), + ((1, 0, 0), "pass_through_middle", -12.0, [1, -1, 1], (2, 0, 0)), + ((1, 0, 0), "inactive_middle", -12.0, [1, 0, 1], (0, 0, 0)), + ((1, 0, 0), "inactive_middle", -14.0, [1, 0, 1], (2, 0, 0)), ], ) -def test_drain_from_df_rejects_3d_inactive_or_pass_through_cellids( - cellid, setup_idomain, expected_idomain +def test_drain_from_df_remaps_3d_inactive_or_pass_through_cellids( + cellid, setup_idomain, elevation, expected_idomain, expected_cellid ): - """Test explicit 3D cell IDs must target active cells.""" + """Test explicit 3D cell IDs are remapped to the nearest active layer.""" ds = test_010_wells.get_model_ds() if setup_idomain == "inactive_top": ds["botm"].data[0, 0, 0] = ds["top"].data[0, 0] elif setup_idomain == "pass_through_middle": ds["botm"].data[1, 0, 0] = ds["botm"].data[0, 0, 0] + elif setup_idomain == "inactive_middle": + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0, 0] = False assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == expected_idomain _, gwf = test_010_wells.get_sim_and_gwf(ds) drains = pd.DataFrame( { "cellid": [cellid], + "elevation": [elevation], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_remapped_3d_idomain", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == expected_cellid + assert provider_mapping.loc[0, "elev"] == elevation + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_raises_for_3d_cellids_without_active_layers(): + """Test explicit 3D cell IDs still require an active column.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["top"].notnull() + ds["active_domain"].data[0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [0, 0, 0] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0, 0)], "elevation": [-1.0], "cond": [5.0], } ) - with pytest.raises(ValueError, match="inactive or pass-through"): + with pytest.raises(ValueError, match="no active layers"): nlmod.gwf.drain.drain_from_df( drains, gwf, @@ -387,6 +450,102 @@ def test_drain_from_df_rejects_3d_inactive_or_pass_through_cellids( ) +def test_drain_from_df_requires_numeric_elevation_to_remap_3d_cellids(): + """Test inactive explicit 3D cell IDs need numeric elevations for remapping.""" + ds = test_010_wells.get_model_ds() + ds["botm"].data[0, 0, 0] = ds["top"].data[0, 0] + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [0, 1, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0, 0)], + "elevation": ["stage_ts"], + "cond": [5.0], + } + ) + + with pytest.raises(ValueError, match="elevation is not numeric"): + nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_non_numeric_remap", + silent=True, + ) + + +def test_drain_from_df_remaps_vertex_3d_cellids_to_nearest_active_layer(): + """Test explicit vertex 3D cell IDs are remapped like structured cell IDs.""" + ds = util.get_ds_vertex( + model_name="drain_vertex", + top=0.0, + botm=[-10.0, -15.0, -30.0], + kh=[10.0, 0.1, 20.0], + kv=[5.0, 0.05, 10.0], + ) + ds = nlmod.time.set_ds_time(ds, "2023", time="2024") + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0].tolist() == [1, 0, 1] + gwf = util.get_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(1, 0)], + "elevation": [-14.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_vertex_remap", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (2, 0) + assert provider_mapping.loc[0, "elev"] == -14.0 + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + +def test_drain_from_df_preserves_active_vertex_3d_cellids(): + """Test active explicit vertex 3D cell IDs are preserved exactly.""" + ds = util.get_ds_vertex( + model_name="drn_vtx_active", + top=0.0, + botm=[-10.0, -15.0, -30.0], + kh=[10.0, 0.1, 20.0], + kv=[5.0, 0.05, 10.0], + ) + ds = nlmod.time.set_ds_time(ds, "2023", time="2024") + assert nlmod.dims.layers.get_idomain(ds).data[:, 1].tolist() == [1, 1, 1] + gwf = util.get_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 1)], + "elevation": [-14.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_vertex_active", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (0, 1) + assert provider_mapping.loc[0, "elev"] == -14.0 + assert provider_mapping.loc[0, "cond"] == 5.0 + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + def test_drain_from_df_preserves_point_conductance(): """Test that point drains use supplied integrated conductance unchanged.""" ds = test_010_wells.get_model_ds() @@ -444,7 +603,10 @@ def _assert_mapping_matches_stress_period_data(drn, provider_mapping): for _, row in provider_mapping.iterrows(): record = spd[int(row["mvr_provider_id"])] assert record["cellid"] == row["cellid"] - assert record["elev"] == pytest.approx(row["elev"]) + if isinstance(row["elev"], str): + assert record["elev"] == row["elev"] + else: + assert record["elev"] == pytest.approx(row["elev"]) assert record["cond"] == pytest.approx(row["cond"]) if "boundname" in record.dtype.names and pd.notna(row["boundname"]): assert record["boundname"] == row["boundname"] From db1bb43c1af422ef95f0ceba9df35dd0d6a9eedd Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Thu, 9 Jul 2026 11:25:31 +0200 Subject: [PATCH 4/7] Lean explicit drain remapping Share structured and vertex remapping logic and cover layered top datasets. --- nlmod/gwf/drain.py | 47 ++++++++++++++++++----------------------- tests/test_027_drain.py | 28 ++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 27 deletions(-) diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index 0dc5acaa..603bde60 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -363,6 +363,7 @@ def _base_celldata(df, elev, boundnames, mover_destinations): def _build_spd_with_provider_mapping(celldata, ds, layer_method, silent): spd = [] provider_mapping = [] + idomain = None for cellid, row in tqdm( celldata.iterrows(), total=celldata.index.size, @@ -370,7 +371,9 @@ def _build_spd_with_provider_mapping(celldata, ds, layer_method, silent): disable=silent, ): if _is_3d_cellid(cellid, ds): - row_spd = [_record_from_3d_cellid(cellid, row, ds)] + if idomain is None: + idomain = get_idomain(ds) + row_spd = [_record_from_3d_cellid(cellid, row, ds, idomain)] else: index = np.empty(1, dtype=object) index[0] = cellid @@ -411,8 +414,8 @@ def _is_3d_cellid(cellid, ds): raise ValueError(f"Unsupported gridtype: {ds.gridtype}") -def _record_from_3d_cellid(cellid, row, ds): - cellid = _get_active_cellid_for_3d_cellid(cellid, row, ds) +def _record_from_3d_cellid(cellid, row, ds, idomain): + cellid = _get_active_cellid_for_3d_cellid(cellid, row, ds, idomain) if np.isnan(row["cond"]): raise ValueError(f"Conductance is NaN in cell {cellid}") @@ -425,9 +428,8 @@ def _record_from_3d_cellid(cellid, row, ds): return [cellid, row["stage"], row["cond"]] + auxlist -def _get_active_cellid_for_3d_cellid(cellid, row, ds): - idomain = get_idomain(ds).data - if idomain[cellid] > 0: +def _get_active_cellid_for_3d_cellid(cellid, row, ds, idomain): + if idomain.data[cellid] > 0: return cellid try: @@ -443,24 +445,17 @@ def _get_active_cellid_for_3d_cellid(cellid, row, ds): "because the drain elevation is not finite." ) - layer = cellid[0] - if ds.gridtype == "vertex": - icell2d = cellid[1] - idomain_column = idomain[:, icell2d] - layer_tops = np.r_[ds["top"].data[icell2d], ds["botm"].data[:-1, icell2d]] - layer_botms = ds["botm"].data[:, icell2d] - column_cellid = (icell2d,) - elif ds.gridtype == "structured": - row_index, column_index = cellid[1:] - idomain_column = idomain[:, row_index, column_index] - layer_tops = np.r_[ - ds["top"].data[row_index, column_index], - ds["botm"].data[:-1, row_index, column_index], - ] - layer_botms = ds["botm"].data[:, row_index, column_index] - column_cellid = (row_index, column_index) + column_cellid = cellid[1:] + column_indexer = dict(zip(ds["botm"].dims[1:], column_cellid, strict=True)) + idomain_column = idomain.isel(column_indexer).data + layer_botms = ds["botm"].isel(column_indexer).data + top_indexer = { + dim: index for dim, index in column_indexer.items() if dim in ds["top"].dims + } + if "layer" in ds["top"].dims: + layer_tops = ds["top"].isel(top_indexer).data else: - raise ValueError(f"Unsupported gridtype: {ds.gridtype}") + layer_tops = np.r_[ds["top"].isel(top_indexer).data, layer_botms[:-1]] active_layers = np.where(idomain_column > 0)[0] if len(active_layers) == 0: @@ -477,7 +472,7 @@ def _get_active_cellid_for_3d_cellid(cellid, row, ds): layer_top=layer_tops[active_layer], layer_botm=layer_botms[active_layer], ), - abs(active_layer - layer), + abs(active_layer - cellid[0]), active_layer, ), ) @@ -487,9 +482,7 @@ def _get_active_cellid_for_3d_cellid(cellid, row, ds): def _distance_to_layer_interval(elevation, layer_top, layer_botm): upper = max(layer_top, layer_botm) lower = min(layer_top, layer_botm) - if lower <= elevation <= upper: - return 0.0 - return min(abs(elevation - lower), abs(elevation - upper)) + return abs(elevation - np.clip(elevation, lower, upper)) def _empty_celldata(): diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py index d2bd88e3..cdd999ec 100644 --- a/tests/test_027_drain.py +++ b/tests/test_027_drain.py @@ -425,6 +425,34 @@ def test_drain_from_df_remaps_3d_inactive_or_pass_through_cellids( _assert_mapping_matches_stress_period_data(drn, provider_mapping) +def test_drain_from_df_remaps_3d_cellids_with_layered_top(): + """Test remapping uses layer-specific top when top has a layer dimension.""" + ds = test_010_wells.get_model_ds() + ds["botm"].data[1, 0, 0] = ds["botm"].data[0, 0, 0] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + ds = nlmod.layers.add_layer_dim_to_top(ds) + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [1, -1, 1] + drains = pd.DataFrame( + { + "cellid": [(1, 0, 0)], + "elevation": [-12.0], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_layered_top", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping.loc[0, "cellid"] == (2, 0, 0) + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + def test_drain_from_df_raises_for_3d_cellids_without_active_layers(): """Test explicit 3D cell IDs still require an active column.""" ds = test_010_wells.get_model_ds() From 681ecf093c026069d3b1c26c38218588a6881961 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Thu, 9 Jul 2026 11:37:11 +0200 Subject: [PATCH 5/7] Tighten drain remapping helpers Reuse IDOMAIN during explicit 3D remapping and cover layered-top datasets. --- nlmod/gwf/drain.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index 603bde60..3d90c202 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -297,7 +297,16 @@ def _drain_celldata_from_df( raise TypeError(f"Unsupported drain geometry types: {unsupported_geom_types}") if len(parts) == 0: - return _empty_celldata() + return pd.DataFrame( + columns=[ + "stage", + "rbot", + "cond", + "area", + "len_estimate", + "source_index", + ] + ) return pd.concat(parts, axis=0) @@ -485,19 +494,6 @@ def _distance_to_layer_interval(elevation, layer_top, layer_botm): return abs(elevation - np.clip(elevation, lower, upper)) -def _empty_celldata(): - return pd.DataFrame( - columns=[ - "stage", - "rbot", - "cond", - "area", - "len_estimate", - "source_index", - ] - ) - - def _validate_columns(df, columns, name): missing = set(columns).difference(df.columns) if missing: From 419e0ad92283efd808294333bb7ff57415ffa9d9 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Thu, 9 Jul 2026 11:57:41 +0200 Subject: [PATCH 6/7] Move nearest active layer lookup to layers Expose nearest active layer selection for boundary packages and use it from drain_from_df. --- nlmod/dims/layers.py | 67 ++++++++++++++++++++++++++++++++++++++++ nlmod/gwf/drain.py | 55 ++++----------------------------- tests/test_009_layers.py | 27 ++++++++++++++-- tests/test_027_drain.py | 3 +- 4 files changed, 100 insertions(+), 52 deletions(-) diff --git a/nlmod/dims/layers.py b/nlmod/dims/layers.py index acce3c8c..97c28d4d 100644 --- a/nlmod/dims/layers.py +++ b/nlmod/dims/layers.py @@ -1611,6 +1611,73 @@ def get_layer_of_z(ds, z, above_model=-999, below_model=-999): return layer +def get_nearest_active_layer(ds, cellid, z, idomain=None, preferred_layer=None): + """Get the active layer nearest to a z-value in one vertical column. + + Parameters + ---------- + ds : xarray.Dataset + Model Dataset. + cellid : int or tuple of int + Two-dimensional cellid of the vertical column. Use an integer for vertex + grids and ``(row, column)`` for structured grids. + z : float + Elevation for which the nearest active layer is determined. + idomain : xarray.DataArray, optional + Idomain array. If None, it is calculated from ``ds``. The default is + None. + preferred_layer : int, optional + Layer used as a secondary tie-breaker when multiple active layers are + equally near to ``z``. The default is None, which breaks ties to the + shallower layer. + + Returns + ------- + layer : int + Zero-based layer index of the active layer nearest to ``z``. + """ + try: + z = float(z) + except (TypeError, ValueError) as err: + raise ValueError( + "z must be numeric to determine the nearest active layer" + ) from err + if not np.isfinite(z): + raise ValueError("z must be finite to determine the nearest active layer") + + if idomain is None: + idomain = get_idomain(ds) + + column_cellid = (cellid,) if not isinstance(cellid, tuple) else cellid + column_indexer = dict(zip(ds["botm"].dims[1:], column_cellid, strict=True)) + idomain_column = idomain.isel(column_indexer).data + active_layers = np.where(idomain_column > 0)[0] + if len(active_layers) == 0: + raise ValueError("the vertical column has no active layers") + + layer_botms = ds["botm"].isel(column_indexer).data + top_indexer = { + dim: index for dim, index in column_indexer.items() if dim in ds["top"].dims + } + if "layer" in ds["top"].dims: + layer_tops = ds["top"].isel(top_indexer).data + else: + layer_tops = np.r_[ds["top"].isel(top_indexer).data, layer_botms[:-1]] + + distances = np.abs( + z + - np.clip( + z, + np.minimum(layer_tops[active_layers], layer_botms[active_layers]), + np.maximum(layer_tops[active_layers], layer_botms[active_layers]), + ) + ) + if preferred_layer is None: + return int(active_layers[np.lexsort((active_layers, distances))[0]]) + layer_distance = np.abs(active_layers - preferred_layer) + return int(active_layers[np.lexsort((active_layers, layer_distance, distances))[0]]) + + def update_idomain_from_thickness(idomain, thickness, mask): """Get new idomain from thickness in the cells where mask is 1 (or True). diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index 3d90c202..b0e6fb2a 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -6,7 +6,7 @@ import pandas as pd from ..dims.grid import gdf_to_grid -from ..dims.layers import get_idomain +from ..dims.layers import get_idomain, get_nearest_active_layer from ..util import tqdm from .surface_water import build_spd @@ -441,59 +441,16 @@ def _get_active_cellid_for_3d_cellid(cellid, row, ds, idomain): if idomain.data[cellid] > 0: return cellid - try: - drain_elevation = float(row["stage"]) - except (TypeError, ValueError) as err: - raise ValueError( - f"Cannot remap DRN cellid {cellid} to the nearest active layer " - "because the drain elevation is not numeric." - ) from err - if not np.isfinite(drain_elevation): - raise ValueError( - f"Cannot remap DRN cellid {cellid} to the nearest active layer " - "because the drain elevation is not finite." - ) - column_cellid = cellid[1:] - column_indexer = dict(zip(ds["botm"].dims[1:], column_cellid, strict=True)) - idomain_column = idomain.isel(column_indexer).data - layer_botms = ds["botm"].isel(column_indexer).data - top_indexer = { - dim: index for dim, index in column_indexer.items() if dim in ds["top"].dims - } - if "layer" in ds["top"].dims: - layer_tops = ds["top"].isel(top_indexer).data - else: - layer_tops = np.r_[ds["top"].isel(top_indexer).data, layer_botms[:-1]] - - active_layers = np.where(idomain_column > 0)[0] - if len(active_layers) == 0: - raise ValueError( - f"Cannot remap DRN cellid {cellid}; the vertical column has no active " - "layers." + try: + nearest_layer = get_nearest_active_layer( + ds, column_cellid, row["stage"], idomain=idomain, preferred_layer=cellid[0] ) - - nearest_layer = min( - active_layers, - key=lambda active_layer: ( - _distance_to_layer_interval( - drain_elevation, - layer_top=layer_tops[active_layer], - layer_botm=layer_botms[active_layer], - ), - abs(active_layer - cellid[0]), - active_layer, - ), - ) + except ValueError as err: + raise ValueError(f"Cannot remap DRN cellid {cellid}; {err}") from err return (nearest_layer,) + column_cellid -def _distance_to_layer_interval(elevation, layer_top, layer_botm): - upper = max(layer_top, layer_botm) - lower = min(layer_top, layer_botm) - return abs(elevation - np.clip(elevation, lower, upper)) - - def _validate_columns(df, columns, name): missing = set(columns).difference(df.columns) if missing: diff --git a/tests/test_009_layers.py b/tests/test_009_layers.py index d89cea5a..72ad139b 100644 --- a/tests/test_009_layers.py +++ b/tests/test_009_layers.py @@ -1,16 +1,16 @@ # %% import os -import pytest import matplotlib.pyplot as plt import numpy as np +import pytest import test_001_model +import util from pandas import DataFrame from shapely.geometry import LineString import nlmod from nlmod.plot import DatasetCrossSection -import util MODEL_DATA_ENV_VAR = "NLMOD_TEST_MODEL_DATA_DIR" @@ -200,6 +200,29 @@ def test_get_layer_of_z_below_model(): assert (layer == layer.attrs["below_model"]).all() +def test_get_nearest_active_layer(): + ds = nlmod.get_ds([0, 1000, 0, 500], top=0, botm=[-10, -15, -30]) + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0, 0] = False + idomain = nlmod.layers.get_idomain(ds) + + assert idomain.data[:, 0, 0].tolist() == [1, 0, 1] + assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -12.0, idomain) == 0 + assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -14.0, idomain) == 2 + assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -12.5, idomain) == 0 + assert ( + nlmod.layers.get_nearest_active_layer( + ds, (0, 0), -12.5, idomain, preferred_layer=2 + ) + == 2 + ) + + ds["active_domain"].data[:, 0, 1] = False + idomain = nlmod.layers.get_idomain(ds) + with pytest.raises(ValueError, match="no active layers"): + nlmod.layers.get_nearest_active_layer(ds, (0, 1), -12.0, idomain) + + def test_aggregate_by_weighted_mean_to_ds(): regis = get_regis_horstermeer() regis2 = regis.copy(deep=True) diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py index cdd999ec..ec4b4ab9 100644 --- a/tests/test_027_drain.py +++ b/tests/test_027_drain.py @@ -386,6 +386,7 @@ def test_drain_from_df_omits_2d_cellids_without_active_layers(): ((1, 0, 0), "pass_through_middle", -12.0, [1, -1, 1], (2, 0, 0)), ((1, 0, 0), "inactive_middle", -12.0, [1, 0, 1], (0, 0, 0)), ((1, 0, 0), "inactive_middle", -14.0, [1, 0, 1], (2, 0, 0)), + ((1, 0, 0), "inactive_middle", -12.5, [1, 0, 1], (0, 0, 0)), ], ) def test_drain_from_df_remaps_3d_inactive_or_pass_through_cellids( @@ -492,7 +493,7 @@ def test_drain_from_df_requires_numeric_elevation_to_remap_3d_cellids(): } ) - with pytest.raises(ValueError, match="elevation is not numeric"): + with pytest.raises(ValueError, match="numeric"): nlmod.gwf.drain.drain_from_df( drains, gwf, From 317ada435e5a791d0de3572f75c4272db62b4819 Mon Sep 17 00:00:00 2001 From: Bas des Tombe Date: Thu, 9 Jul 2026 14:42:09 +0200 Subject: [PATCH 7/7] Extend get_layer_of_z for sparse active lookups Support scalar and vector cellids with nearest-active layer selection and use it for batched drain remapping. --- nlmod/dims/layers.py | 197 +++++++++++++++++++++++++-------------- nlmod/gwf/drain.py | 85 +++++++++-------- tests/test_009_layers.py | 85 +++++++++++++++-- tests/test_027_drain.py | 81 +++++++++++++++- 4 files changed, 329 insertions(+), 119 deletions(-) diff --git a/nlmod/dims/layers.py b/nlmod/dims/layers.py index 97c28d4d..ff6e0e52 100644 --- a/nlmod/dims/layers.py +++ b/nlmod/dims/layers.py @@ -1572,7 +1572,16 @@ def get_last_active_layer_from_idomain(idomain, nodata=-999): return last_active_layer -def get_layer_of_z(ds, z, above_model=-999, below_model=-999): +def get_layer_of_z( + ds, + z, + above_model=-999, + below_model=-999, + cellid=None, + idomain=None, + nearest_active=False, + preferred_layer=None, +): """Get the layer of a certain z-value in all cells from a model ds. Parameters @@ -1585,13 +1594,126 @@ def get_layer_of_z(ds, z, above_model=-999, below_model=-999): value used for cells where z is above the top of the model. The default is -999. below_model : int, optional value used for cells where z is below the top of the model. The default is -999. + cellid : int, tuple of int or sequence, optional + Two-dimensional cellid(s) of vertical column(s). Use integers for vertex grids + and ``(row, column)`` tuples for structured grids. If None, layers are returned + for all model cells. The default is None. + idomain : xarray.DataArray, optional + Idomain array. Only used when ``cellid`` is not None and ``nearest_active`` is + True. If None, it is calculated from ``ds``. The default is None. + nearest_active : bool, optional + If True, return the active layer nearest to ``z`` in each vertical column, + ignoring inactive and pass-through cells. The default is False. + preferred_layer : int or sequence of int, optional + Layer(s) used as a secondary tie-breaker when multiple active layers are + equally near to ``z``. The default is None, which breaks ties to the shallower + layer. Returns ------- - layer : xr.DataArray - DataArray with values representing the integer layer index. Shape can be (y, x) - or (icell2d) + layer : xr.DataArray, int or np.ndarray + DataArray with values representing the integer layer index when ``cellid`` is + None. If ``cellid`` is supplied, returns an int for one cellid or an array for + multiple cellids. """ + if cellid is not None: + scalar_cellid = isinstance(cellid, (int, np.integer)) or ( + isinstance(cellid, tuple) + and all(isinstance(part, (int, np.integer)) for part in cellid) + ) + cellids = [cellid] if scalar_cellid else list(cellid) + column_cellids = [ + (item,) if isinstance(item, (int, np.integer)) else tuple(item) + for item in cellids + ] + try: + z = np.asarray(z, dtype=float) + except (TypeError, ValueError) as err: + raise ValueError("z must be numeric to determine the layer") from err + if z.ndim == 0: + z = np.full(len(column_cellids), float(z)) + elif z.size == 1 and len(column_cellids) != 1: + z = np.full(len(column_cellids), float(z.ravel()[0])) + elif z.size != len(column_cellids): + raise ValueError("z must be scalar or have the same length as cellid") + else: + z = z.ravel() + if not np.isfinite(z).all(): + raise ValueError("z must be finite to determine the layer") + + cellids_array = np.asarray(column_cellids, dtype=int) + if cellids_array.ndim == 1: + cellids_array = cellids_array[:, np.newaxis] + if cellids_array.shape[1] != len(ds["botm"].dims) - 1: + raise ValueError("cellid does not match the model grid dimensions") + + layer_botms = ds["botm"].data[(slice(None), *cellids_array.T)] + if "layer" in ds["top"].dims: + layer_tops = ds["top"].data[(slice(None), *cellids_array.T)] + top0 = layer_tops[0] + else: + top0 = ds["top"].data[tuple(cellids_array.T)] + layer_tops = np.vstack((top0, layer_botms[:-1])) + + if not nearest_active: + below_z = layer_botms < z + layer = np.where( + below_z.any(axis=0), np.argmax(below_z, axis=0), below_model + ) + layer = np.where(z > top0, above_model, layer) + return int(layer[0]) if scalar_cellid else layer + + if idomain is None: + idomain = get_idomain(ds) + idomain_column = idomain.data[(slice(None), *cellids_array.T)] + active = idomain_column > 0 + if not active.any(axis=0).all(): + bad = np.where(~active.any(axis=0))[0] + raise ValueError( + "the vertical column has no active layers for cellid positions " + f"{bad.tolist()}" + ) + + distances = np.abs( + z + - np.clip( + z, + np.minimum(layer_tops, layer_botms), + np.maximum(layer_tops, layer_botms), + ) + ) + distances[~active] = np.inf + candidates = distances == distances.min(axis=0) + + if preferred_layer is not None: + preferred_layer = np.asarray(preferred_layer, dtype=int) + if preferred_layer.ndim == 0: + preferred_layer = np.full(len(column_cellids), int(preferred_layer)) + elif preferred_layer.size == 1 and len(column_cellids) != 1: + preferred_layer = np.full( + len(column_cellids), int(preferred_layer.ravel()[0]) + ) + elif preferred_layer.size != len(column_cellids): + raise ValueError( + "preferred_layer must be scalar or have the same length as cellid" + ) + else: + preferred_layer = preferred_layer.ravel() + + layer_number = np.arange(len(ds.layer))[:, np.newaxis] + layer_distance = np.abs(layer_number - preferred_layer) + layer_distance[~candidates] = np.iinfo(layer_distance.dtype).max + candidates &= layer_distance == layer_distance.min(axis=0) + unresolved = np.where(candidates.sum(axis=0) > 1)[0] + if len(unresolved) > 0: + raise ValueError( + "multiple active layers are equally near to z for cellid " + f"positions {unresolved.tolist()}" + ) + + layer = np.argmax(candidates, axis=0) + return int(layer[0]) if scalar_cellid else layer + layer = xr.where(ds["botm"][0] < z, 0, below_model) for i in range(1, len(ds.layer)): layer = xr.where((layer == below_model) & (ds["botm"][i] < z), i, layer) @@ -1611,73 +1733,6 @@ def get_layer_of_z(ds, z, above_model=-999, below_model=-999): return layer -def get_nearest_active_layer(ds, cellid, z, idomain=None, preferred_layer=None): - """Get the active layer nearest to a z-value in one vertical column. - - Parameters - ---------- - ds : xarray.Dataset - Model Dataset. - cellid : int or tuple of int - Two-dimensional cellid of the vertical column. Use an integer for vertex - grids and ``(row, column)`` for structured grids. - z : float - Elevation for which the nearest active layer is determined. - idomain : xarray.DataArray, optional - Idomain array. If None, it is calculated from ``ds``. The default is - None. - preferred_layer : int, optional - Layer used as a secondary tie-breaker when multiple active layers are - equally near to ``z``. The default is None, which breaks ties to the - shallower layer. - - Returns - ------- - layer : int - Zero-based layer index of the active layer nearest to ``z``. - """ - try: - z = float(z) - except (TypeError, ValueError) as err: - raise ValueError( - "z must be numeric to determine the nearest active layer" - ) from err - if not np.isfinite(z): - raise ValueError("z must be finite to determine the nearest active layer") - - if idomain is None: - idomain = get_idomain(ds) - - column_cellid = (cellid,) if not isinstance(cellid, tuple) else cellid - column_indexer = dict(zip(ds["botm"].dims[1:], column_cellid, strict=True)) - idomain_column = idomain.isel(column_indexer).data - active_layers = np.where(idomain_column > 0)[0] - if len(active_layers) == 0: - raise ValueError("the vertical column has no active layers") - - layer_botms = ds["botm"].isel(column_indexer).data - top_indexer = { - dim: index for dim, index in column_indexer.items() if dim in ds["top"].dims - } - if "layer" in ds["top"].dims: - layer_tops = ds["top"].isel(top_indexer).data - else: - layer_tops = np.r_[ds["top"].isel(top_indexer).data, layer_botms[:-1]] - - distances = np.abs( - z - - np.clip( - z, - np.minimum(layer_tops[active_layers], layer_botms[active_layers]), - np.maximum(layer_tops[active_layers], layer_botms[active_layers]), - ) - ) - if preferred_layer is None: - return int(active_layers[np.lexsort((active_layers, distances))[0]]) - layer_distance = np.abs(active_layers - preferred_layer) - return int(active_layers[np.lexsort((active_layers, layer_distance, distances))[0]]) - - def update_idomain_from_thickness(idomain, thickness, mask): """Get new idomain from thickness in the cells where mask is 1 (or True). diff --git a/nlmod/gwf/drain.py b/nlmod/gwf/drain.py index b0e6fb2a..1a9bbf5d 100644 --- a/nlmod/gwf/drain.py +++ b/nlmod/gwf/drain.py @@ -6,7 +6,7 @@ import pandas as pd from ..dims.grid import gdf_to_grid -from ..dims.layers import get_idomain, get_nearest_active_layer +from ..dims.layers import get_idomain, get_layer_of_z from ..util import tqdm from .surface_water import build_spd @@ -372,19 +372,52 @@ def _base_celldata(df, elev, boundnames, mover_destinations): def _build_spd_with_provider_mapping(celldata, ds, layer_method, silent): spd = [] provider_mapping = [] - idomain = None - for cellid, row in tqdm( - celldata.iterrows(), - total=celldata.index.size, - desc="Building stress period data DRN", - disable=silent, + index = np.empty(1, dtype=object) + cellids = celldata.index.to_list() + is_3d_cellid = np.array([_is_3d_cellid(cellid, ds) for cellid in cellids]) + if is_3d_cellid.any(): + idomain = get_idomain(ds) + positions = np.where(is_3d_cellid)[0] + original_cellids = [cellids[position] for position in positions] + active = np.array([idomain.data[cellid] > 0 for cellid in original_cellids]) + remap_positions = positions[~active] + if len(remap_positions) > 0: + try: + remap_layers = get_layer_of_z( + ds, + celldata.iloc[remap_positions]["stage"].to_numpy(), + cellid=[cellids[position][1:] for position in remap_positions], + idomain=idomain, + nearest_active=True, + preferred_layer=[ + cellids[position][0] for position in remap_positions + ], + ) + except ValueError as err: + raise ValueError(f"Cannot remap DRN cellids; {err}") from err + for position, layer in zip(remap_positions, remap_layers, strict=True): + cellids[position] = (layer,) + cellids[position][1:] + + for position, (cellid, row) in enumerate( + tqdm( + celldata.iterrows(), + total=celldata.index.size, + desc="Building stress period data DRN", + disable=silent, + ) ): - if _is_3d_cellid(cellid, ds): - if idomain is None: - idomain = get_idomain(ds) - row_spd = [_record_from_3d_cellid(cellid, row, ds, idomain)] + if is_3d_cellid[position]: + cellid = cellids[position] + if pd.isna(row["stage"]): + row_spd = [] + continue + if np.isnan(row["cond"]): + raise ValueError(f"Conductance is NaN in cell {cellid}") + if row["cond"] < 0: + raise ValueError(f"Conductance is negative in cell {cellid}") + auxlist = [row["boundname"]] if "boundname" in row else [] + row_spd = [[cellid, row["stage"], row["cond"]] + auxlist] else: - index = np.empty(1, dtype=object) index[0] = cellid row_df = pd.DataFrame( [row], @@ -423,34 +456,6 @@ def _is_3d_cellid(cellid, ds): raise ValueError(f"Unsupported gridtype: {ds.gridtype}") -def _record_from_3d_cellid(cellid, row, ds, idomain): - cellid = _get_active_cellid_for_3d_cellid(cellid, row, ds, idomain) - - if np.isnan(row["cond"]): - raise ValueError(f"Conductance is NaN in cell {cellid}") - if row["cond"] < 0: - raise ValueError(f"Conductance is negative in cell {cellid}") - - auxlist = [] - if "boundname" in row: - auxlist.append(row["boundname"]) - return [cellid, row["stage"], row["cond"]] + auxlist - - -def _get_active_cellid_for_3d_cellid(cellid, row, ds, idomain): - if idomain.data[cellid] > 0: - return cellid - - column_cellid = cellid[1:] - try: - nearest_layer = get_nearest_active_layer( - ds, column_cellid, row["stage"], idomain=idomain, preferred_layer=cellid[0] - ) - except ValueError as err: - raise ValueError(f"Cannot remap DRN cellid {cellid}; {err}") from err - return (nearest_layer,) + column_cellid - - def _validate_columns(df, columns, name): missing = set(columns).difference(df.columns) if missing: diff --git a/tests/test_009_layers.py b/tests/test_009_layers.py index 72ad139b..46383d5d 100644 --- a/tests/test_009_layers.py +++ b/tests/test_009_layers.py @@ -200,27 +200,98 @@ def test_get_layer_of_z_below_model(): assert (layer == layer.attrs["below_model"]).all() -def test_get_nearest_active_layer(): +def test_get_layer_of_z_nearest_active_cellids(): ds = nlmod.get_ds([0, 1000, 0, 500], top=0, botm=[-10, -15, -30]) ds["active_domain"] = ds["botm"].notnull() ds["active_domain"].data[1, 0, 0] = False idomain = nlmod.layers.get_idomain(ds) assert idomain.data[:, 0, 0].tolist() == [1, 0, 1] - assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -12.0, idomain) == 0 - assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -14.0, idomain) == 2 - assert nlmod.layers.get_nearest_active_layer(ds, (0, 0), -12.5, idomain) == 0 assert ( - nlmod.layers.get_nearest_active_layer( - ds, (0, 0), -12.5, idomain, preferred_layer=2 + nlmod.layers.get_layer_of_z( + ds, -12.0, cellid=(0, 0), idomain=idomain, nearest_active=True + ) + == 0 + ) + assert ( + nlmod.layers.get_layer_of_z( + ds, -14.0, cellid=(0, 0), idomain=idomain, nearest_active=True + ) + == 2 + ) + assert ( + nlmod.layers.get_layer_of_z( + ds, -12.5, cellid=(0, 0), idomain=idomain, nearest_active=True + ) + == 0 + ) + assert ( + nlmod.layers.get_layer_of_z( + ds, + -12.5, + cellid=(0, 0), + idomain=idomain, + nearest_active=True, + preferred_layer=2, ) == 2 ) + with pytest.raises(ValueError, match="equally near"): + nlmod.layers.get_layer_of_z( + ds, + -12.5, + cellid=(0, 0), + idomain=idomain, + nearest_active=True, + preferred_layer=1, + ) ds["active_domain"].data[:, 0, 1] = False idomain = nlmod.layers.get_idomain(ds) with pytest.raises(ValueError, match="no active layers"): - nlmod.layers.get_nearest_active_layer(ds, (0, 1), -12.0, idomain) + nlmod.layers.get_layer_of_z( + ds, -12.0, cellid=(0, 1), idomain=idomain, nearest_active=True + ) + + +def test_get_layer_of_z_nearest_active_vector_cellids(): + ds = nlmod.get_ds([0, 1000, 0, 500], top=0, botm=[-10, -15, -30]) + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0, 0] = False + ds["active_domain"].data[0, 0, 1] = False + idomain = nlmod.layers.get_idomain(ds) + + layers = nlmod.layers.get_layer_of_z( + ds, + [-12.0, -14.0, -1.0, 1.0, -40.0], + cellid=[(0, 0), (0, 0), (0, 1), (0, 2), (0, 2)], + idomain=idomain, + nearest_active=True, + preferred_layer=[1, 1, 0, 0, 2], + ) + + np.testing.assert_array_equal(layers, [0, 2, 1, 0, 2]) + + layers = nlmod.layers.get_layer_of_z( + ds, + [-12.5, -12.5], + cellid=[(0, 0), (0, 0)], + idomain=idomain, + nearest_active=True, + preferred_layer=[0, 2], + ) + np.testing.assert_array_equal(layers, [0, 2]) + + ds["active_domain"].data[:, 0, 2] = False + idomain = nlmod.layers.get_idomain(ds) + with pytest.raises(ValueError, match="no active layers"): + nlmod.layers.get_layer_of_z( + ds, + [-12.0, -12.0], + cellid=[(0, 0), (0, 2)], + idomain=idomain, + nearest_active=True, + ) def test_aggregate_by_weighted_mean_to_ds(): diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py index ec4b4ab9..6a23e370 100644 --- a/tests/test_027_drain.py +++ b/tests/test_027_drain.py @@ -237,6 +237,31 @@ def test_drain_from_df_preserves_active_3d_cellid_with_timeseries_elevation(): _assert_mapping_matches_stress_period_data(drn, provider_mapping) +def test_drain_from_df_omits_3d_cellids_without_elevation(): + """Test explicit 3D cell IDs without elevation are omitted.""" + ds = test_010_wells.get_model_ds() + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0, 0)], + "elevation": [np.nan], + "cond": [5.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_3d_nan", + silent=True, + return_provider_mapping=True, + ) + + assert drn is None + assert provider_mapping.empty + + def test_drain_from_df_places_2d_cellids_in_layer_from_elevation(): """Test that 2D cell IDs use drain elevation for layer placement.""" ds = test_010_wells.get_model_ds() @@ -386,7 +411,6 @@ def test_drain_from_df_omits_2d_cellids_without_active_layers(): ((1, 0, 0), "pass_through_middle", -12.0, [1, -1, 1], (2, 0, 0)), ((1, 0, 0), "inactive_middle", -12.0, [1, 0, 1], (0, 0, 0)), ((1, 0, 0), "inactive_middle", -14.0, [1, 0, 1], (2, 0, 0)), - ((1, 0, 0), "inactive_middle", -12.5, [1, 0, 1], (0, 0, 0)), ], ) def test_drain_from_df_remaps_3d_inactive_or_pass_through_cellids( @@ -426,6 +450,61 @@ def test_drain_from_df_remaps_3d_inactive_or_pass_through_cellids( _assert_mapping_matches_stress_period_data(drn, provider_mapping) +def test_drain_from_df_raises_for_ambiguous_nearest_active_layer(): + """Test unresolved nearest-active ties require explicit modeler input.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [1, 0, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(1, 0, 0)], + "elevation": [-12.5], + "cond": [5.0], + } + ) + + with pytest.raises(ValueError, match="equally near"): + nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_ambiguous_remap", + silent=True, + ) + + +def test_drain_from_df_remaps_multiple_3d_cellids_in_one_call(): + """Test batched explicit 3D remapping keeps row-specific layer choices.""" + ds = test_010_wells.get_model_ds() + ds["active_domain"] = ds["botm"].notnull() + ds["active_domain"].data[1, 0, 0] = False + assert nlmod.dims.layers.get_idomain(ds).data[:, 0, 0].tolist() == [1, 0, 1] + _, gwf = test_010_wells.get_sim_and_gwf(ds) + drains = pd.DataFrame( + { + "cellid": [(0, 0, 0), (1, 0, 0), (1, 0, 0)], + "elevation": [-1.0, -12.0, -14.0], + "cond": [3.0, 5.0, 7.0], + } + ) + + drn, provider_mapping = nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_multi_remap", + silent=True, + return_provider_mapping=True, + ) + + assert provider_mapping["cellid"].tolist() == [(0, 0, 0), (0, 0, 0), (2, 0, 0)] + assert provider_mapping["elev"].tolist() == [-1.0, -12.0, -14.0] + assert provider_mapping["cond"].tolist() == [3.0, 5.0, 7.0] + _assert_mapping_matches_stress_period_data(drn, provider_mapping) + + def test_drain_from_df_remaps_3d_cellids_with_layered_top(): """Test remapping uses layer-specific top when top has a layer dimension.""" ds = test_010_wells.get_model_ds()