diff --git a/imod/wq/ani.py b/imod/wq/ani.py index 29381ee44..d576d854f 100644 --- a/imod/wq/ani.py +++ b/imod/wq/ani.py @@ -1,79 +1,21 @@ -import pathlib -import re -import shutil - import numpy as np import imod from imod.wq.pkgbase import Package +FLOAT_FORMAT = "%.18G" -class HorizontalAnisotropyFile(Package): - """ - Horizontal Anisotropy package. - Parameters - ---------- - anifile: str - is the file location of the imod-wq ani-file. This file contains the - anisotropy factor and angle of each layer, either as a constant or a - reference to the file location of an '.arr' file. No checks are - implemented for this file, user is responsible for consistency with - model. +def _write_arr(path, da, nodata=1.0e20, *_): """ - - _pkg_id = "ani" - - _template = "[ani]\n" " anifile = {anifile}\n" - - def __init__( - self, - anifile, - ): - super().__init__() - self["anifile"] = anifile - - def _render(self, modelname, directory, nlayer): - # write ani file - # in render function, as here we know where the model is run from - # and how many layers: store this info for later use in save - self.anifile = f"{modelname}.ani" - self.rendir = directory - self.nlayer = nlayer - - d = {"anifile": f"{directory.as_posix()}/{modelname}.ani"} - - return self._template.format(**d) - - def save(self, directory): - """Overload save function. - Saves anifile to location, along with referenced .arr files - assumes _render() to have run previously""" - directory.mkdir(exist_ok=True) # otherwise handled by idf.save - - path_ani = pathlib.Path(str(self["anifile"].values)) - - # regex adapted from stackoverflow: https://stackoverflow.com/questions/54990405/a-general-regex-to-extract-file-paths-not-urls - rgx = r"((?:[a-zA-Z]:|(? 0)).all( + axis=1 + ) + else: + layer, edge = np.argwhere(notnull).transpose() + layer2d = np.repeat(layer, 2).reshape((-1, 2)) + cell2d = edge_faces[edge] + valid = ((cell2d != grid.fill_value) & (idomain2d[layer2d, cell2d] > 0)).all( + axis=1 + ) + layer = layer[valid] + + # Skip the exterior edges (marked with a fill value). + cell2d = cell2d[valid] + c = resistance[notnull][valid] + + # Define the numpy structured array dtype + field_spec = [ + ("layer_1", np.int32), + ("row_1", np.int32), + ("column_1", np.int32), + ("layer_2", np.int32), + ("row_2", np.int32), + ("column_2", np.int32), + ("resistance", np.float64), + ] + dtype = np.dtype(field_spec) + shape = (ibound["y"].size, ibound["x"].size) + row_1, column_1 = np.unravel_index(cell2d[:, 0], shape) + row_2, column_2 = np.unravel_index(cell2d[:, 1], shape) + # Set the indices + recarr = np.empty(len(cell2d), dtype=dtype) + recarr["layer_1"] = layer + 1 + recarr["row_1"] = row_1 + 1 + recarr["column_1"] = column_1 + 1 + recarr["row_2"] = row_2 + 1 + recarr["column_2"] = column_2 + 1 + recarr["resistance"] = c + return recarr + + class HorizontalFlowBarrier(Package): """ Horizontal Flow Barrier package. Parameters ---------- - hfbfile: str - is the file location of the imod-wq hfb-file. This file contains cell- - to-cell resistance values. The hfb file can be constructed from generate - files using imod-batch. No checks are implemented for this file, user is - responsible for consistency with model. + resistance: xu.UgridDataArray + ibound: xr.DataArray """ _pkg_id = "hfb6" - _template = "[hfb6]\n" " hfbfile = {hfbfile}\n" + _template = "[hfb6]\n hfbfile = {hfbfile}\n" def __init__( self, - hfbfile, + resistance, + ibound, ): - super().__init__() - self["hfbfile"] = hfbfile - - def _render(self, modelname, directory, *args, **kwargs): - self.hfbfile = f"{modelname}.hfb" - d = {"hfbfile": f"{directory.as_posix()}/{modelname}.hfb"} + self.dataset = xu.UgridDataset() + self["resistance"] = resistance + self["ibound"] = ibound.rename({"layer": "ibound_layer"}) + def _render(self, directory, *args, **kwargs): + d = {"hfbfile": (directory / "horizontal_flow_barrier.hfb").as_posix()} return self._template.format(**d) def save(self, directory): - """Overloads save function. - Saves hfbfile to directory - assumes _render() to have run previously""" directory.mkdir(exist_ok=True) # otherwise handled by idf.save + path = (directory / "horizontal_flow_barrier.hfb").as_posix() - path_hfb = pathlib.Path(str(self["hfbfile"].values)) + hfb_array = create_hfb_array( + notnull=self["resistance"].notnull().values, + resistance=self["resistance"].values, + layer=self["resistance"].coords["layer"].values, + ibound=self["ibound"], + grid=self.dataset.ugrid.grid, + ) + nhfbnp = len(hfb_array) + header = f"0 0 {nhfbnp} 0" + fmt = ("%i",) * 5 + ("%.18G",) - shutil.copyfile(path_hfb, directory / self.hfbfile) + with open(path, "w") as f: + np.savetxt(fname=f, X=hfb_array, fmt=fmt, header=header) def _pkgcheck(self, ibound=None): pass diff --git a/imod/wq/model.py b/imod/wq/model.py index aeff8a94e..0351c32df 100644 --- a/imod/wq/model.py +++ b/imod/wq/model.py @@ -487,18 +487,21 @@ def _render_gen(self, modelname, globaltimes, writehelp, result_dir): d["start_date"] = start_date return self._gen_template.render(d) - def _render_pkg(self, key, directory, globaltimes, nlayer): + def _render_required_pkg(self, key, directory, globaltimes, nlayer): """ Rendering method for straightforward packages """ - # Get name of pkg, e.g. lookup "recharge" for rch _pkg_id pkgkey = self._get_pkgkey(key) if pkgkey is None: - # Maybe do enum look for full package name? - if (key == "rch") or (key == "evt"): # since recharge is optional - return "" - else: - raise ValueError(f"No {key} package provided.") + raise ValueError(f"No {key} package provided.") + return self[pkgkey]._render( + directory=directory / pkgkey, globaltimes=globaltimes, nlayer=nlayer + ) + + def _render_optional_pkg(self, key, directory, globaltimes, nlayer): + pkgkey = self._get_pkgkey(key) + if pkgkey is None: + return "" return self[pkgkey]._render( directory=directory / pkgkey, globaltimes=globaltimes, nlayer=nlayer ) @@ -638,21 +641,16 @@ def render(self, directory, result_dir, writehelp): ) ) # Modflow - for key in ("bas6", "oc", "lpf", "rch", "evt"): + for key in ("bas6", "oc", "lpf"): content.append( - self._render_pkg( + self._render_required_pkg( key=key, directory=directory, globaltimes=globaltimes, nlayer=nlayer ) ) - - # ani and hfb packages - for key in ("ani", "hfb6"): + for key in ("rch", "evt", "ani", "hfb"): content.append( - self._render_anihfb( - key=key, - modelname=self.modelname, - directory=directory, - nlayer=nlayer, + self._render_optional_pkg( + key=key, directory=directory, globaltimes=globaltimes, nlayer=nlayer ) )