diff --git a/nlmod/dims/layers.py b/nlmod/dims/layers.py index acce3c8c..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) 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..1a9bbf5d --- /dev/null +++ b/nlmod/gwf/drain.py @@ -0,0 +1,462 @@ +import logging + +import flopy +import geopandas as gpd +import numpy as np +import pandas as pd + +from ..dims.grid import gdf_to_grid +from ..dims.layers import get_idomain, get_layer_of_z +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)``. + + 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. 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") + + 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 pd.DataFrame( + columns=[ + "stage", + "rbot", + "cond", + "area", + "len_estimate", + "source_index", + ] + ) + 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 = [] + 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[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[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 _validate_columns(df, columns, name): + missing = set(columns).difference(df.columns) + if missing: + raise ValueError(f"Missing columns in {name}: {missing}") 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_009_layers.py b/tests/test_009_layers.py index d89cea5a..46383d5d 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,100 @@ def test_get_layer_of_z_below_model(): assert (layer == layer.attrs["below_model"]).all() +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_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_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(): regis = get_regis_horstermeer() regis2 = regis.copy(deep=True) diff --git a/tests/test_027_drain.py b/tests/test_027_drain.py new file mode 100644 index 00000000..6a23e370 --- /dev/null +++ b/tests/test_027_drain.py @@ -0,0 +1,720 @@ +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +import test_010_wells +import util +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": [(0, 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"] == (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_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() + _, 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_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", "elevation", "expected_idomain", "expected_cellid"), + [ + ((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_remaps_3d_inactive_or_pass_through_cellids( + cellid, setup_idomain, elevation, expected_idomain, expected_cellid +): + """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_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() + 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() + 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="no active layers"): + nlmod.gwf.drain.drain_from_df( + drains, + gwf, + ds, + pname="drn_bad_3d_idomain", + silent=True, + ) + + +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="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() + _, 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"] + 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"]