diff --git a/imod/mf6/__init__.py b/imod/mf6/__init__.py index b183fc772..3785b10ca 100644 --- a/imod/mf6/__init__.py +++ b/imod/mf6/__init__.py @@ -4,6 +4,7 @@ from imod.mf6.chd import ConstantHead from imod.mf6.dis import StructuredDiscretization +from imod.mf6.disu import LowLevelUnstructuredDiscretization from imod.mf6.disv import VerticesDiscretization from imod.mf6.drn import Drainage from imod.mf6.evt import Evapotranspiration @@ -18,7 +19,7 @@ from imod.mf6.model import GroundwaterFlowModel from imod.mf6.npf import NodePropertyFlow from imod.mf6.oc import OutputControl -from imod.mf6.out import open_cbc, open_hds, read_cbc_headers, read_grb +from imod.mf6.out import open_cbc, open_hds, open_hds_like, read_cbc_headers, read_grb from imod.mf6.rch import Recharge from imod.mf6.riv import River from imod.mf6.simulation import Modflow6Simulation diff --git a/imod/mf6/bas.py b/imod/mf6/bas.py deleted file mode 100644 index 265a5d44a..000000000 --- a/imod/mf6/bas.py +++ /dev/null @@ -1,215 +0,0 @@ -import jinja2 -import numpy as np -import scipy.ndimage.morphology -import xarray as xr - -from imod import util -from imod.wq.pkgbase import Package - - -class BasicFlow(Package): - """ - The Basic package is used to specify certain data used in all models. - These include: - - 1. the locations of acitve, inactive, and specified head in cells, - 2. the head stored in inactive cells, - 3. the initial head in all cells, and - 4. the top and bottom of the aquifer - - The number of layers (NLAY) is automatically calculated using the IBOUND. - Thickness is calculated using the specified tops en bottoms. - The Basic package input file is required in all models. - - Parameters - ---------- - ibound: xr.DataArray of integers - is the boundary variable. - If IBOUND(J,I,K) < 0, cell J,I,K has a constant head. - If IBOUND(J,I,K) = 0, cell J,I,K is inactive. - If IBOUND(J,I,K) > 0, cell J,I,K is active. - top: float or xr.DataArray of floats - is the top elevation of layer 1. For the common situation in which the - top layer represents a water-table aquifer, it may be reasonable to set - `top` equal to land-surface elevation. - bottom: xr.DataArray of floats - is the bottom elevation of model layers or Quasi-3d confining beds. The - DataArray should at least include the `layer` dimension. - starting_head: float or xr.DataArray of floats - is initial (starting) head—that is, head at the beginning of the - simulation (STRT). starting_head must be specified for all simulations, - including steady-state simulations. One value is read for every model - cell. Usually, these values are read a layer at a time. - inactive_head: float, optional - is the value of head to be assigned to all inactive (no flow) cells - (IBOUND = 0) throughout the simulation (HNOFLO). Because head at - inactive cells is unused in model calculations, this does not affect - model results but serves to identify inactive cells when head is - printed. This value is also used as drawdown at inactive cells if the - drawdown option is used. Even if the user does not anticipate having - inactive cells, a value for inactive_head must be entered. Default - value is 1.0e30. - """ - - _pkg_id = "bas6" - _template = jinja2.Template( - "[bas6]\n" - " {%- for layer, value in ibound.items() %}\n" - " ibound_l{{layer}} = {{value}}\n" - " {%- endfor %}\n" - " hnoflo = {{inactive_head}}\n" - " {%- for layer, value in starting_head.items() %}\n" - " strt_l{{layer}} = {{value}}\n" - " {%- endfor -%}" - ) - - def __init__(self, ibound, top, bottom, starting_head, inactive_head=1.0e30): - self._check_ibound(ibound) - super(__class__, self).__init__() - self["ibound"] = ibound - self["top"] = top - self["bottom"] = bottom - self["starting_head"] = starting_head - self["inactive_head"] = inactive_head - - def _check_ibound(self, ibound): - if not isinstance(ibound, xr.DataArray): - raise TypeError("ibound must be xarray.DataArray") - dims = ibound.dims - if not (dims == ("layer", "y", "x") or dims == ("z", "y", "x")): - raise ValueError( - f'ibound dimensions must be ("layer", "y", "x") or ("z", "y", "x"),' - f" got instead {dims}" - ) - - def _render(self, directory, nlayer, *args, **kwargs): - """ - Renders part of runfile that ends up under [bas] section. - """ - d = {} - for varname in ("ibound", "starting_head"): - d[varname] = self._compose_values_layer(varname, directory, nlayer) - d["inactive_head"] = self["inactive_head"].values - - return self._template.render(d) - - def _compose_top(self, directory): - """ - Composes paths to file, or gets the appropriate scalar value for - a top of model domain. - - Parameters - ---------- - directory : str - """ - da = self["top"] - if "x" in da.coords and "y" in da.coords: - if not len(da.shape) == 2: - raise ValueError("Top should either be 2d or a scalar value") - d = {} - d["name"] = "top" - d["directory"] = directory - d["extension"] = ".idf" - value = self._compose(d) - else: - if not da.shape == (): - raise ValueError("Top should either be 2d or a scalar value") - value = float(da) - return value - - @staticmethod - def _cellsizes(dx): - ncell = dx.size - index_ends = np.argwhere(np.diff(dx) != 0.0) + 1 - index_ends = np.append(index_ends, ncell) - index_starts = np.insert(index_ends[:-1], 0, 0) + 1 - - d = {} - for s, e in zip(index_starts, index_ends): - value = abs(float(dx[s - 1])) - if s == e: - d[f"{s}"] = value - else: - d[f"{s}:{e}"] = value - return d - - def _render_dis(self, directory, nlayer): - """ - Renders part of runfile that ends up under [dis] section. - """ - d = {} - d["top"] = self._compose_top(directory) - d["bottom"] = self._compose_values_layer("bottom", directory, nlayer) - d["nlay"], d["nrow"], d["ncol"] = self["ibound"].shape - # TODO: check dx > 0, dy < 0? - if "dx" not in self or "dy" not in self: # assume equidistant - dx, _, _ = util.coord_reference(self["x"]) - dy, _, _ = util.coord_reference(self["y"]) - else: - dx = self.coords["dx"] - dy = self.coords["dy"] - - if isinstance(dy, (float, int)) or dy.shape in ((), (1,)): - d["dy"] = {"?": abs(float(dy))} - else: - d["dy"] = self._cellsizes(dy) - - if isinstance(dx, (float, int)) or dx.shape in ((), (1,)): - d["dx"] = {"?": abs(float(dx))} - else: - d["dx"] = self._cellsizes(dx) - - # Non-time dependent part of dis - # Can be inferred from ibound - _dis_template = jinja2.Template( - "[dis]\n" - " nlay = {{nlay}}\n" - " nrow = {{nrow}}\n" - " ncol = {{ncol}}\n" - " {%- for row, value in dy.items() %}\n" - " delc_r{{row}} = {{value}}\n" - " {%- endfor %}\n" - " {%- for col, value in dx.items() %}\n" - " delr_c{{col}} = {{value}}\n" - " {%- endfor %}\n" - " top = {{top}}\n" - " {%- for layer, value in bottom.items() %}\n" - " botm_l{{layer}} = {{value}}\n" - " {%- endfor %}\n" - " laycbd_l? = 0" - ) - - return _dis_template.render(d) - - def thickness(self): - """ - Computes layer thickness from top and bottom data. - - Returns - ------- - thickness : xr.DataArray - """ - th = xr.concat( - [ - self["top"] - self["bottom"].sel(layer=1), - -1.0 * self["bottom"].diff("layer"), - ], - dim="layer", - ) - return th - - def _pkgcheck(self, ibound=None): - if (self["top"] < self["bottom"]).any(): - raise ValueError(f"top should be larger than bottom in {self}") - - active_cells = self["ibound"] != 0 - if (active_cells & np.isnan(self["starting_head"])).any(): - raise ValueError( - f"Active cells in ibound may not have a nan value in starting_head in {self}" - ) - - _, nlabels = scipy.ndimage.label(active_cells.values) - if nlabels > 1: - raise ValueError( - f"{nlabels} disconnected model domain detected in the ibound in {self}" - ) diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index 0dbbdbea6..b37e9abf7 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -1,8 +1,13 @@ +from pathlib import Path + import numpy as np +import xarray as xr import imod from imod.mf6.pkgbase import Package, VariableMetaData +from .read_input import read_dis_blockfile + class StructuredDiscretization(Package): """ @@ -60,6 +65,40 @@ def _delrc(self, dx): else: raise ValueError(f"Unhandled type of {dx}") + @classmethod + def open(cls, path: str, simroot: Path) -> "StructuredDiscretization": + content = read_dis_blockfile(simroot / path, simroot) + nlay = content["nlay"] + + griddata = content["griddata"] + xmin = float(content.get("xorigin", 0.0)) + ymin = float(content.get("yorigin", 0.0)) + dx = griddata.pop("delr") + dy = griddata.pop("delc") + x = (xmin + np.cumsum(dx)) - 0.5 * dx + y = ((ymin + np.cumsum(dy)) - 0.5 * dy)[::-1] + + # Top is a 2D array unlike the others + top = xr.DataArray( + data=griddata.pop("top"), + coords={"y": y, "x": x}, + dims=("y", "x"), + ) + + coords = {"layer": np.arange(1, nlay + 1), "y": y, "x": x} + dims = ("layer", "y", "x") + inverse_keywords = {v: k for k, v in cls._keyword_map.items()} + variables = {"top": top} + for key, value in griddata.items(): + invkey = inverse_keywords.get(key, key) + variables[invkey] = xr.DataArray( + value, + coords, + dims, + ) + + return cls(**variables) + def render(self, directory, pkgname, globaltimes, binary): disdirectory = directory / pkgname d = {} diff --git a/imod/mf6/disu.py b/imod/mf6/disu.py new file mode 100644 index 000000000..fe7fc3ac8 --- /dev/null +++ b/imod/mf6/disu.py @@ -0,0 +1,304 @@ +import numba as nb +import numpy as np +import xarray as xr +from scipy import sparse + +import imod +from imod.mf6.pkgbase import Package + +IntArray = np.ndarray + + +@nb.njit(inline="always") +def _number(k, i, j, nrow, ncolumn): + return k * (nrow * ncolumn) + i * ncolumn + j + + +@nb.njit +def _structured_connectivity(idomain: IntArray): + nlayer, nrow, ncolumn = idomain.shape + # Pre-allocate: structured connectivity implies maximum of 8 neighbors + nconnection = idomain.size * 8 + ii = np.empty(nconnection, dtype=np.int32) + jj = np.empty(nconnection, dtype=np.int32) + + connection = 0 + for k in range(nlayer): + for i in range(nrow): + for j in range(ncolumn): + # Skip inactive or pass-through cells + if idomain[k, i, j] <= 0: + continue + + if j < ncolumn - 1: + if idomain[k, i, j + 1] > 0: + ii[connection] = _number(k, i, j, nrow, ncolumn) + jj[connection] = _number(k, i, j + 1, nrow, ncolumn) + connection += 1 + + if i < nrow - 1: + if idomain[k, i + 1, j] > 0: + ii[connection] = _number(k, i, j, nrow, ncolumn) + jj[connection] = _number(k, i + 1, j, nrow, ncolumn) + connection += 1 + + if k < nlayer - 1: + kk = k + while kk < nlayer - 1: + kk += 1 + below = idomain[kk, i, j] + if below > 0: + ii[connection] = _number(k, i, j, nrow, ncolumn) + jj[connection] = _number(kk, i, j, nrow, ncolumn) + connection += 1 + break + elif below == 0: + break + + return ii[:connection], jj[:connection] + + +class LowLevelUnstructuredDiscretization(Package): + """ + Unstructured Discretization (DISU). + + Parameters + ---------- + xorigin: float + yorigin: float + top: xr.DataArray of floats + bot: xr.DataArray of floats + area: xr.DataArray of floats + iac: xr.DataArray of integers + ja: xr.DataArray of integers + ihc: xr.DataArray of integers + cl12: xr.DataArray of floats + hwva: xr.DataArray of floats + angldegx: xr.DataArray of floats + """ + + _pkg_id = "disu" + _grid_data = { + "top": np.float64, + "bot": np.float64, + "area": np.float64, + "iac": np.int32, + "ja": np.int32, + "ihc": np.int32, + "cl12": np.float64, + "hwva": np.float64, + "idomain": np.int32, + "angldegx": np.float64, + } + _keyword_map = {} + _template = Package._initialize_template(_pkg_id) + + def __init__( + self, + xorigin, + yorigin, + top, + bot, + area, + iac, + ja, + ihc, + cl12, + hwva, + angldegx, + idomain=None, + ): + super().__init__(locals()) + self.dataset["xorigin"] = xorigin + self.dataset["yorigin"] = yorigin + self.dataset["top"] = top + self.dataset["bot"] = bot + self.dataset["area"] = area + self.dataset["iac"] = iac + self.dataset["ja"] = ja + self.dataset["ihc"] = ihc + self.dataset["cl12"] = cl12 + self.dataset["hwva"] = hwva + self.dataset["angldegx"] = angldegx + if idomain is not None: + self.dataset["idomain"] = idomain + + def render(self, directory, pkgname, globaltimes, binary): + disdirectory = directory / pkgname + d = {} + d["xorigin"] = float(self.dataset["xorigin"]) + d["yorigin"] = float(self.dataset["yorigin"]) + + # Dimensions + d["nodes"] = self.dataset["top"].size + d["nja"] = int(self.dataset["iac"].sum()) + + # Grid data + for varname in self._grid_data: + if varname in self.dataset: + key = self._keyword_map.get(varname, varname) + d[varname] = self._compose_values( + self.dataset[varname], disdirectory, key, binary=binary + )[1][0] + + return self._template.render(d) + + @staticmethod + def from_dis( + top, + bottom, + idomain, + reduce_nodes=False, + ): + """ + Parameters + ---------- + reduce_nodes: bool, optional. Default: False. + Reduces the node numbering, discards cells when idomain <= 0. + + Returns + ------- + disu: LowLevelUnstructuredDiscretization + cell_ids: ndarray of integers. + Only provided if ``reduce_nodes`` is ``True``. + """ + x = idomain.coords["x"] + y = idomain.coords["y"] + layer = idomain.coords["layer"] + active = idomain.values.ravel() > 0 + + ncolumn = x.size + nrow = y.size + nlayer = layer.size + size = idomain.size + + dx, xmin, _ = imod.util.coord_reference(x) + dy, ymin, _ = imod.util.coord_reference(y) + + # MODFLOW6 expects the ja values to contain the cell number first while + # the row should be otherwise sorted ascending. scipy.spare.csr_matrix + # will sort the values ascending, but would not put the cell number + # first. To ensure this, we use the values as well as i and j; we sort + # on the zeros (thereby ensuring it results as a first value per + # column), but the actual value is the (negative) cell number (in v). + ii, jj = _structured_connectivity(idomain.values) + ii += 1 + jj += 1 + nodes = np.arange(1, size + 1, dtype=np.int32) + if reduce_nodes: + nodes = nodes[active.ravel()] + + zeros = np.zeros_like(nodes) + i = np.concatenate([nodes, ii, jj]) + j = np.concatenate([zeros, jj, ii]) + v = np.concatenate([-nodes, jj, ii]) + csr = sparse.csr_matrix((v, (i, j)), shape=(size + 1, size + 1)) + # The first column can be identified by its negative (node) number. + # This entry does not require data in ihc, cl12, hwva. + is_node = csr.data < 0 + + nnz = csr.getnnz(axis=1) + if reduce_nodes: + # Constructing the CSR matrix will have sorted all the values are + # required by MODFLOW6. However, we're using the original structured + # numbering, which includes inactive cells. + # For MODFLOW6, we use the reduced numbering if reduce_nodes is True, + # excluding all inactive cells. This means getting rid of empty rows + # (iac), generating (via cumsum) new numbers, and extracting them in + # the right order. + iac = nnz[nnz > 0] + ja_index = np.abs(csr.data) - 1 # Get rid of negative values temporarily. + ja = active.cumsum()[ja_index] + else: + # In this case, inactive cells are included as well. They have no + # connections to other cells and form empty rows (0 in iac), but + # are still included. There is no need to update the cell numbers + # in this case. + iac = nnz[1:] # Cell 0 does not exist. + ja = csr.data + + # From CSR back to COO form + # connectivity for every cell: n -> m + n = np.repeat(np.arange(size + 1), nnz) - 1 + m = csr.indices - 1 + # Ignore the values that do not represent n -> m connections + n[is_node] = 0 + m[is_node] = 0 + + # Based on the row and column number differences we can derive the type + # of connection (unless the model is a single row or single column!). + diff = np.abs(n - m) + is_vertical = (diff > 0) & (diff % (nrow * ncolumn) == 0) # one or more layers + is_x = diff == 1 + is_y = diff == ncolumn + is_horizontal = is_x | is_y + + # We need the indexes twice. Store for re-use. + # As the input is structured, we need only look at cell n, not m. + # (n = row, m = column off the connectivity matrix.) + index_x = n[is_x] + index_y = n[is_y] + index_v = n[is_vertical] + + # Create flat arrays for easy indexing. + cellheight = top.values.ravel() - bottom.values.ravel() + dyy, dxx = np.meshgrid( + np.ones(ncolumn) * np.abs(dx), + np.ones(nrow) * np.abs(dy), + indexing="ij", + ) + dyy = np.repeat(dyy, nlayer).ravel() + dxx = np.repeat(dxx, nlayer).ravel() + area = dyy * dxx + + # Allocate connectivity geometry arrays, all size nja. + ihc = is_horizontal.astype(np.int32) + cl12 = np.zeros_like(ihc, dtype=np.float64) + hwva = np.zeros_like(ihc, dtype=np.float64) + angldegx = np.zeros_like(ihc, dtype=np.float64) + # Fill. + cl12[is_x] = 0.5 * dxx[index_x] # cell center to vertical face + cl12[is_y] = 0.5 * dyy[index_y] # cell center to vertical face + cl12[is_vertical] = 0.5 * cellheight[index_v] # cell center to horizontal face + hwva[is_x] = dyy[index_x] # width + hwva[is_y] = dxx[index_y] # width + hwva[is_vertical] = area[index_v] # area + angldegx[is_y] = 90.0 # angle between connection normal and x-axis. + + # Set "node" and "nja" as the dimension in accordance with MODFLOW6. + # Should probably be updated if we could move to UGRID3D... + if reduce_nodes: + # If we reduce nodes, we should only take active cells from top, + # bottom, area. There is no need to include an idomain: all defined + # cells are active. + disu = LowLevelUnstructuredDiscretization( + xorigin=xmin, + yorigin=ymin, + top=xr.DataArray(top.values.ravel()[active], dims=["node"]), + bot=xr.DataArray(bottom.values.ravel()[active], dims=["node"]), + area=xr.DataArray(area[active], dims=["node"]), + iac=xr.DataArray(iac, dims=["node"]), + ja=xr.DataArray(ja, dims=["nja"]), + ihc=xr.DataArray(ihc, dims=["nja"]), + cl12=xr.DataArray(cl12, dims=["nja"]), + hwva=xr.DataArray(hwva, dims=["nja"]), + angldegx=xr.DataArray(angldegx, dims=["nja"]), + ) + cell_ids = np.cumsum(active) - 1 + cell_ids[~active] = -1 + return disu, cell_ids + else: + return LowLevelUnstructuredDiscretization( + xorigin=xmin, + yorigin=ymin, + top=xr.DataArray(top.values.ravel(), dims=["node"]), + bot=xr.DataArray(bottom.values.ravel(), dims=["node"]), + area=xr.DataArray(area, dims=["node"]), + iac=xr.DataArray(iac, dims=["node"]), + ja=xr.DataArray(ja, dims=["nja"]), + ihc=xr.DataArray(ihc, dims=["nja"]), + cl12=xr.DataArray(cl12, dims=["nja"]), + hwva=xr.DataArray(hwva, dims=["nja"]), + angldegx=xr.DataArray(angldegx, dims=["nja"]), + idomain=xr.DataArray(active.astype(np.int32), dims=["node"]), + ) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 2db089847..2c814bd03 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -1,12 +1,16 @@ import collections import pathlib +from typing import Dict, Type, Union import cftime import jinja2 import numpy as np +from imod import mf6 from imod.mf6 import qgs_util +from .read_input import read_gwf_namefile + class Model(collections.UserDict): def __setitem__(self, key, value): @@ -25,6 +29,40 @@ class GroundwaterFlowModel(Model): _pkg_id = "model" + @staticmethod + def _PACKAGE_CLASSES(distype: str) -> Dict[str, Type]: + # DIS dependent packages do not have a completely topology description + # to a Python equivalent. These are packages such as: + # + # * Well + # * Multi-Aquifer Well + # * Horizontal Flow Barrier + # * Stream Flow Routing + # + # Instead, they are returned in a lower level form, directly related to + # the MODFLOW6 input. + dis_dependent = { + "dis": (mf6.WellDisStructured,), + "disv": (mf6.WellDisVertices,), + } + + # mf6.OutputControl is skipped + # mf6.StorageCoefficient handled by SpecificStorage + packages = dis_dependent[distype] + ( + mf6.ConstantHead, + mf6.StructuredDiscretization, + mf6.VerticesDiscretization, + mf6.Drainage, + mf6.Evapotranspiration, + mf6.GeneralHeadBoundary, + mf6.InitialConditions, + mf6.NodePropertyFlow, + mf6.Recharge, + mf6.River, + mf6.SpecificStorage, + ) + return {package._pkg_id: package for package in packages} + def _initialize_template(self): loader = jinja2.PackageLoader("imod", "templates/mf6") env = jinja2.Environment(loader=loader, keep_trailing_newline=True) @@ -109,6 +147,77 @@ def render(self, modelname): d["packages"] = packages return self._template.render(d) + @classmethod + def open( + cls, + path: Union[pathlib.Path, str], + simroot: pathlib.Path, + globaltimes: np.ndarray, + ): + content = read_gwf_namefile(simroot / path) + model = cls( + newton=content.get("newton", False), + under_relaxation=content.get("under_relaxation", False), + ) + + # Search for the DIS/DISV/DISU package first. This provides us with + # the coordinates and dimensions to instantiate the other packages. + dis_packages = [ + tup for tup in content["packages"] if tup[0] in ("dis6", "disv6", "disu6") + ] + if len(dis_packages) != 1: + raise ValueError(f"Exactly one DIS/DISV/DISU package required in {path}") + + disftype, disfname, dispname = dis_packages[0] + diskey = disftype[:-1] + + classes = cls._PACKAGE_CLASSES(diskey) + package = classes[diskey] + dis_package = package.open( + disfname, + simroot, + ) + model[dispname] = dis_package + shape = dis_package["idomain"].shape + coords = dict(dis_package["idomain"].coords) + dims = dis_package["idomain"].dims + + # Non-dis packages + packages = [ + tup + for tup in content["packages"] + if tup[0] not in ("dis6", "disv6", "disu6") + ] + # Make sure names are unique by counting first, and appending numbers + # if they are not unique. + occurence = collections.Counter([tup[2] for tup in packages]) + counter = collections.defaultdict(int) + for ftype, fname, pname in packages: + key = ftype[:-1] + package = classes.get(key) + if package is None: + continue + if occurence[pname] > 1: + pkgname = f"{pname}_{counter[pname]}" + counter[pname] += 1 + else: + pkgname = pname + + try: + # Create the package and add it to the model. + model[pkgname] = package.open( + path=fname, + simroot=simroot, + shape=shape, + coords=coords, + dims=dims, + globaltimes=globaltimes, + ) + except Exception as e: + raise type(e)(f"{e}\n Error reading {fname}") from e + + return model + def write(self, directory, modelname, globaltimes, binary=True): """ Write model namefile diff --git a/imod/mf6/out/__init__.py b/imod/mf6/out/__init__.py index 1ba37bc5f..3807e58d7 100644 --- a/imod/mf6/out/__init__.py +++ b/imod/mf6/out/__init__.py @@ -151,10 +151,12 @@ def open_hds_like( # TODO: check shape with hds metadata. if isinstance(like, xr.DataArray): d = dis.grid_info(like) + d["name"] = "head" return dis.open_hds(path, d, dry_nan) elif isinstance(like, xu.UgridDataArray): d = disv.grid_info(like) + d["name"] = "head" return disv.open_hds(path, d, dry_nan) else: diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py index dcbd051eb..2d1cf32a1 100644 --- a/imod/mf6/out/dis.py +++ b/imod/mf6/out/dis.py @@ -74,11 +74,11 @@ def read_times( """ times = np.empty(ntime, dtype=np.float64) - # Compute how much to skip to the next timestamp - start_of_header = 16 - rest_of_header = 28 + # Compute how much to skip to the next timestamp. + start_of_header = 16 # KSTP(4), KPER(4), PERTIM(8) + rest_of_header = 28 # TEXT(16), NCOL(4), NROW(4), ILAY(4) data_single_layer = nrow * ncol * 8 - header = 52 + header = 52 # start_of_header + TOTIM(8) + rest_of_header nskip = ( rest_of_header + data_single_layer @@ -117,6 +117,7 @@ def read_hds_timestep( def open_hds(path: FilePath, d: Dict[str, Any], dry_nan: bool) -> xr.DataArray: nlayer, nrow, ncol = d["nlayer"], d["nrow"], d["ncol"] filesize = os.path.getsize(path) + # 52 is header size; 8 is size of double. ntime = filesize // (nlayer * (52 + (nrow * ncol * 8))) times = read_times(path, ntime, nlayer, nrow, ncol) coords = d["coords"] diff --git a/imod/mf6/out/disu.py b/imod/mf6/out/disu.py index f7c0372d9..f884d8157 100644 --- a/imod/mf6/out/disu.py +++ b/imod/mf6/out/disu.py @@ -1,12 +1,14 @@ +import os import struct from typing import Any, BinaryIO, Dict, List +import dask import numpy as np import xarray as xr import xugrid as xu from . import cbc -from .common import FilePath, _grb_text +from .common import FilePath, _grb_text, _to_nan def read_grb(f: BinaryIO, ntxt: int, lentxt: int) -> Dict[str, Any]: @@ -58,17 +60,48 @@ def read_grb(f: BinaryIO, ntxt: int, lentxt: int) -> Dict[str, Any]: def read_times(path: FilePath, ntime: int, ncell: int): - raise NotImplementedError - - -def read_hds_timestep( - path: FilePath, nlayer: int, ncell_per_layer: int, dry_nan: bool, pos: int -): - raise NotImplementedError + """ + Reads all total simulation times. + """ + times = np.empty(ntime, dtype=np.float64) + + # Compute how much to skip to the next timestamp + start_of_header = 16 # KSTP(4), KPER(4), PERTIM(8) + rest_of_header = 28 # TEXT(16), NCOL(4), NROW(4), ILAY(4) + data = ncell * 8 + nskip = rest_of_header + data + start_of_header + with open(path, "rb") as f: + f.seek(start_of_header) + for i in range(ntime): + times[i] = struct.unpack("d", f.read(8))[0] + f.seek(nskip, 1) + return times + + +def read_hds_timestep(path: FilePath, ncell: int, dry_nan: bool, pos: int): + with open(path, "rb") as f: + f.seek(pos + 52) + a1d = np.fromfile(f, np.float64, ncell) + return _to_nan(a1d, dry_nan) def open_hds(path: FilePath, d: Dict[str, Any], dry_nan: bool) -> xr.DataArray: - raise NotImplementedError + ncell = d["ncell"] + filesize = os.path.getsize(path) + ntime = filesize // (52 + ncell * 8) + times = read_times(path, ntime, ncell) + coords = d["coords"] + coords["time"] = times + + dask_list = [] + for i in range(ntime): + pos = i * (52 + ncell * 8) + a = dask.delayed(read_hds_timestep)(path, ncell, dry_nan, pos) + x = dask.array.from_delayed(a, shape=(ncell,), dtype=np.float64) + dask_list.append(x) + + daskarr = dask.array.stack(dask_list, axis=0) + return xr.DataArray(daskarr, coords, ("time", "node"), name=d["name"]) def open_imeth1_budgets( diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py index 6026fefe4..8119fe422 100644 --- a/imod/mf6/out/disv.py +++ b/imod/mf6/out/disv.py @@ -130,7 +130,7 @@ def read_hds_timestep( f.seek(pos) a1d = np.empty(nlayer * ncells_per_layer, dtype=np.float64) for k in range(nlayer): - f.seek(52, 1) # skip kstp, kper, pertime + f.seek(52, 1) # skip header a1d[k * ncells_per_layer : (k + 1) * ncells_per_layer] = np.fromfile( f, np.float64, ncells_per_layer ) diff --git a/imod/mf6/pkgbase.py b/imod/mf6/pkgbase.py index 9e5bbcfa6..61b76d653 100644 --- a/imod/mf6/pkgbase.py +++ b/imod/mf6/pkgbase.py @@ -1,12 +1,16 @@ import abc +import inspect import pathlib from dataclasses import dataclass +from typing import Any, Dict import jinja2 import numpy as np import xarray as xr import xugrid as xu +from .read_input import read_blockfile, read_boundary_blockfile, shape_to_max_rows + def dis_recarr(arrdict, layer, notnull): # Define the numpy structured array dtype @@ -47,6 +51,22 @@ def disv_recarr(arrdict, layer, notnull): return recarr +def disu_recarr(arrdict, node, notnull): + index_spec = [("node", np.int32)] + field_spec = [(key, np.float64) for key in arrdict] + sparse_dtype = np.dtype(index_spec + field_spec) + # Initialize the structured array + nrow = notnull.sum() + recarr = np.empty(nrow, dtype=sparse_dtype) + # Argwhere returns an index_array with dims (N, a.ndims) + index = np.argwhere(notnull)[:, 0] + if node is None: + recarr["node"] = index + 1 + else: + recarr["node"] = node[index] + 1 + return recarr + + @dataclass class VariableMetaData: """ @@ -165,13 +185,18 @@ def sel(self): def _valid(self, value): """ - Filters values that are None, False, or a numpy.bool_ False. + Filters values that are None, False, np.nan, or a numpy.bool_ False. Needs to be this specific, since 0.0 and 0 are valid values, but are equal to a boolean False. + + Intercepting single NaNs is practical because xarray replaces None by + NaN when writing. """ # Test singletons if value is False or value is None: return False + elif isinstance(value, float) and np.isnan(value): + return False # Test numpy bool (not singleton) elif isinstance(value, np.bool_) and not value: return False @@ -220,7 +245,12 @@ def to_sparse(self, arrdict, layer): notnull = ~np.isnan(data) if isinstance(self.dataset, xr.Dataset): - recarr = dis_recarr(arrdict, layer, notnull) + # TODO + if ("node" in self.dataset.dims) or ("nja" in self.dataset.dims): + recarr = disu_recarr(arrdict, layer, notnull) + else: + recarr = dis_recarr(arrdict, layer, notnull) + elif isinstance(self.dataset, xu.UgridDataset): recarr = disv_recarr(arrdict, layer, notnull) else: @@ -301,7 +331,12 @@ def render(self, directory, pkgname, globaltimes, binary): @staticmethod def _is_xy_data(obj): if isinstance(obj, (xr.DataArray, xr.Dataset)): - xy = "x" in obj.dims and "y" in obj.dims + # "nja" is not really xy_data: arguably the method name should be + # changed. This method is exclusively used whether the contents of + # the object should be written to an external file. + xy = ("x" in obj.dims and "y" in obj.dims) or ( + "node" in obj.dims or "nja" in obj.dims + ) elif isinstance(obj, (xu.UgridDataArray, xu.UgridDataset)): xy = obj.ugrid.grid.face_dimension in obj.dims else: @@ -355,14 +390,60 @@ def write(self, directory, pkgname, globaltimes, binary): pkgdirectory.mkdir(exist_ok=True, parents=True) for varname, dtype in self._grid_data.items(): key = self._keyword_map.get(varname, varname) - da = self.dataset[varname] - if self._is_xy_data(da): - if binary: - path = pkgdirectory / f"{key}.bin" - self.write_binary_griddata(path, da, dtype) - else: - path = pkgdirectory / f"{key}.dat" - self.write_text_griddata(path, da, dtype) + if varname in self.dataset: + da = self.dataset[varname] + if self._is_xy_data(da): + if binary: + path = pkgdirectory / f"{key}.bin" + self.write_binary_griddata(path, da, dtype) + else: + path = pkgdirectory / f"{key}.dat" + self.write_text_griddata(path, da, dtype) + + @classmethod + def filter_and_rename(cls, content: Dict[str, Any]) -> Dict[str, Any]: + """ + Filter content for keywords that are required by the ``__init__`` function + of a class. + + Rename keywords to the expected form for the class, as given by + ``_keyword_map``. + + Arguments + --------- + init: Callable + __init__ of the + """ + inverse_keywords = {v: k for k, v in cls._keyword_map.items()} + init_kwargs = inspect.getfullargspec(cls.__init__).args + + filtered = {} + for key, value in content.items(): + newkey = inverse_keywords.get(key, key) + if newkey in init_kwargs: + filtered[newkey] = value + + return filtered + + @classmethod + def open(cls, path, simroot, shape, coords, dims, globaltimes): + sections = { + cls._keyword_map.get(field, field): (dtype, shape_to_max_rows) + for field, dtype in cls._grid_data.items() + } + content = read_blockfile( + simroot / path, + simroot=simroot, + sections=sections, + shape=shape, + ) + + griddata = content.pop("griddata") + for field, data in griddata.items(): + content[field] = xr.DataArray(data, coords, dims) + + filtered_content = cls.filter_and_rename(content) + return cls(**filtered_content) def _netcdf_path(self, directory, pkgname): """create path for netcdf, this function is also used to create paths to use inside the qgis projectfiles""" @@ -413,6 +494,58 @@ def write_netcdf(self, directory, pkgname, aggregate_layers=False): spatial_ds.to_netcdf(path) return has_dims + def _dis_to_disu(self, cell_ids=None): + structured = self.dataset + if cell_ids is not None and not np.issubdtype(cell_ids.dtype, np.integer): + raise TypeError(f"cell_ids should be integer, received: {cell_ids.dtype}") + if "layer" not in structured.coords: + raise ValueError("layer coordinate is required") + if "layer" not in structured.dims: + structured = self.dataset.expand_dims("layer") + + # Stack will automatically broadcast to layer + dataset = structured.stack(node=("layer", "y", "x")) + layers = structured["layer"].values + ncell_per_layer = structured["y"].size * structured["x"].size + offset = (layers - 1) * ncell_per_layer + index = np.add.outer(offset, np.arange(ncell_per_layer)).ravel() + + if cell_ids is not None: + # node has the shape of index, i.e. one value for every cell in DIS + # form, including no data padding, which are values of -1. + node = cell_ids[index] + # Create a subselection without the inactive cells. + active = [node != -1] + dataset = dataset.isel(node=index[active]) + dataset = dataset.assign_coords(node=node[active]) + else: + dataset = dataset.assign_coords(node=index) + + return self.__class__(**dataset) + + def to_disu(self, cell_ids=None): + + """ + Parameters + ---------- + cell_ids: np.ndarray of integers, optional. + DISU cell IDs. Should contain values of -1 for inactive cells. + + Returns + ------- + disu_package: Any + Package in DISU form. + """ + if isinstance(self.dataset, xr.Dataset): + return self._dis_to_disu(cell_ids) + elif isinstance(self.dataset.xu.UgridDataset): + raise NotImplementedError + else: + raise TypeError( + "Expected xarray.Dataset or xugrid.UgridDataset. " + f"Found: {type(self.dataset)}" + ) + def _check_types(self): for varname, metadata in self._metadata_dict.items(): expected_dtype = metadata.dtype @@ -467,10 +600,19 @@ def write_datafile(self, outpath, ds, binary): """ Writes a modflow6 binary data file """ - layer = ds["layer"].values + if "layer" in ds: + layer = ds["layer"].values + # TODO: change layer argument name to cell_id or something? + # Also forward the node number in case of DISU input. + elif "node" in ds: + layer = ds["node"].values + else: + layer = None + arrdict = self._ds_to_arrdict(ds) sparse_data = self.to_sparse(arrdict, layer) outpath.parent.mkdir(exist_ok=True, parents=True) + if binary: self._write_binaryfile(outpath, sparse_data) else: @@ -555,6 +697,27 @@ def write(self, directory, pkgname, globaltimes, binary): binary=binary, ) + @classmethod + def open(cls, path, simroot, shape, coords, dims, globaltimes): + content = read_boundary_blockfile( + simroot / path, + simroot, + fields=cls._period_data, + shape=shape, + ) + + period_index = content.pop("period_index") + period_data = content.pop("period_data") + coords = coords.copy() + coords["time"] = globaltimes[period_index] + dims = ("time",) + dims + + for field, data in period_data.items(): + content[field] = xr.DataArray(data, coords, dims) + + filtered_content = cls.filter_and_rename(content) + return cls(**filtered_content) + class AdvancedBoundaryCondition(BoundaryCondition, abc.ABC): """ diff --git a/imod/mf6/read_input/__init__.py b/imod/mf6/read_input/__init__.py new file mode 100644 index 000000000..b7bcb5116 --- /dev/null +++ b/imod/mf6/read_input/__init__.py @@ -0,0 +1,478 @@ +""" +Utilities for reading MODFLOW6 input files. +""" +import warnings +from collections import defaultdict +from pathlib import Path +from typing import IO, Any, Callable, Dict, List, Tuple, Union + +import dask.array +import numpy as np + +from .common import ( + advance_to_header, + advance_to_period, + parse_dimension, + parse_option, + read_iterable_block, + read_key_value_block, + split_line, + to_float, +) +from .grid_data import read_griddata, shape_to_max_rows +from .list_input import read_listinput + + +def parse_model(stripped: str, fname: str) -> Tuple[str, str, str]: + """Parse model entry in the simulation name file.""" + separated = split_line(stripped) + nwords = len(separated) + if nwords == 3: + return separated + else: + raise ValueError( + "ftype, fname and pname expected. Received instead: " + f"{','.join(separated)} in file {fname}" + ) + + +def parse_exchange(stripped: str, fname: str) -> Tuple[str, str, str, str]: + """Parse exchange entry in the simulation name file.""" + separated = split_line(stripped) + nwords = len(separated) + if nwords == 4: + return separated + else: + raise ValueError( + "exgtype, exgfile, exgmnamea, exgmnameb expected. Received instead: " + f"{','.join(separated)} in file {fname}" + ) + + +def parse_solutiongroup(stripped: str, fname: str) -> Tuple[str, str]: + """Parse solution group entry iyn the simulation name file.""" + separated = split_line(stripped) + if "mxiter" in stripped: + return separated + + nwords = len(separated) + if nwords > 2: + return separated[0], separated[1], separated[2:] + else: + raise ValueError( + "Expected at least three entries: slntype, slnfname, and one model name. " + f"Received instead: {','.join(separated)} in file {fname}" + ) + + +def parse_package(stripped: str, fname: str) -> Tuple[str, str, str]: + """Parse package entry in model name file.""" + separated = split_line(stripped) + nwords = len(separated) + if nwords == 2: + ftype, fname = separated + pname = ftype[:-1] # split the number, e.g. riv6 -> riv + elif nwords == 3: + ftype, fname, pname = separated + else: + raise ValueError( + "Expected ftype, fname, and optionally pname. " + f"Received instead: {','.join(separated)} in file {fname}" + ) + return ftype, fname, pname + + +def parse_tdis_perioddata(stripped: str, fname: str) -> Tuple[float, int, float]: + """Parse a single period data entry in the time discretization file.""" + separated = split_line(stripped) + nwords = len(separated) + if nwords >= 3: + return to_float(separated[0]), int(separated[1]), to_float(separated[2]) + else: + raise ValueError( + "perlen, nstp, tsmult expected. Received instead: " + f"{','.join(separated)} in file {fname}" + ) + + +def read_tdis(path: Path) -> Dict[str, Any]: + """Read and parse the content of the time discretization file.""" + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "dimensions") + dimensions = read_key_value_block(f, parse_dimension) + advance_to_header(f, "perioddata") + content["perioddata"] = read_iterable_block(f, parse_tdis_perioddata) + return {**content, **dimensions} + + +def read_dis_blockfile(path: Path, simroot: Path) -> Dict[str, Any]: + """Read and parse structured discretization file.""" + ROW = 1 + COLUMN = 2 + + def delr_max_rows(shape, layered) -> Tuple[int, Tuple[int]]: + if layered: + raise ValueError(f"DELR section in {path} is LAYERED") + return 1, (shape[COLUMN],) + + def delc_max_rows(shape, layered) -> Tuple[int, Tuple[int]]: + if layered: + raise ValueError(f"DELC section in {path} is LAYERED") + return 1, (shape[ROW],) + + def top_max_rows(shape, layered) -> Tuple[int, Tuple[int]]: + if layered: + raise ValueError(f"TOP section in {path} is LAYERED") + _, nrow, ncolumn = shape + return nrow, (nrow, ncolumn) + + sections = { + "delr": (np.float64, delr_max_rows), + "delc": (np.float64, delc_max_rows), + "top": (np.float64, top_max_rows), + "botm": (np.float64, shape_to_max_rows), + "idomain": (np.int32, shape_to_max_rows), + } + + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "dimensions") + dimensions = read_key_value_block(f, parse_dimension) + shape = (dimensions["nlay"], dimensions["nrow"], dimensions["ncol"]) + advance_to_header(f, "griddata") + content["griddata"] = read_griddata(f, simroot, sections, shape) + + return {**content, **dimensions} + + +def tdis_time(tdis: Dict[str, Any]) -> np.ndarray: + """ + Use start_date, time_units, and period duration to create datetime + timestaps for the stress periods. + """ + # https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units + TIME_UNITS = { + "unknown": None, + "seconds": "s", + "minutes": "m", + "hours": "h", + "days": "D", + "years": "Y", + } + unit = TIME_UNITS.get(tdis["time_units"]) + + start = None + if "start_date_time" in tdis: + try: + start = np.datetime64(tdis["start_date_time"]) + except ValueError: + pass + + cumulative_length = np.cumsum([entry[0] for entry in tdis["perioddata"]]) + if unit is not None and start is not None: + timedelta = np.timedelta64(cumulative_length, unit) + times = np.full(timedelta.size, start) + times[1:] += timedelta[:-1] + else: + possibilities = ", ".join(list(TIME_UNITS.keys())[1:]) + warnings.warn( + "Cannot convert time coordinate to datetime. " + "Falling back to (unitless) floating point time coordinate. \n" + f"time_units should be one of: {possibilities}.\n" + "start_date_time should be according to ISO 8601." + ) + times = np.full(cumulative_length.size, 0.0) + times[1:] += cumulative_length[:-1] + + return times + + +def read_solver(path: Path) -> Dict[str, str]: + """Read and parse content of solver config file (IMS).""" + with open(path, "r") as f: + advance_to_header(f, "options") + options = read_key_value_block(f, parse_option) + advance_to_header(f, "nonlinear") + nonlinear = read_key_value_block(f, parse_option) + advance_to_header(f, "linear") + linear = read_key_value_block(f, parse_option) + return {**options, **nonlinear, **linear} + + +def read_sim_namefile(path: Path) -> Dict[str, str]: + """Read and parse content of simulation name file.""" + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "timing") + timing = read_key_value_block(f, parse_option) + advance_to_header(f, "models") + content["models"] = read_iterable_block(f, parse_model) + advance_to_header(f, "exchanges") + content["exchanges"] = read_iterable_block(f, parse_exchange) + advance_to_header(f, "solutiongroup 1") + content["solutiongroup 1"] = read_iterable_block(f, parse_solutiongroup) + return {**content, **timing} + + +def read_gwf_namefile(path: Path) -> Dict[str, Any]: + """Read and parse content of groundwater flow name file.""" + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "packages") + content["packages"] = read_iterable_block(f, parse_package) + return content + + +def read_blockfile( + path: Path, + simroot: Path, + sections: Dict[str, Tuple[type, Callable]], + shape: Tuple[int], +) -> Dict[str, Any]: + """ + Read blockfile of a "standard" MODFLOW6 package: NPF, IC, etc. + + External files are lazily read using dask, constants are lazily allocated, + and internal values are eagerly read and then converted to dask arrays for + consistency. + + Parameters + ---------- + path: Path + simroot: Path + Root path of the simulation, used for reading external files. + sections: Dict[str, Tuple[type, Callable]] + Contains for every array entry its type, and a function to compute from + shape the number of rows to read in case of internal values. + shape: Tuple[int] + DIS: 3-sized, DISV: 2-sized, DISU: 1-sized. + + Returns + ------- + content: Dict[str, Any] + Content of the block file. Grid data arrays are stored as dask arrays. + """ + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "griddata") + content["griddata"] = read_griddata(f, simroot, sections, shape) + + return content + + +def read_package_periods( + f: IO[str], + simroot: Path, + dtype: Union[type, np.dtype], + index_columns: Tuple[str], + fields: Tuple[str], + max_rows: int, + shape: Tuple[int], + sparse_to_dense: bool = True, +) -> Tuple[List[int], Dict[str, dask.array.Array]]: + """ + Read blockfile periods section of a "standard" MODFLOW6 boundary + conditions: RIV, DRN, GHB, etc. + + External files are lazily read using dask and internal values are eagerly + read and then converted to dask arrays for consistency. + + Parameters + ---------- + f: IO[str] + File handle. + simroot: + Root path of the simulation, used for reading external files. + dtype: Union[type, np.dtype] + Generally a numpy structured dtype. + index_columns: Tuple[str] + The index columns (np.int32) of the dtype. + fields: Tuple[str] + The fields (generally np.float64) of the dtype. + max_rows: int + Number of rows to read in case of internal values. + shape: Tuple[int] + DIS: 3-sized, DISV: 2-sized, DISU: 1-sized. + + Returns + ------- + period_index: np.ndarray of integers + variable_values: Dict[str, dask.array.Array] + """ + periods = [] + dask_lists = defaultdict(list) + key = advance_to_period(f) + + if sparse_to_dense: + variables = fields + else: + variables = index_columns + fields + + while key is not None: + # Read the recarrays, already separated into dense arrays. + variable_values = read_listinput( + f=f, + simroot=simroot, + dtype=dtype, + fields=fields, + index_columns=index_columns, + max_rows=max_rows, + shape=shape, + sparse_to_dense=sparse_to_dense, + ) + + # Group them by field (e.g. cond, head, etc.) + for var, values in zip(variables, variable_values): + dask_lists[var].append(values) + + # Store number and advance to next period + periods.append(key) + key = advance_to_period(f) + + # Create a dictionary of arrays + variable_values = { + field: dask.array.stack(dask_list, axis=0) + for field, dask_list in dask_lists.items() + } + return np.array(periods) - 1, variable_values + + +def read_boundary_blockfile( + path: Path, + simroot: Path, + fields: Tuple[str], + shape: Tuple[int], + sparse_to_dense: bool = True, +) -> Dict[str, Any]: + """ + Read blockfile of a "standard" MODFLOW6 boundary condition package: RIV, + DRN, GHB, etc. + + External files are lazily read using dask and internal values are eagerly + read and then converted to dask arrays for consistency. + + Parameters + ---------- + path: Path + simroot: Path + Root path of the simulation, used for reading external files. + fields: Tuple[str] + The fields (generally np.float64) of the dtype. + shape: Tuple[int] + DIS: 3-sized, DISV: 2-sized, DISU: 1-sized. + sparse_to_dense: bool, default: True + Whether to convert "sparse" COO (cell id) data into "dense" grid form + (with implicit location). False for packages such as Well, which are + not usually in grid form. + + Returns + ------- + content: Dict[str, Any] + """ + ndim = len(shape) + index_columns = { + 1: ("node",), + 2: ("layer", "cell2d"), + 3: ("layer", "row", "column"), + }.get(ndim) + if index_columns is None: + raise ValueError(f"Length of dimensions should be 1, 2, or 3. Received: {ndim}") + dtype_fields = [(column, np.int32) for column in index_columns] + [ + (field, np.float64) for field in fields + ] + + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + + # Create the dtype from the required variables, and with potential + # auxiliary variables. + auxiliary = content.pop("auxiliary", None) + if auxiliary: + # Make sure it's iterable in case of a single value + if isinstance(auxiliary, str): + auxiliary = (auxiliary,) + for aux in auxiliary: + dtype_fields.append((aux, np.float64)) + + boundnames = content.pop("boundnames", False) + if boundnames: + dtype_fields.append(("boundnames", str)) + + # Create np.dtype, make fields and columns immutable. + index_columns = tuple(index_columns) + fields = tuple(fields) + dtype = np.dtype(dtype_fields) + + advance_to_header(f, "dimensions") + dimensions = read_key_value_block(f, parse_dimension) + content["period_index"], content["period_data"] = read_package_periods( + f=f, + simroot=simroot, + dtype=dtype, + index_columns=index_columns, + fields=fields, + max_rows=dimensions["maxbound"], + shape=shape, + sparse_to_dense=sparse_to_dense, + ) + + return content + + +def read_sto_blockfile( + path: Path, + simroot: Path, + sections: Dict[str, Tuple[type, Callable]], + shape: Tuple[int], +) -> Dict[str, Any]: + """ + Read blockfile of MODFLOW6 Storage package. + + Parameters + ---------- + path: Path + simroot: Path + Root path of the simulation, used for reading external files. + sections: Dict[str, Tuple[type, Callable]] + Contains for every array entry its type, and a function to compute from + shape the number of rows to read in case of internal values. + shape: Tuple[int] + DIS: 3-sized, DISV: 2-sized, DISU: 1-sized. + + Returns + ------- + content: Dict[str, Any] + Content of the block file. Grid data arrays are stored as dask arrays. + """ + with open(path, "r") as f: + advance_to_header(f, "options") + content = read_key_value_block(f, parse_option) + advance_to_header(f, "griddata") + content["griddata"] = read_griddata(f, simroot, sections, shape) + + periods = {} + key = advance_to_period(f) + while key is not None: + entry = read_key_value_block(f, parse_option) + value = next(iter(entry.keys())) + if value == "steady-state": + value = False + elif value == "transient": + value = True + else: + raise ValueError( + f'Expected "steady-state" or "transient". Received: {value}' + ) + periods[key] = value + key = advance_to_period(f) + + content["periods"] = periods + + return content diff --git a/imod/mf6/read_input/common.py b/imod/mf6/read_input/common.py new file mode 100644 index 000000000..5bb50c6c1 --- /dev/null +++ b/imod/mf6/read_input/common.py @@ -0,0 +1,222 @@ +""" +Commonly shared utilities for reading MODFLOW6 input files. +""" +from pathlib import Path +from typing import IO, Any, Callable, Dict, List, Tuple + +import numpy as np + + +def end_of_file(line: str) -> bool: + return line == "" + + +def strip_line(line: str) -> str: + # remove possible comments + before, _, _ = line.partition("#") + return before.strip().lower() + + +def flatten(iterable: Any) -> List[Any]: + out = [] + for part in iterable: + out.extend(part) + return out + + +def split_line(line: str) -> List[str]: + # Maybe check: https://stackoverflow.com/questions/36165050/python-equivalent-of-fortran-list-directed-input + # Split on comma and whitespace, like a FORTRAN read would do. + flat = flatten([part.split(",") for part in line.split()]) + return [part for part in flat if part != ""] + + +def to_float(string: str) -> float: + # Fortran float may contain d (e.g. 1.0d0), Python only accepts e. + string = string.replace("d", "e") + # Fortran may specify exponents without a letter, e.g. 1.0+1 for 1.0e+1 + if "e" not in string: + string = string.replace("+", "e+").replace("-", "e-") + return float(string) + + +def find_entry(line: str, entry: str, cast: Callable) -> str: + if entry not in line: + return None + else: + _, after = line.split(entry) + return cast(after.split()[0]) + + +def read_internal(f: IO[str], dtype: type, max_rows: int) -> np.ndarray: + return np.loadtxt( + fname=f, + dtype=dtype, + max_rows=max_rows, + ) + + +def read_external_binaryfile(path: Path, dtype: type, max_rows: int) -> np.ndarray: + return np.fromfile( + file=path, + dtype=dtype, + count=max_rows, + sep="", + ) + + +def read_fortran_deflated_text_array( + path: Path, dtype: type, max_rows: int +) -> np.ndarray: + """ + The Fortran read intrinsic is capable of parsing arrays in many forms. + One of those is: + + 1.0 + 2*2.0 + 3.0 + + Which should be interpreted as: [1.0, 2.0, 2.0, 3.0] + + This function attempts this. + """ + out = np.empty(max_rows, dtype) + with open(path, "r") as f: + lines = [line.strip() for line in f] + + iterable_lines = iter(lines) + start = 0 + while start < max_rows: + + line = next(iterable_lines) + if "*" in line: + n, value = line.split("*") + n = int(n) + value = dtype(value) + else: + n = 1 + value = dtype(line) + + end = start + n + out[start:end] = value + start = end + + return out + + +def read_external_textfile(path: Path, dtype: type, max_rows: int) -> np.ndarray: + try: + return np.loadtxt( + fname=path, + dtype=dtype, + max_rows=max_rows, + ) + except ValueError as e: + if str(e).startswith("could not convert string to float"): + return read_fortran_deflated_text_array(path, dtype, max_rows) + else: + raise e + + +def advance_to_header(f: IO[str], section) -> None: + line = None + start = f"begin {section}" + # Empty line is end-of-file + while not end_of_file(line): + line = f.readline() + stripped = strip_line(line) + # Return if start has been found + if stripped == start: + return + # Continue if line is empty + elif stripped == "": + continue + # Raise if the line has (unexpected) content + else: + break + # Also raise if no further content is found + raise ValueError(f'"{start}" is not present in file {f.name}') + + +def parse_option(stripped: str, fname: str) -> Tuple[str, Any]: + separated = stripped.split() + nwords = len(separated) + if nwords == 0: + raise ValueError(f"Cannot parse option in {fname}") + elif nwords == 1: + key = separated[0] + value = True + elif nwords == 2: + key, value = separated + else: + key = separated[0] + value = separated[1:] + return key, value + + +def read_key_value_block(f: IO[str], parse: Callable) -> Dict[str, str]: + fname = f.name + content = {} + line = None + while not end_of_file(line): + line = f.readline() + stripped = strip_line(line) + # Return if end has been found + if stripped[:3] == "end": + return content + # Continue in case of an empty line + elif stripped == "": + continue + # Valid entry + else: + key, value = parse(stripped, fname) + content[key] = value + + # Also raise if no further content is found + raise ValueError(f'"end" of block is not present in file {fname}') + + +def read_iterable_block(f: IO[str], parse: Callable) -> List[Any]: + fname = f.name + content = [] + line = None + while not end_of_file(line): + line = f.readline() + stripped = strip_line(line) + # Return if end has been found + if stripped[:3] == "end": + return content + # Continue in case of an empty line + elif stripped == "": + continue + # Valid entry + else: + content.append(parse(stripped, fname)) + + # Also raise if no further content is found + raise ValueError(f'"end" of block is not present in file {fname}') + + +def parse_dimension(stripped: str, fname: str) -> Tuple[str, int]: + key, value = parse_option(stripped, fname) + return key, int(value) + + +def advance_to_period(f: IO[str]) -> int: + line = None + start = "begin period" + # Empty line is end-of-file + while not end_of_file(line): + line = f.readline() + stripped = strip_line(line) + # Return if start has been found + if stripped[:12] == start: + return int(stripped.split()[2]) + # Continue if line is empty + elif stripped == "": + continue + # Raise if the line has (unexpected) content + else: + break + else: + return None diff --git a/imod/mf6/read_input/grid_data.py b/imod/mf6/read_input/grid_data.py new file mode 100644 index 000000000..8532f42d0 --- /dev/null +++ b/imod/mf6/read_input/grid_data.py @@ -0,0 +1,249 @@ +""" +Utilities for reading MODFLOW6 GRID DATA input. +""" +from pathlib import Path +from typing import IO, Any, Callable, Dict, Tuple + +import dask.array +import numpy as np + +from .common import ( + end_of_file, + find_entry, + read_external_binaryfile, + read_external_textfile, + read_internal, + split_line, + strip_line, + to_float, +) + + +def advance_to_griddata_section(f: IO[str]) -> Tuple[str, bool]: + line = None + # Empty line is end-of-file. + # Should always exit early in while loop, err otherwise. + while not end_of_file(line): + line = f.readline() + stripped = strip_line(line) + if stripped == "": + continue + elif "end" in stripped: + return None, False + elif "layered" in stripped: + layered = True + section = split_line(stripped)[0] + return section, layered + else: + layered = False + section = stripped + return section, layered + raise ValueError(f"No end of griddata specified in {f.name}") + + +def shape_to_max_rows(shape: Tuple[int], layered: bool) -> Tuple[int, Tuple[int]]: + """ + Compute the number of rows to read in case the data is internal. In case + of DIS, the number of (table) rows equals the number of layers times the + number of (model) rows; a single row then contains ncolumn values. + + A DISV does not have rows and columns, the number of rows is equal to the + number of layers. + + In case of LAYERED, each layer is written on a separate row. + + In case of DISU, all values are written one a single row, and LAYERED input + is not allowed. + + Parameters + ---------- + shape: Tuple[int] + layered: bool + + Returns + ---------- + max_rows: int + Reduced number if layered is True. + read_shape: Tuple[int] + First dimension removed if layered is True. + """ + ndim = len(shape) + if ndim == 3: + nlayer, nrow, _ = shape + max_rows_layered = nrow + max_rows = nlayer * nrow + + elif ndim == 2: + nlayer, _ = shape + max_rows_layered = 1 + max_rows = nlayer + + elif ndim == 1: + if layered: + raise ValueError( + "LAYERED section detected. DISU does not support LAYERED input." + ) + nlayer = None + max_rows_layered = None + max_rows = 1 + + else: + raise ValueError( + f"length of shape should be 1, 2, or 3. Received shape of length: {ndim}" + ) + + if layered: + return max_rows_layered, shape[1:] + else: + return max_rows, shape + + +def constant(value: Any, shape: Tuple[int], dtype: type) -> dask.array.Array: + return dask.array.full(shape, value, dtype=dtype) + + +def read_internal_griddata( + f: IO[str], dtype: type, shape: Tuple[int], max_rows: int +) -> np.ndarray: + return read_internal(f=f, dtype=dtype, max_rows=max_rows).reshape(shape) + + +def read_external_griddata( + path: Path, dtype: type, shape: Tuple[int], binary: bool +) -> np.ndarray: + max_rows = np.product(shape) + if binary: + a = read_external_binaryfile(path, dtype, max_rows) + else: + a = read_external_textfile(path, dtype, max_rows) + return a.reshape(shape) + + +def read_array( + f: IO[str], + simroot: Path, + dtype: type, + max_rows: int, + shape: Tuple[int], +) -> dask.array.Array: + """ + MODFLOW6 READARRAY functionality for grid data. + + External files are lazily read using dask, constants are lazily allocated, + and internal values are eagerly read and then converted to dask arrays for + consistency. + + Parameters + ---------- + f: IO[str] + simroot: Path + dtype: type + max_rows: int + Number of rows to read in case of internal data values. + shape: Tuple[int] + + Returns + ------- + array: dask.array.Array of size ``shape`` + """ + firstline = f.readline() + stripped = strip_line(firstline) + separated = split_line(stripped) + first = separated[0] + + if first == "constant": + factor = None + array = constant(separated[1], shape, dtype) + elif first == "internal": + factor = find_entry(stripped, "factor", to_float) + a = read_internal_griddata(f, dtype, shape, max_rows) + array = dask.array.from_array(a) + elif first == "open/close": + factor = find_entry(stripped, "factor", to_float) + fname = separated[1] + binary = "(binary)" in stripped + path = simroot / fname + a = dask.delayed(read_external_griddata)(path, dtype, shape, binary) + array = dask.array.from_delayed(a, shape=shape, dtype=dtype) + else: + raise ValueError( + 'Expected "constant", "internal", or "open/close". Received instead: ' + f"{stripped}" + ) + + if factor is not None: + # Cast the factor as well to make sure type is preserved. + array = array * dtype(factor) + + return array + + +def read_griddata( + f: IO[str], + simroot: Path, + sections: Dict[str, Tuple[type, Callable]], + shape: Tuple[int], +) -> Dict[str, dask.array.Array]: + """ + Read GRID DATA section. + + External files are lazily read using dask, constants are lazily allocated, + and internal values are eagerly read and then converted to dask arrays for + consistency. + + Parameters + ---------- + f: IO[str] + simroot: Path + Root path of simulation. Used for reading external files. + sections: Dict[str, (type, Callabe)] + Dictionary containing type of grid data entry, and function to compute + number of max_rows to read. + shape: Tuple[int] + Full shape of model read from DIS file. Should have length: + * DIS: 3 + * DISV: 2 + * DISU: 1 + + Returns + ------- + variables: Dict[str, dask.array.Array] + """ + content = {} + try: + section, layered = advance_to_griddata_section(f) + while section is not None: + try: + dtype, compute_max_rows = sections[section] + except KeyError: + raise ValueError( + f"Unexpected section: {section}. " + f"Expected one of: {', '.join(sections.keys())}" + ) + max_rows, read_shape = compute_max_rows(shape, layered) + kwargs = { + "f": f, + "simroot": simroot, + "dtype": dtype, + "shape": read_shape, + "max_rows": max_rows, + } + + if layered: + nlayer = shape[0] + data = dask.array.stack( + [read_array(**kwargs) for _ in range(nlayer)], + axis=0, + ) + else: + data = read_array(**kwargs) + + content[section] = data + + # Advance to next section + section, layered = advance_to_griddata_section(f) + + except Exception as e: + raise type(e)(f"{e}\n Error reading {f.name}") from e + + return content diff --git a/imod/mf6/read_input/list_input.py b/imod/mf6/read_input/list_input.py new file mode 100644 index 000000000..090ddd351 --- /dev/null +++ b/imod/mf6/read_input/list_input.py @@ -0,0 +1,181 @@ +""" +Utilities for reading MODFLOW6 list input. +""" +from pathlib import Path +from typing import IO, List, Tuple + +import dask.array +import numpy as np +import pandas as pd + +from .common import read_external_binaryfile, split_line, strip_line + + +def field_values( + recarr: np.ndarray, + fields: Tuple[str], +): + """ + Return the record array columns as a list of separate arrays. + """ + return [recarr[field] for field in fields] + + +def recarr_to_dense( + recarr: np.ndarray, + index_columns: Tuple[str], + fields: Tuple[str], + shape: Tuple[int], +) -> List[np.ndarray]: + """ + Convert a record array to separate numpy arrays. Uses the index columns to + place the values in a dense array form. + """ + # MODFLOW6 is 1-based, Python is 0-based + indices = [recarr[column] - 1 for column in index_columns] + variables = [] + for column in field_values(recarr, fields): + data = np.full(shape, np.nan) + data[indices] = column + variables.append(data) + return variables + + +def read_text_listinput( + path: Path, + dtype: np.dtype, + max_rows: int, +) -> np.ndarray: + # This function is over three times faster than: + # recarr = np.loadtxt(path, dtype, max_rows=max_rows) + # I guess MODFLOW6 will also accept comma delimited? + d = {key: value[0] for key, value in dtype.fields.items()} + df = pd.read_csv( + path, + header=None, + index_col=False, + dtype=d, + names=d.keys(), + delim_whitespace=True, + comment="#", + nrows=max_rows, + ) + return df.to_records(index=False) + + +def read_internal_listinput( + f: IO[str], + dtype: np.dtype, + index_columns: Tuple[str], + fields: Tuple[str], + shape: Tuple[int], + max_rows: int, + sparse_to_dense: bool, +) -> List[np.ndarray]: + # recarr = read_internal(f, max_rows, dtype) + recarr = read_text_listinput(f, dtype, max_rows) + if sparse_to_dense: + return recarr_to_dense(recarr, index_columns, fields, shape) + else: + return field_values(recarr, index_columns + fields) + + +def read_external_listinput( + path: Path, + dtype: np.dtype, + index_columns: Tuple[str], + fields: Tuple[str], + shape: Tuple[int], + binary: bool, + max_rows: int, + sparse_to_dense: bool, +) -> List[np.ndarray]: + """ + Read external list input, separate and reshape to a dense array form. + """ + if binary: + recarr = read_external_binaryfile(path, dtype, max_rows) + else: + recarr = read_text_listinput(path, dtype, max_rows) + if sparse_to_dense: + return recarr_to_dense(recarr, index_columns, fields, shape) + else: + return field_values(recarr, index_columns + fields) + + +def read_listinput( + f: IO[str], + simroot: Path, + dtype: type, + index_columns: Tuple[str], + fields: Tuple[str], + shape: Tuple[int], + max_rows: int, + sparse_to_dense: bool = True, +) -> List[dask.array.Array]: + """ + MODFLOW6 list input reading functionality. + + MODFLOW6 list input is "sparse": it consists of a cell id and a number of + values. Depending on whether the model is discretized according to DIS, + DISV, or DISU; this cell id may be a tuple of size 3, 2, or 1. + + Parameters + ---------- + f: IO[str] + File handle. + simroot: Path + Root path of simulation. Used for reading external files. + dtype: np.dtype + index_columns: Tuple[str] + fields: Tuple[str] + shape: Tuple[int] + max_rows: int + sparse_to_dense: bool + + Returns + ------- + variable_values: List[dask.array.Array] + A dask array for every entry in ``fields``. + """ + # Store position so week can move back in the file if data is stored + # internally. + position = f.tell() + + # Read and check the first line. + firstline = f.readline() + stripped = strip_line(firstline) + separated = split_line(stripped) + first = separated[0] + + nout = len(fields) + fieldtypes = [dtype.fields[field][0] for field in fields] + if not sparse_to_dense: + shape = (max_rows,) + nout += len(index_columns) + fieldtypes = [dtype.fields[field][0] for field in index_columns] + fieldtypes + + if first == "open/close": + fname = separated[1] + path = simroot / fname + binary = "(binary)" in stripped + + if binary and "boundnames" in dtype.fields: + raise ValueError("(BINARY) does not support BOUNDNAMES") + + x = dask.delayed(read_external_listinput, nout=nout)( + path, dtype, index_columns, fields, shape, binary, max_rows, sparse_to_dense + ) + variable_values = [ + dask.array.from_delayed(a, shape=shape, dtype=dt) + for a, dt in zip(x, fieldtypes) + ] + else: + # Move file position back one line, and try reading internal values. + f.seek(position) + x = read_internal_listinput( + f, dtype, index_columns, fields, shape, max_rows, sparse_to_dense + ) + variable_values = [dask.array.from_array(a) for a in x] + + return variable_values diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index 5e74e92ac..5b86630d2 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -2,6 +2,7 @@ import pathlib import subprocess import warnings +from typing import Union import jinja2 import numpy as np @@ -9,6 +10,8 @@ import imod +from .read_input import read_sim_namefile, read_tdis, tdis_time + class Modflow6Simulation(collections.UserDict): def _initialize_template(self): @@ -131,6 +134,39 @@ def render(self): d["solutiongroups"] = [[("ims6", f"{solvername}.ims", modelnames)]] return self._template.render(d) + @classmethod + def open(cls, path: Union[str, pathlib.Path]): + simroot = pathlib.Path(path).parent + name = simroot.stem + content = read_sim_namefile(path) + + # Get the times from the time discretization file. + tdis_path = simroot / content["tdis6"] + tdis = read_tdis(tdis_path) + globaltimes = tdis_time(tdis) + + # Initialize the simulation + simulation = cls(name=name) + + # Collect the models + MODEL_CLASSES = { + "gwf6": imod.mf6.GroundwaterFlowModel, + } + models = {} + for i, (ftype, fname, pname) in enumerate(content["models"]): + if pname is None: + modelname = f"{ftype}_{i+1}" + # Check for duplicate entry + elif pname in models: + modelname = f"{pname}_{i+1}" + else: + modelname = pname + + _class = MODEL_CLASSES[ftype] + simulation[modelname] = _class.open(fname, simroot, globaltimes) + + return simulation + def write(self, directory=".", binary=True): # Check models for required content for key, model in self.items(): diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 1ca1f9ca2..583646bce 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -1,7 +1,13 @@ +from pathlib import Path +from typing import Tuple + import numpy as np +import xarray as xr from imod.mf6.pkgbase import Package, VariableMetaData +from .read_input import read_sto_blockfile, shape_to_max_rows + class Storage(Package): def __init__(*args, **kwargs): @@ -10,7 +16,73 @@ def __init__(*args, **kwargs): ) -class SpecificStorage(Package): +class StorageBase(Package): + def _render(self, directory, pkgname, globaltimes, binary): + d = {} + stodirectory = directory / pkgname + for varname in self._grid_data.keys(): + key = self._keyword_map.get(varname, varname) + layered, value = self._compose_values( + self[varname], stodirectory, key, binary=binary + ) + if self._valid(value): # skip False or None + d[f"{key}_layered"], d[key] = layered, value + + periods = {} + if "time" in self.dataset["transient"].coords: + package_times = self.dataset["transient"].coords["time"].values + starts = np.searchsorted(globaltimes, package_times) + 1 + for i, s in enumerate(starts): + periods[s] = self.dataset["transient"].isel(time=i).values[()] + else: + periods[1] = self.dataset["transient"].values[()] + + d["periods"] = periods + return d + + @staticmethod + def open( + path: Path, + simroot: Path, + shape: Tuple[int], + coords: Tuple[str], + dims: Tuple[str], + globaltimes: np.ndarray, + ): + sections = { + "iconvert": (np.int32, shape_to_max_rows), + "ss": (np.float64, shape_to_max_rows), + "sy": (np.float64, shape_to_max_rows), + } + content = read_sto_blockfile( + path=simroot / path, + simroot=simroot, + sections=sections, + shape=shape, + ) + + griddata = content.pop("griddata") + for field, data in griddata.items(): + content[field] = xr.DataArray(data, coords, dims) + periods = content.pop("periods") + time_index = np.fromiter(periods.keys(), dtype=int) - 1 + content["transient"] = xr.DataArray( + list(periods.values()), + coords={"time": globaltimes[time_index]}, + dims=("time",), + ) + + storagecoefficient = content.pop("storagecoefficient", False) + if storagecoefficient: + cls = StorageCoefficient + else: + cls = SpecificStorage + + filtered_content = cls.filter_and_rename(content) + return cls(**filtered_content) + + +class SpecificStorage(StorageBase): """ Storage Package with specific storage. @@ -71,34 +143,14 @@ def __init__(self, specific_storage, specific_yield, transient, convertible): self._pkgcheck() def render(self, directory, pkgname, globaltimes, binary): - d = {} - stodirectory = directory / "sto" - for varname in ["specific_storage", "specific_yield", "convertible"]: - key = self._keyword_map.get(varname, varname) - layered, value = self._compose_values( - self[varname], stodirectory, key, binary=binary - ) - if self._valid(value): # skip False or None - d[f"{key}_layered"], d[key] = layered, value - - periods = {} - if "time" in self.dataset["transient"].coords: - package_times = self.dataset["transient"].coords["time"].values - starts = np.searchsorted(globaltimes, package_times) + 1 - for i, s in enumerate(starts): - periods[s] = self.dataset["transient"].isel(time=i).values[()] - else: - periods[1] = self.dataset["transient"].values[()] - - d["periods"] = periods - + d = self._render(directory, pkgname, globaltimes, binary) return self._template.render(d) -class StorageCoefficient(Package): +class StorageCoefficient(StorageBase): """ - Storage Package with a storage coefficient. Be careful, - this is not the same as the specific storage. + Storage Package with a storage coefficient. Be careful, this is not the + same as the specific storage. From wikipedia (https://en.wikipedia.org/wiki/Specific_storage): @@ -169,26 +221,6 @@ def __init__(self, storage_coefficient, specific_yield, transient, convertible): self._pkgcheck() def render(self, directory, pkgname, globaltimes, binary): - d = {} - stodirectory = directory / "sto" - for varname in ["storage_coefficient", "specific_yield", "convertible"]: - key = self._keyword_map.get(varname, varname) - layered, value = self._compose_values( - self[varname], stodirectory, key, binary=binary - ) - if self._valid(value): # skip False or None - d[f"{key}_layered"], d[key] = layered, value - - periods = {} - if "time" in self.dataset["transient"].coords: - package_times = self.dataset["transient"].coords["time"].values - starts = np.searchsorted(globaltimes, package_times) + 1 - for i, s in enumerate(starts): - periods[s] = self.dataset["transient"].isel(time=i).values[()] - else: - periods[1] = self.dataset["transient"].values[()] - - d["periods"] = periods + d = self._render(directory, pkgname, globaltimes, binary) d["storagecoefficient"] = True - return self._template.render(d) diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index b0c77af3d..551436a6f 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -5,6 +5,8 @@ from imod.mf6.pkgbase import BoundaryCondition, VariableMetaData +from .read_input import read_boundary_blockfile + def assign_dims(arg) -> Dict: is_da = isinstance(arg, xr.DataArray) @@ -109,6 +111,30 @@ def to_sparse(self, arrdict, layer): recarr[key] = arr return recarr + @classmethod + def open(cls, path, simroot, shape, coords, dims, globaltimes): + # read_boundary_blockfile calls: + # read_package_periods, which calls: + # read_listinput, which transforms recarrays to dense arrays by default. + content = read_boundary_blockfile( + simroot / path, + simroot, + fields=("rate",), + shape=shape, + sparse_to_dense=False, + ) + + period_index = content.pop("period_index") + period_data = content.pop("period_data") + coords = {"time": globaltimes[period_index]} + dims = ("time", "index") + + for field, data in period_data.items(): + content[field] = xr.DataArray(data, coords, dims) + + filtered_content = cls.filter_and_rename(content) + return cls(**filtered_content) + class WellDisVertices(BoundaryCondition): """ diff --git a/imod/templates/mf6/gwf-disu.j2 b/imod/templates/mf6/gwf-disu.j2 index c00936b55..82e8bdf66 100644 --- a/imod/templates/mf6/gwf-disu.j2 +++ b/imod/templates/mf6/gwf-disu.j2 @@ -1,49 +1,47 @@ begin options -{%- if length_units is defined -%} length_units {{length_units}}{%- endif -%} -{%- if nogrb is defined -%} nogrb{%- endif -%} -{%- if xorigin is defined -%} xorigin {{xorigin}}{%- endif -%} -{%- if yorigin is defined -%} yorigin {{yorigin}}{%- endif -%} -{%- if angrot is defined -%} angrot {{angrot}}{%- endif -%} +{% if length_units is defined %} length_units {{length_units}} +{% endif %} +{%- if nogrb is defined %} nogrb +{% endif %} +{%- if xorigin is defined %} xorigin {{xorigin}} +{% endif %} +{%- if yorigin is defined %} yorigin {{yorigin}} +{% endif %} +{%- if angrot is defined %} angrot {{angrot}} +{% endif -%} end options begin dimensions nodes {{nodes}} nja {{nja}} -{%- if nvert is defined -%} nvert {{nvert}}{%- endif -%} +{%- if nvert is defined %} nvert {{nvert}} +{% endif %} end dimensions begin griddata top - {top} + {{top}} bot - {bot} + {{bot}} area - {area} + {{area}} +{% if idomain is defined %} idomain + {{idomain}} +{% endif -%} end griddata begin connectiondata iac - {iac} + {{iac}} ja - {ja} + {{ja}} ihc - {ihc} + {{ihc}} cl12 - {cl12} + {{cl12}} hwva - {hwva} -{%- if angldegx is defined -%} angldegx - {angldegx}{%- endif -%} -end connectiondata - -begin vertices - {{iv}} {{xv}} {{yv}} - {{iv}} {{xv}} {{yv}} - ... -end vertices - -begin cell2d - {{icell2d}} {{xc}} {{yc}} {{ncvert}} {{icvert(ncvert)}} - {{icell2d}} {{xc}} {{yc}} {{ncvert}} {{icvert(ncvert)}} - ... -end cell2d + {{hwva}} +{% if angldegx is defined %} angldegx + {{angldegx}} +{% endif -%} +end connectiondata \ No newline at end of file diff --git a/imod/tests/test_mf6/test_ex01_twri.py b/imod/tests/test_mf6/test_ex01_twri.py index 29cf41c05..01eac3637 100644 --- a/imod/tests/test_mf6/test_ex01_twri.py +++ b/imod/tests/test_mf6/test_ex01_twri.py @@ -2,6 +2,7 @@ import sys import textwrap +import dask.array import numpy as np import pytest import xarray as xr @@ -460,3 +461,28 @@ def test_simulation_write_errors(twri_model, tmp_path): expected_message = "No sto package found in model GWF_1" with pytest.raises(ValueError, match=re.escape(expected_message)): simulation.write(modeldir, binary=True) + + +@pytest.mark.parametrize("binary", [True, False]) +@pytest.mark.usefixtures("transient_twri_model") +@pytest.mark.skipif(sys.version_info < (3, 7), reason="capture_output added in 3.7") +def test_simulation_write_and_open(transient_twri_model, tmp_path, binary): + simulation = transient_twri_model + modeldir = tmp_path / "ex01-twri-transient-binary" + simulation.write(modeldir, binary=binary) + + back = imod.mf6.Modflow6Simulation.open(modeldir / "mfsim.nam") + assert isinstance(back, imod.mf6.Modflow6Simulation) + + gwf = back["gwf_1"] + for name in ["chd", "drn", "ic", "npf", "rch", "sto", "wel"]: + assert name in gwf + + chd = gwf["chd"] + assert isinstance(chd, imod.mf6.ConstantHead) + assert tuple(chd.dataset["head"].dims) == ("time", "layer", "y", "x") + assert isinstance(chd.dataset["head"].data, dask.array.Array) + + head = chd["head"].dropna("layer", how="all").isel(time=0, drop=True).compute() + original = simulation["GWF_1"]["chd"]["head"] + assert head.equals(original) diff --git a/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading.py b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading.py new file mode 100644 index 000000000..88df4f67e --- /dev/null +++ b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading.py @@ -0,0 +1,52 @@ +import pytest + +from imod.mf6 import read_input as ri + + +def test_parse_model(): + ri.parse_model("gwf6 GWF_1.nam GWF_1", "sim.nam") == ["gwf6", "GWF_1_nam", "GWF_1"] + with pytest.raises(ValueError, match="ftype, fname and pname expected"): + ri.parse_model("gwf6 GWF_1.nam", "sim.nam") + + +def test_parse_exchange(): + ri.parse_exchange( + "GWF6-GWF6 simulation.exg GWF_Model_1 GWF_Model_2", "sim.nam" + ) == ["GWF6-GWF6, simulation.exg, GWF_Model_1, GWF_Model_2"] + with pytest.raises(ValueError, match="exgtype, exgfile, exgmnamea, exgmnameb"): + ri.parse_exchange("GWF6-GWF6 simulation.exg GWF_Model_1", "sim.nam") + + +def test_parse_solutiongroup(): + ri.parse_solutiongroup("ims6 solver.ims GWF_1", "sim.nam") == [ + "ims6", + "solver.ims", + "GWF_1", + ] + with pytest.raises(ValueError, match="Expected at least three entries"): + ri.parse_solutiongroup("ims6 solver.ims", "sim.name") + + +def test_parse_package(): + ri.parse_package("dis6 dis.dis", "gwf.nam") == ["dis6", "dis.dis", "dis"] + ri.parse_package("dis6 dis.dis discretization", "gwf.nam") == [ + "dis6", + "dis.dis", + "discretization", + ] + with pytest.raises(ValueError, match="Expected ftype, fname"): + ri.parse_package("dis6", "gwf.nam") + with pytest.raises(ValueError, match="Expected ftype, fname"): + ri.parse_package("dis6 dis.dis abc def", "gwf.nam") + + +def test_parse_tdis_perioddata(): + perlen, nstp, tsmult = ri.parse_tdis_perioddata("1.0 1 1.0", "timedis.tdis") + assert isinstance(perlen, float) + assert perlen == 1.0 + assert isinstance(nstp, int) + assert nstp == 1 + assert isinstance(tsmult, float) + assert tsmult == 1.0 + with pytest.raises(ValueError, match="perlen, nstp, tsmult expected"): + ri.parse_tdis_perioddata("1.0 1.0", "timedis.tdis") diff --git a/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_common.py b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_common.py new file mode 100644 index 000000000..2e40d1641 --- /dev/null +++ b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_common.py @@ -0,0 +1,249 @@ +import textwrap + +import numpy as np +import pytest + +from imod.mf6.read_input import common as cm + + +def isa(value, kind, expected): + return isinstance(value, kind) and value == expected + + +def test_end_of_file(): + assert cm.end_of_file("") is True + assert cm.end_of_file(" ") is False + assert cm.end_of_file("\n") is False + + +def test_strip_line(): + assert cm.strip_line("abc") == "abc" + assert cm.strip_line("abc ") == "abc" + assert cm.strip_line("ABC ") == "abc" + assert cm.strip_line("abc#def") == "abc" + assert cm.strip_line("abc #def") == "abc" + assert cm.strip_line(" abc ##def") == "abc" + + +def test_flatten(): + assert cm.flatten([[1, 2], [3, 4]]) == [1, 2, 3, 4] + assert cm.flatten([["abc", "def"], ["ghi"]]) == ["abc", "def", "ghi"] + + +def test_split_line(): + assert cm.split_line("abc def ghi") == ["abc", "def", "ghi"] + assert cm.split_line("abc,def,ghi") == ["abc", "def", "ghi"] + assert cm.split_line("abc,def, ghi") == ["abc", "def", "ghi"] + assert cm.split_line("abc ,def, ghi ") == ["abc", "def", "ghi"] + + +def test_to_float(): + assert isa(cm.to_float("1"), float, 1.0) + assert isa(cm.to_float("1.0"), float, 1.0) + assert isa(cm.to_float("1.0d0"), float, 1.0) + assert isa(cm.to_float("1.0e0"), float, 1.0) + assert isa(cm.to_float("1.0+0"), float, 1.0) + assert isa(cm.to_float("1.0-0"), float, 1.0) + assert isa(cm.to_float("1.0e-0"), float, 1.0) + assert isa(cm.to_float("1.0e+0"), float, 1.0) + + +def test_find_entry(): + assert isa(cm.find_entry("abc factor 1", "factor", float), float, 1.0) + assert isa(cm.find_entry("abc factor 1", "factor", int), int, 1) + assert isa(cm.find_entry("abc factor 1", "factor", str), str, "1") + assert cm.find_entry("abc 1", "factor", float) is None + + +def test_read_internal(tmp_path): + path1 = tmp_path / "internal-1.dat" + with open(path1, "w") as f: + f.write("1 2 3 4") + # max_rows should read all lines, unless max_rows is exceeded. + with open(path1) as f: + a = cm.read_internal(f, int, 1) + with open(path1) as f: + b = cm.read_internal(f, int, 2) + assert np.issubdtype(a.dtype, np.integer) + assert np.array_equal(a, b) + + path2 = tmp_path / "internal-2.dat" + with open(path2, "w") as f: + f.write("1 2\n3 4") + with open(path2) as f: + a = cm.read_internal(f, float, 1) + with open(path2) as f: + b = cm.read_internal(f, int, 2) + assert np.issubdtype(a.dtype, np.floating) + assert np.issubdtype(b.dtype, np.integer) + assert a.size == 2 + assert b.size == 4 + + +def test_read_external_binaryfile(tmp_path): + path1 = tmp_path / "external-1.bin" + a = np.ones((5, 5)) + a.tofile(path1) + b = cm.read_external_binaryfile(path1, np.float64, 25) + assert b.shape == (25,) + b = cm.read_external_binaryfile(path1, np.float64, 10) + assert b.shape == (10,) + + dtype = np.dtype([("node", np.int32), ("value", np.float32)]) + a = np.ones((5,), dtype) + path2 = tmp_path / "external-2.bin" + a.tofile(path2) + b = cm.read_external_binaryfile(path2, dtype, 5) + assert np.array_equal(a, b) + + +def test_read_fortran_deflated_text_array(tmp_path): + path1 = tmp_path / "deflated-1.txt" + with open(path1, "w") as f: + f.write("1.0\n2*2.0\n3*3.0") + a = cm.read_fortran_deflated_text_array(path1, float, 3) + b = np.array([1.0, 2.0, 2.0]) + assert np.array_equal(a, b) + + a = cm.read_fortran_deflated_text_array(path1, float, 6) + b = np.array([1.0, 2.0, 2.0, 3.0, 3.0, 3.0]) + assert np.array_equal(a, b) + + +def test_read_external_textfile(tmp_path): + path1 = tmp_path / "external.dat" + with open(path1, "w") as f: + f.write("1.0 2.0 3.0") + a = cm.read_external_textfile(path1, float, 6) + b = np.array([1.0, 2.0, 3.0]) + assert np.array_equal(a, b) + + path2 = tmp_path / "deflated-1.txt" + with open(path2, "w") as f: + f.write("1.0\n2*2.0\n3*3.0") + a = cm.read_external_textfile(path2, float, 6) + b = np.array([1.0, 2.0, 2.0, 3.0, 3.0, 3.0]) + assert np.array_equal(a, b) + + +def test_advance_to_header(tmp_path): + path = tmp_path / "header.txt" + content = textwrap.dedent( + """\ + begin options + 1 + end options + + begin griddata + 2 + end griddata + """ + ) + with open(path, "w") as f: + f.write(content) + + with open(path) as f: + cm.advance_to_header(f, "options") + assert f.readline() == "1\n" + f.readline() # read "end options" line + cm.advance_to_header(f, "griddata") + assert f.readline() == "2\n" + + with pytest.raises(ValueError, match='"begin perioddata" is not present'): + cm.advance_to_header(f, "perioddata") + + +def test_parse_option(): + with pytest.raises(ValueError, match="Cannot parse option in options.txt"): + cm.parse_option("", "options.txt") + assert cm.parse_option("save_flows", "options.txt") == ("save_flows", True) + assert cm.parse_option("variablecv dewatered", "options.txt") == ( + "variablecv", + "dewatered", + ) + assert cm.parse_option("multiple a b c", "options.txt") == ( + "multiple", + ["a", "b", "c"], + ) + + +def test_read_key_value_block(tmp_path): + path = tmp_path / "values.txt" + content = "\n".join(["", "save_flows", "variablecv dewatered", "", "end options"]) + + with open(path, "w") as f: + f.write(content) + + with open(path) as f: + d = cm.read_key_value_block(f, cm.parse_option) + assert d == {"save_flows": True, "variablecv": "dewatered"} + + path2 = tmp_path / "values-unterminated.txt" + content = "\n".join( + [ + "", + "save_flows", + "variablecv dewatered", + "", + ] + ) + + with open(path2, "w") as f: + f.write(content) + + with open(path2) as f: + with pytest.raises(ValueError, match='"end" of block is not present in file'): + cm.read_key_value_block(f, cm.parse_option) + + +def test_read_iterable_block(tmp_path): + def parse(line, _): # second arg is normally file name + return line.split() + + path = tmp_path / "iterable-values.txt" + content = "\n".join(["1.0 1 1.0", "1.0 1 1.0", "1.0 1 1.0", "end perioddata"]) + with open(path, "w") as f: + f.write(content) + + with open(path) as f: + back = cm.read_iterable_block(f, parse) + assert back == [ + ["1.0", "1", "1.0"], + ["1.0", "1", "1.0"], + ["1.0", "1", "1.0"], + ] + + path2 = tmp_path / "iterable-values-unterminated.txt" + content = textwrap.dedent( + """\ + 1.0 1 1.0 + 1.0 1 1.0 + 1.0 1 1.0 + """ + ) + with open(path2, "w") as f: + f.write(content) + with open(path2) as f: + with pytest.raises(ValueError, match='"end" of block is not present in file'): + back = cm.read_iterable_block(f, parse) + + +def test_parse_dimension(): + assert cm.parse_dimension(" nper 3 ", "fname") == ("nper", 3) + + +def test_advance_to_period(tmp_path): + path = tmp_path / "period.txt" + content = textwrap.dedent( + """\ + + begin period 1 + 2 + """ + ) + with open(path, "w") as f: + f.write(content) + + with open(path) as f: + cm.advance_to_period(f) + assert f.readline() == "2\n" diff --git a/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_griddata.py b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_griddata.py new file mode 100644 index 000000000..c41eabec7 --- /dev/null +++ b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_griddata.py @@ -0,0 +1,174 @@ +import textwrap + +import dask.array +import numpy as np +import pytest + +from imod.mf6.read_input import grid_data as gd + + +def test_advance_to_griddata_section(tmp_path): + path = tmp_path / "griddata.txt" + content = textwrap.dedent( + """\ + idomain + open/close GWF_1/dis/idomain.dat", + botm layered + """ + ) + + with open(path, "w") as f: + f.write(content) + + with open(path) as f: + assert gd.advance_to_griddata_section(f) == ("idomain", False) + f.readline() + assert gd.advance_to_griddata_section(f) == ("botm", True) + with pytest.raises(ValueError, match="No end of griddata specified"): + gd.advance_to_griddata_section(f) + + +def test_shape_to_max_rows(): + assert gd.shape_to_max_rows(shape=(4, 3, 2), layered=False) == (12, (4, 3, 2)) + assert gd.shape_to_max_rows(shape=(4, 3, 2), layered=True) == (3, (3, 2)) + assert gd.shape_to_max_rows(shape=(3, 2), layered=False) == (3, (3, 2)) + assert gd.shape_to_max_rows(shape=(3, 2), layered=True) == (1, (2,)) + assert gd.shape_to_max_rows(shape=(6,), layered=False) == (1, (6,)) + with pytest.raises( + ValueError, match="LAYERED section detected. DISU does not support LAYERED" + ): + gd.shape_to_max_rows(shape=(6,), layered=True) + with pytest.raises(ValueError, match="length of shape should be 1, 2, or 3"): + gd.shape_to_max_rows(shape=(5, 4, 3, 2), layered=True) + + +def test_constant(): + shape = (3, 4, 5) + a = gd.constant(2.0, shape, np.float64) + assert isinstance(a, dask.array.Array) + b = a.compute() + assert np.allclose(b, 2.0) + + +def test_read_internal_griddata(tmp_path): + path = tmp_path / "internal.txt" + with open(path, "w") as f: + f.write("1 2 3 4\n5 6 7 8") + + with open(path) as f: + a = gd.read_internal_griddata(f, np.int32, (2, 4), 2) + + assert a.shape == (2, 4) + assert np.array_equal(a.ravel(), np.arange(1, 9)) + + +def test_read_external_griddata(tmp_path): + path1 = tmp_path / "external.dat" + path2 = tmp_path / "external.bin" + + a = np.ones((4, 2)) + a.tofile(path1, sep=" ") + a.tofile(path2) # binary + + b = gd.read_external_griddata(path1, np.float64, (4, 2), False) + c = gd.read_external_griddata(path2, np.float64, (4, 2), True) + assert np.array_equal(a, b) + assert np.array_equal(a, c) + + +def test_read_array(tmp_path): + blockfile_path = tmp_path / "blockfile.txt" + external_path = tmp_path / "external.dat" + external_binary_path = tmp_path / "external.bin" + + content = textwrap.dedent( + """\ + top + constant 200.0 + idomain + open/close external.dat + botm + open/close external.bin (binary) + abc + """ + ) + with open(blockfile_path, "w") as f: + f.write(content) + + shape = (3, 4, 5) + a = np.ones(shape, dtype=np.float64) + a.tofile(external_path, sep=" ") + a.tofile(external_binary_path) + + with open(blockfile_path) as f: + f.readline() + top = gd.read_array(f, tmp_path, np.float64, max_rows=12, shape=shape) + f.readline() + idomain = gd.read_array(f, tmp_path, np.float64, max_rows=12, shape=shape) + f.readline() + botm = gd.read_array(f, tmp_path, np.float64, max_rows=12, shape=shape) + with pytest.raises( + ValueError, match='Expected "constant", "internal", or "open/close"' + ): + gd.read_array(f, tmp_path, np.float64, max_rows=12, shape=shape) + + for a in (top, idomain, botm): + assert isinstance(a, dask.array.Array) + assert a.shape == shape + + assert np.allclose(top, 200.0) + assert np.allclose(idomain, 1) + assert np.allclose(botm, 1.0) + + +def test_read_griddata(tmp_path): + blockfile_path = tmp_path / "blockfile.txt" + external_path = tmp_path / "external.dat" + external_binary_path = tmp_path / "external.bin" + + content = textwrap.dedent( + """\ + top + constant 200.0 + idomain + open/close external.dat + botm + open/close external.bin (binary) + end + """ + ) + + with open(blockfile_path, "w") as f: + f.write(content) + + sections = { + "top": (np.float64, gd.shape_to_max_rows), + "idomain": (np.int32, gd.shape_to_max_rows), + "botm": (np.float64, gd.shape_to_max_rows), + } + shape = (3, 4, 5) + a = np.ones(shape, dtype=np.int32) + a.tofile(external_path, sep=" ") + b = np.ones(shape, dtype=np.float64) + b.tofile(external_binary_path) + + with open(blockfile_path) as f: + d = gd.read_griddata(f, tmp_path, sections, shape) + + for key in ("top", "idomain", "botm"): + assert key in d + a = d[key] + assert isinstance(a, dask.array.Array) + assert a.shape == shape + + assert np.allclose(d["top"], 200.0) + assert np.allclose(d["idomain"], 1) + assert np.allclose(d["botm"], 1.0) + + dummy_path = tmp_path / "dummy.txt" + with open(dummy_path, "w") as f: + f.write("\n") + + with open(dummy_path) as f: + with pytest.raises(ValueError, match="Error reading"): + d = gd.read_griddata(f, tmp_path, sections, shape) diff --git a/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_listinput.py b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_listinput.py new file mode 100644 index 000000000..87925119a --- /dev/null +++ b/imod/tests/test_mf6/test_input_reading/test_mf6_input_reading_listinput.py @@ -0,0 +1,242 @@ +import textwrap + +import dask.array +import numpy as np + +from imod.mf6.read_input import list_input as li + +DIS_DTYPE = np.dtype( + [ + ("layer", np.int32), + ("row", np.int32), + ("column", np.int32), + ("head", np.float64), + ("conductance", np.float64), + ] +) + + +def test_recarr_to_dense__dis(): + dtype = DIS_DTYPE + recarr = np.array( + [ + (1, 1, 1, 1.0, 10.0), + (2, 1, 1, 2.0, 10.0), + (3, 1, 1, 3.0, 10.0), + ], + dtype=dtype, + ) + + variables = li.recarr_to_dense( + recarr, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + shape=(3, 4, 5), + ) + assert isinstance(variables, list) + assert len(variables) == 2 + a, b = variables + assert np.isfinite(a).sum() == 3 + assert np.isfinite(b).sum() == 3 + assert a[0, 0, 0] == 1.0 + assert a[1, 0, 0] == 2.0 + assert a[2, 0, 0] == 3.0 + + +def test_recarr_to_dense__disv(): + dtype = np.dtype( + [ + ("layer", np.int32), + ("cell2d", np.int32), + ("head", np.float64), + ("conductance", np.float64), + ] + ) + recarr = np.array( + [ + (1, 1, 1.0, 10.0), + (2, 1, 2.0, 10.0), + (3, 1, 3.0, 10.0), + ], + dtype=dtype, + ) + variables = li.recarr_to_dense( + recarr, + index_columns=["layer", "cell2d"], + fields=["head", "conductance"], + shape=(3, 20), + ) + assert isinstance(variables, list) + assert len(variables) == 2 + a, b = variables + assert np.isfinite(a).sum() == 3 + assert np.isfinite(b).sum() == 3 + assert a[0, 0] == 1.0 + assert a[1, 0] == 2.0 + assert a[2, 0] == 3.0 + + +def test_recarr_to_dense__disu(): + dtype = np.dtype( + [ + ("node", np.int32), + ("head", np.float64), + ("conductance", np.float64), + ] + ) + recarr = np.array( + [ + (1, 1.0, 10.0), + (21, 2.0, 10.0), + (41, 3.0, 10.0), + ], + dtype=dtype, + ) + variables = li.recarr_to_dense( + recarr, + index_columns=["node"], + fields=["head", "conductance"], + shape=(60,), + ) + assert isinstance(variables, list) + assert len(variables) == 2 + a, b = variables + assert np.isfinite(a).sum() == 3 + assert np.isfinite(b).sum() == 3 + assert a[0] == 1.0 + assert a[20] == 2.0 + assert a[40] == 3.0 + + +def test_read_text_listinput(tmp_path): + dtype = DIS_DTYPE + path = tmp_path / "listinput.dat" + content = textwrap.dedent( + """\ + # layer row column head conductance + 1 1 1 1.0 10.0 + 2 1 1 2.0 10.0 + 3 1 1 3.0 10.0 + """ + ) + + with open(path, "w") as f: + f.write(content) + + # Test for internal input as well, for an already opened file. + with open(path) as f: + variables = li.read_internal_listinput( + f, + dtype, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + max_rows=3, + shape=(3, 4, 5), + sparse_to_dense=True, + ) + assert isinstance(variables, list) + assert len(variables) == 2 + + variables = li.read_external_listinput( + path, + dtype, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + shape=(3, 4, 5), + binary=False, + max_rows=3, + sparse_to_dense=True, + ) + assert isinstance(variables, list) + assert len(variables) == 2 + + +def test_read_binary_listinput(tmp_path): + path = tmp_path / "listinput.bin" + dtype = DIS_DTYPE + recarr = np.array( + [ + (1, 1, 1, 1.0, 10.0), + (2, 1, 1, 2.0, 10.0), + (3, 1, 1, 3.0, 10.0), + ], + dtype=dtype, + ) + recarr.tofile(path) + + variables = li.read_external_listinput( + path, + dtype, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + shape=(3, 4, 5), + binary=True, + max_rows=3, + sparse_to_dense=True, + ) + assert isinstance(variables, list) + assert len(variables) == 2 + + +def test_read_listinput(tmp_path): + dtype = DIS_DTYPE + path = tmp_path / "package-binary.txt" + content = "\n".join( + [ + "open/close listinput.bin (binary)", + ] + ) + with open(path, "w") as f: + f.write(content) + + binpath = tmp_path / "listinput.bin" + dtype = DIS_DTYPE + recarr = np.array( + [ + (1, 1, 1, 1.0, 10.0), + (2, 1, 1, 2.0, 10.0), + (3, 1, 1, 3.0, 10.0), + ], + dtype=dtype, + ) + recarr.tofile(binpath) + + with open(path) as f: + variables = li.read_listinput( + f, + tmp_path, + dtype, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + shape=(3, 4, 5), + max_rows=3, + ) + + assert len(variables) == 2 + for a in variables: + assert isinstance(a, dask.array.Array) + assert a.shape == (3, 4, 5) + notnull = np.isfinite(a) + assert notnull[0, 0, 0] + assert notnull[1, 0, 0] + assert notnull[2, 0, 0] + + # Test another round, this time without converting from COO to a dense + # form. + with open(path) as f: + variables = li.read_listinput( + f, + tmp_path, + dtype, + index_columns=["layer", "row", "column"], + fields=["head", "conductance"], + shape=(3, 4, 5), + max_rows=3, + sparse_to_dense=False, + ) + + assert len(variables) == 5 + for a in variables: + assert isinstance(a, dask.array.Array) + assert a.shape == (3,) + assert np.isfinite(a).all() diff --git a/imod/tests/test_mf6/test_mf6_disu.py b/imod/tests/test_mf6/test_mf6_disu.py new file mode 100644 index 000000000..4b8e1d005 --- /dev/null +++ b/imod/tests/test_mf6/test_mf6_disu.py @@ -0,0 +1,167 @@ +import pathlib +import textwrap + +import numpy as np +import pytest +import xarray as xr + +from imod.mf6 import disu + + +@pytest.fixture(scope="function") +def structured_dis(): + coords = { + "layer": [1, 2, 3], + "y": [25.0, 15.0, 5.0], + "x": [50.0, 70.0, 90.0, 110.0], + } + dims = ["layer", "y", "x"] + idomain = xr.DataArray(np.ones((3, 3, 4), dtype=np.int32), coords, dims) + top = xr.DataArray(np.ones((3, 3, 4)), coords, dims) + bottom = xr.DataArray(np.ones((3, 3, 4)), coords, dims) + top.values[0] = 45.0 + top.values[1] = 30.0 + top.values[2] = 15.0 + bottom.values[0] = 30.0 + bottom.values[1] = 15.0 + bottom.values[2] = 0.0 + return idomain, top, bottom + + +def connectivity_checks(dis): + ncell = dis.dataset["node"].size + i = np.repeat(np.arange(1, ncell + 1), dis.dataset["iac"].values) - 1 + j = dis.dataset["ja"].values - 1 + diff = np.abs(i - j) + ihc = dis.dataset["ihc"].values + hwva = dis.dataset["hwva"].values + cl12 = dis.dataset["cl12"].values + + # The node number itself is marked by a negative number. + connection = j > 0 + vertical = (ihc == 0) & connection + + assert dis.dataset["iac"].values.sum() == j.size + assert i.min() == 0 + assert i.max() == ncell - 1 + assert j.max() == ncell - 1 + assert (diff[connection] > 0).all() + assert (cl12[connection] != 0).all() + assert (hwva[connection] != 0).all() + assert (diff[vertical] > 1).all() + assert np.allclose(cl12[vertical], 7.5) + assert np.allclose(hwva[vertical], 200.0) + + +def test_cell_number(): + nrow = 4 + ncol = 3 + assert disu._number(0, 0, 0, nrow, ncol) == 0 + assert disu._number(0, 0, 1, nrow, ncol) == 1 + assert disu._number(0, 1, 0, nrow, ncol) == 3 + assert disu._number(0, 1, 1, nrow, ncol) == 4 + assert disu._number(1, 0, 0, nrow, ncol) == 12 + assert disu._number(1, 0, 1, nrow, ncol) == 13 + assert disu._number(1, 1, 0, nrow, ncol) == 15 + assert disu._number(1, 1, 1, nrow, ncol) == 16 + + +def test_structured_connectivity_full(): + idomain = np.ones((3, 3, 4), dtype=np.int32) + i, j = disu._structured_connectivity(idomain) + + expected_i = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) + expected_j = np.array([1, 4, 12, 2, 5, 13, 3, 6, 14, 7, 15]) + assert np.array_equal(i[:11], expected_i) + assert np.array_equal(j[:11], expected_j) + + idomain[1, 0, 1] = -1 + idomain[1, 0, 2] = -1 + i, j = disu._structured_connectivity(idomain) + + expected_i = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) + expected_j = np.array([1, 4, 12, 2, 5, 25, 3, 6, 26, 7, 15]) + assert np.array_equal(i[:11], expected_i) + assert np.array_equal(j[:11], expected_j) + + +def test_from_structured(structured_dis): + idomain, top, bottom = structured_dis + dis = disu.LowLevelUnstructuredDiscretization.from_dis(top, bottom, idomain) + assert np.allclose(dis.dataset["xorigin"], 40.0) + assert np.allclose(dis.dataset["yorigin"], 0.0) + connectivity_checks(dis) + + # Now disable some cells, create one pass-through + idomain.values[1, 0, 1] = -1 + idomain.values[:, 0, 0] = 0 + dis = disu.LowLevelUnstructuredDiscretization.from_dis(top, bottom, idomain) + connectivity_checks(dis) + + +def test_render(structured_dis): + idomain, top, bottom = structured_dis + dis = disu.LowLevelUnstructuredDiscretization.from_dis(top, bottom, idomain) + + directory = pathlib.Path("mymodel") + actual = dis.render(directory, "disu", None, True) + + expected = textwrap.dedent( + """\ + begin options + xorigin 40.0 + yorigin 0.0 + end options + + begin dimensions + nodes 36 + nja 186 + end dimensions + + begin griddata + top + open/close mymodel/disu/top.bin (binary) + bot + open/close mymodel/disu/bot.bin (binary) + area + open/close mymodel/disu/area.bin (binary) + idomain + open/close mymodel/disu/idomain.bin (binary) + end griddata + + begin connectiondata + iac + open/close mymodel/disu/iac.bin (binary) + ja + open/close mymodel/disu/ja.bin (binary) + ihc + open/close mymodel/disu/ihc.bin (binary) + cl12 + open/close mymodel/disu/cl12.bin (binary) + hwva + open/close mymodel/disu/hwva.bin (binary) + end connectiondata""" + ) + print(actual) + print(expected) + assert actual == expected + + +def test_write(structured_dis, tmp_path): + idomain, top, bottom = structured_dis + dis = disu.LowLevelUnstructuredDiscretization.from_dis(top, bottom, idomain) + dis.write(tmp_path, "disu", None, True) + + assert (tmp_path / "disu.disu").exists + names = [ + "top.bin", + "bot.bin", + "area.bin", + "iac.bin", + "ja.bin", + "ihc.bin", + "cl12.bin", + "hwva.bin", + ] + for name in names: + assert (tmp_path / name).exists diff --git a/imod/tests/test_mf6/test_mf6_drn.py b/imod/tests/test_mf6/test_mf6_drn.py index 206108f09..f380833e3 100644 --- a/imod/tests/test_mf6/test_mf6_drn.py +++ b/imod/tests/test_mf6/test_mf6_drn.py @@ -84,3 +84,23 @@ def test_3d_singelayer(): arrdict = drn._ds_to_arrdict(bin_ds) sparse_data = drn.to_sparse(arrdict, layer) assert isinstance(sparse_data, np.ndarray) + + +def test_to_disu(drainage): + drn = drainage + disu_drn = drn.to_disu() + assert tuple(disu_drn["conductance"].dims) == ("node",) + assert np.array_equal(disu_drn.dataset["node"], np.arange(75)) + + drn["layer"] = [1, 3, 5] + disu_drn = drn.to_disu() + expected = np.concatenate([np.arange(25), np.arange(50, 75), np.arange(100, 125)]) + assert np.array_equal(disu_drn.dataset["node"], expected) + + onelayer = imod.mf6.Drainage(**drainage.dataset.isel(layer=0)) + disu_drn = onelayer.to_disu() + assert np.array_equal(disu_drn.dataset["node"], np.arange(25)) + + nolayer = imod.mf6.Drainage(**drainage.dataset.isel(layer=0, drop=True)) + with pytest.raises(ValueError, match="layer coordinate is required"): + nolayer.to_disu() diff --git a/imod/tests/test_mf6/test_mf6_npf.py b/imod/tests/test_mf6/test_mf6_npf.py index 735f16897..607807017 100644 --- a/imod/tests/test_mf6/test_mf6_npf.py +++ b/imod/tests/test_mf6/test_mf6_npf.py @@ -68,3 +68,43 @@ def test_wrong_dtype(): perched=True, save_flows=True, ) + + +def test_to_disu(): + layer = [1, 2, 3] + template = xr.full_like( + imod.util.empty_3d(10.0, 0.0, 100.0, 10.0, 0.0, 100.0, layer), 1, dtype=np.int32 + ) + icelltype = xr.DataArray([1, 0, 0], {"layer": layer}, ("layer",)) * template + k = xr.DataArray([1.0e-3, 1.0e-4, 2.0e-4], {"layer": layer}, ("layer",)) * template + k33 = ( + xr.DataArray([2.0e-8, 2.0e-8, 2.0e-8], {"layer": layer}, ("layer",)) * template + ) + + npf = imod.mf6.NodePropertyFlow( + icelltype=icelltype, + k=k, + k33=k33, + variable_vertical_conductance=True, + dewatered=True, + perched=True, + save_flows=True, + ) + + disu_npf = npf.to_disu() + assert np.array_equal(disu_npf.dataset["node"], np.arange(300)) + + # Test again, with partially active domain + active = xr.full_like(template, True, dtype=bool) + active[:, :, 0] = False + n_active = active.sum() + cell_ids = np.full((3, 10, 10), -1) + cell_ids[active.values] = np.arange(n_active) + cell_ids = cell_ids.ravel() + + disu_npf = npf.to_disu(cell_ids) + assert np.array_equal(disu_npf.dataset["node"], np.arange(270)) + + with pytest.raises(TypeError, match="cell_ids should be integer"): + cell_ids = np.arange(3 * 5 * 5, dtype=np.float64) + npf.to_disu(cell_ids) diff --git a/imod/util.py b/imod/util.py index bf6d6cf04..2d069343b 100644 --- a/imod/util.py +++ b/imod/util.py @@ -1073,6 +1073,38 @@ def _replace( ) +def reshape(da: xr.DataArray, coords: Dict[str, Any]): + """ + Gives a new shape and coordinates to a DataArray without changing its data. + + Parameters + ---------- + da: xr.DataArray + coords: dict + Coordinates of the reshaped DataArray. The key order determines + dimension order of the reshaped array. The sizes of the values + determines the size of the dimensions. Scalar values are expanded into + size 1 arrays. Multi-dimensional coordinate values are not allowed. + + Returns + ------- + reshaped: xr.DataArray + + See also + -------- + DataArray.stack + DataArray.unstack + """ + coords = {k: np.atleast_1d(v) for k, v in coords.items()} + shape = [v.size for v in coords.values()] + for key, value in coords.items(): + if value.ndim != 1: + raise ValueError( + f"{key} has ndim {value.ndim}. Can only reshape with 1D coordinates." + ) + return xr.DataArray(da.data.reshape(shape), coords=coords, dims=list(coords.keys())) + + class MissingOptionalModule: """ Presents a clear error for optional modules.