From 55c1e893adca954bb1862a66160212ab4d9dade0 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 6 Mar 2024 13:50:16 +0100 Subject: [PATCH 01/29] work started --- imod/mf6/utilities/regrid.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 3fb6a87df..1dbc7c81e 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -5,9 +5,10 @@ import xarray as xr import xugrid as xu from xugrid.regrid.regridder import BaseRegridder - +from fastcore.dispatch import typedispatch from imod.typing.grid import GridDataArray - +from imod.mf6.utilities.clip import clip_by_grid +from imod.mf6.interfaces.ilinedatapackage import ILineDataPackage class RegridderType(Enum): """ @@ -145,3 +146,15 @@ def assign_coord_if_present( {coordname: target_grid.coords[coordname].values[()]} ) return maybe_has_coords_attr + +@typedispatch # type: ignore[no-redef] +def regrid_like( + package: ILineDataPackage, target_grid: GridDataArray, *_) -> ILineDataPackage: + """ + The regrid_like method is irrelevant for this package as it is + grid-agnostic, instead this method clips the package based on the grid + exterior. + """ + return clip_by_grid(package, target_grid) + + From 44b593ce18aa894058d34246ae266cc98c0dd620 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 6 Mar 2024 16:35:05 +0100 Subject: [PATCH 02/29] working on regridding packages --- imod/mf6/chd.py | 3 +- imod/mf6/cnc.py | 4 +- imod/mf6/dis.py | 3 +- imod/mf6/disv.py | 4 +- imod/mf6/drn.py | 4 +- imod/mf6/dsp.py | 4 +- imod/mf6/evt.py | 3 +- imod/mf6/ghb.py | 3 +- imod/mf6/ic.py | 4 +- imod/mf6/interfaces/igridpackage.py | 11 ++ imod/mf6/ist.py | 4 +- imod/mf6/mst.py | 4 +- imod/mf6/npf.py | 4 +- imod/mf6/package.py | 97 +-------------- imod/mf6/rch.py | 5 +- imod/mf6/riv.py | 4 +- imod/mf6/src.py | 4 +- imod/mf6/sto.py | 4 +- imod/mf6/utilities/regrid.py | 164 ++++++++++++++++++++++--- imod/mf6/utilities/regridding_types.py | 15 +++ 20 files changed, 209 insertions(+), 139 deletions(-) create mode 100644 imod/mf6/interfaces/igridpackage.py create mode 100644 imod/mf6/utilities/regridding_types.py diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py index bb79c84e2..583f9dcdc 100644 --- a/imod/mf6/chd.py +++ b/imod/mf6/chd.py @@ -13,9 +13,10 @@ IndexesSchema, OtherCoordsSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage -class ConstantHead(BoundaryCondition): +class ConstantHead(BoundaryCondition, IGridPackage): """ Constant-Head package. Any number of CHD Packages can be specified for a single groundwater flow model; however, an error will occur if a CHD Package diff --git a/imod/mf6/cnc.py b/imod/mf6/cnc.py index 284ff9612..b27cc19c5 100644 --- a/imod/mf6/cnc.py +++ b/imod/mf6/cnc.py @@ -11,9 +11,9 @@ IndexesSchema, OtherCoordsSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class ConstantConcentration(BoundaryCondition): +class ConstantConcentration(BoundaryCondition, IGridPackage): """ Constant Concentration package. diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index d75075a78..0aa69549a 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -16,8 +16,9 @@ IndexesSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage -class StructuredDiscretization(Package): +class StructuredDiscretization(Package, IGridPackage): """ Discretization information for structered grids is specified using the file. (DIS6) Only one discretization input file (DISU6, DISV6 or DIS6) can be diff --git a/imod/mf6/disv.py b/imod/mf6/disv.py index b062201ef..a48c28bfc 100644 --- a/imod/mf6/disv.py +++ b/imod/mf6/disv.py @@ -14,8 +14,8 @@ IndexesSchema, ) - -class VerticesDiscretization(Package): +from imod.mf6.interfaces.igridpackage import IGridPackage +class VerticesDiscretization(Package, IGridPackage): """ Discretization by Vertices (DISV). diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index 5c6dd7d21..3aa506ceb 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -14,9 +14,9 @@ IndexesSchema, OtherCoordsSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class Drainage(BoundaryCondition): +class Drainage(BoundaryCondition, IGridPackage): """ The Drain package is used to simulate head-dependent flux boundaries. https://water.usgs.gov/ogw/modflow/mf6io.pdf#page=67 diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index 5370a3ae1..a224ef45c 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -9,9 +9,9 @@ IdentityNoDataSchema, IndexesSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class Dispersion(Package): +class Dispersion(Package, IGridPackage): """ Molecular Diffusion and Dispersion. diff --git a/imod/mf6/evt.py b/imod/mf6/evt.py index 6a2e4a909..110e6685d 100644 --- a/imod/mf6/evt.py +++ b/imod/mf6/evt.py @@ -31,8 +31,9 @@ | DimsSchema("segment", "{face_dim}") ) +from imod.mf6.interfaces.igridpackage import IGridPackage -class Evapotranspiration(BoundaryCondition): +class Evapotranspiration(BoundaryCondition, IGridPackage): """ Evapotranspiration (EVT) Package. Any number of EVT Packages can be specified for a single groundwater flow diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index 7fc1c7b5d..d7f4657b3 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -14,9 +14,10 @@ IndexesSchema, OtherCoordsSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage -class GeneralHeadBoundary(BoundaryCondition): +class GeneralHeadBoundary(BoundaryCondition,IGridPackage): """ The General-Head Boundary package is used to simulate head-dependent flux boundaries. diff --git a/imod/mf6/ic.py b/imod/mf6/ic.py index 6160ad233..2714227ba 100644 --- a/imod/mf6/ic.py +++ b/imod/mf6/ic.py @@ -8,7 +8,9 @@ from imod.schemata import DTypeSchema, IdentityNoDataSchema, IndexesSchema -class InitialConditions(Package): +from imod.mf6.interfaces.igridpackage import IGridPackage + +class InitialConditions(Package, IGridPackage): """ Initial Conditions (IC) Package information is read from the file that is specified by "IC6" as the file type. Only one IC Package can be specified diff --git a/imod/mf6/interfaces/igridpackage.py b/imod/mf6/interfaces/igridpackage.py new file mode 100644 index 000000000..5fa46e412 --- /dev/null +++ b/imod/mf6/interfaces/igridpackage.py @@ -0,0 +1,11 @@ + +from imod.mf6.interfaces.ipackage import IPackage +import abc +from typing import Optional +from imod.mf6.utilities.regridding_types import RegridderType + +class IGridPackage(IPackage, abc.ABC): + @property + def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: + if hasattr(self, "_regrid_method"): + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index 696f509da..ebc9867c9 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -11,9 +11,9 @@ IdentityNoDataSchema, IndexesSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class ImmobileStorageTransfer(Package): +class ImmobileStorageTransfer(Package, IGridPackage): """ The Immobile Storage and Transfer (IST) package represents an immobile fraction of groundwater. Any number of IST Packages can be specified for a diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index f46a233d6..1e51f7296 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -8,9 +8,9 @@ IdentityNoDataSchema, IndexesSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class MobileStorageTransfer(Package): +class MobileStorageTransfer(Package, IGridPackage): """ Mobile Storage. diff --git a/imod/mf6/npf.py b/imod/mf6/npf.py index 2107f6743..c32a8da03 100644 --- a/imod/mf6/npf.py +++ b/imod/mf6/npf.py @@ -14,7 +14,7 @@ IndexesSchema, ) from imod.typing import GridDataArray - +from imod.mf6.interfaces.igridpackage import IGridPackage def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: if griddataarray is None or griddataarray.values is None: @@ -33,7 +33,7 @@ def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: return griddataarray.values.item() -class NodePropertyFlow(Package): +class NodePropertyFlow(Package, IGridPackage): """ Node Property Flow package. diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 8e8990724..345f3826a 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -23,10 +23,7 @@ PackageBase, ) from imod.mf6.utilities.regrid import ( - RegridderInstancesCollection, RegridderType, - assign_coord_if_present, - get_non_grid_data, ) from imod.mf6.utilities.schemata import filter_schemata_dict from imod.mf6.validation import validation_pkg_error_message @@ -38,7 +35,7 @@ ValidationError, ) from imod.typing import GridDataArray - +from imod.mf6.utilities.regrid import _regrid_like class Package(PackageBase, IPackage, abc.ABC): """ @@ -601,57 +598,7 @@ def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: return self._regrid_method return None - def _regrid_array( - self, - varname: str, - regridder_collection: RegridderInstancesCollection, - regridder_name: str, - regridder_function: str, - target_grid: GridDataArray, - ) -> Optional[GridDataArray]: - """ - Regrids a data_array. The array is specified by its key in the dataset. - Each data-array can represent: - -a scalar value, valid for the whole grid - -an array of a different scalar per layer - -an array with a value per grid block - -None - """ - - # skip regridding for arrays with no valid values (such as "None") - if not self._valid(self.dataset[varname].values[()]): - return None - # the dataarray might be a scalar. If it is, then it does not need regridding. - if is_scalar(self.dataset[varname]): - return self.dataset[varname].values[()] - - if isinstance(self.dataset[varname], xr.DataArray): - coords = self.dataset[varname].coords - # if it is an xr.DataArray it may be layer-based; then no regridding is needed - if not ("x" in coords and "y" in coords): - return self.dataset[varname] - - # if it is an xr.DataArray it needs the dx, dy coordinates for regridding, which are otherwise not mandatory - if not ("dx" in coords and "dy" in coords): - raise ValueError( - f"DataArray {varname} does not have both a dx and dy coordinates" - ) - - # obtain an instance of a regridder for the chosen method - regridder = regridder_collection.get_regridder( - regridder_name, - regridder_function, - ) - - # store original dtype of data - original_dtype = self.dataset[varname].dtype - - # regrid data array - regridded_array = regridder.regrid(self.dataset[varname]) - - # reconvert the result to the same dtype as the original - return regridded_array.astype(original_dtype) def regrid_like( self, @@ -688,48 +635,8 @@ def regrid_like( a package with the same options as this package, and with all the data-arrays regridded to another discretization, similar to the one used in input argument "target_grid" """ - if not hasattr(self, "_regrid_method"): - raise NotImplementedError( - f"Package {type(self).__name__} does not support regridding" - ) - - regridder_collection = RegridderInstancesCollection( - self.dataset, target_grid=target_grid - ) - - regridder_settings = copy.deepcopy(self._regrid_method) - if regridder_types is not None: - regridder_settings.update(regridder_types) - - new_package_data = get_non_grid_data(self, list(regridder_settings.keys())) - - for ( - varname, - regridder_type_and_function, - ) in regridder_settings.items(): - regridder_name, regridder_function = regridder_type_and_function - - # skip variables that are not in this dataset - if varname not in self.dataset.keys(): - continue - - # regrid the variable - new_package_data[varname] = self._regrid_array( - varname, - regridder_collection, - regridder_name, - regridder_function, - target_grid, - ) - # set dx and dy if present in target_grid - new_package_data[varname] = assign_coord_if_present( - "dx", target_grid, new_package_data[varname] - ) - new_package_data[varname] = assign_coord_if_present( - "dy", target_grid, new_package_data[varname] - ) + return _regrid_like(target_grid,regridder_types) - return self.__class__(**new_package_data) def _skip_masking_dataarray(self, array_name: str) -> bool: if hasattr(self, "_skip_mask_arrays"): diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index 3727d5c48..01769c12d 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -14,9 +14,8 @@ IndexesSchema, OtherCoordsSchema, ) - - -class Recharge(BoundaryCondition): +from imod.mf6.interfaces.igridpackage import IGridPackage +class Recharge(BoundaryCondition, IGridPackage): """ Recharge Package. Any number of RCH Packages can be specified for a single groundwater flow diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index 0271b0454..68d25f4e7 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -14,9 +14,9 @@ IndexesSchema, OtherCoordsSchema, ) +from imod.mf6.interfaces.igridpackage import IGridPackage - -class River(BoundaryCondition): +class River(BoundaryCondition, IGridPackage): """ River package. Any number of RIV Packages can be specified for a single groundwater flow diff --git a/imod/mf6/src.py b/imod/mf6/src.py index 195de09ab..2a83031de 100644 --- a/imod/mf6/src.py +++ b/imod/mf6/src.py @@ -11,8 +11,8 @@ OtherCoordsSchema, ) - -class MassSourceLoading(BoundaryCondition): +from imod.mf6.interfaces.igridpackage import IGridPackage +class MassSourceLoading(BoundaryCondition, IGridPackage): """ Mass Source Loading (SRC) package for structured discretization (DIS) models. Any number of SRC Packages can be specified for a single diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 5b453925e..d9158ddf3 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -12,7 +12,7 @@ IdentityNoDataSchema, IndexesSchema, ) - +from imod.mf6.interfaces.igridpackage import IGridPackage class Storage(Package): _pkg_id = "sto_deprecated" @@ -23,7 +23,7 @@ def __init__(*args, **kwargs): ) -class StorageBase(Package, abc.ABC): +class StorageBase(Package, IGridPackage, abc.ABC): def get_options(self, d): # Skip both variables in grid_data and "transient". not_options = list(self._grid_data.keys()) diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 1dbc7c81e..61f998877 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -1,7 +1,6 @@ import abc -from enum import Enum -from typing import Any, Optional, Tuple, Union +from typing import Any, Optional, Union import xarray as xr import xugrid as xu from xugrid.regrid.regridder import BaseRegridder @@ -9,20 +8,12 @@ from imod.typing.grid import GridDataArray from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.interfaces.ilinedatapackage import ILineDataPackage - -class RegridderType(Enum): - """ - Enumerator referring to regridder types in ``xugrid``. - These can be used safely in scripts, remaining backwards compatible for - when it is decided to rename regridders in ``xugrid``. For an explanation - what each regridder type does, we refer to the `xugrid documentation `_ - """ - - CENTROIDLOCATOR = xu.CentroidLocatorRegridder - BARYCENTRIC = xu.BarycentricInterpolator - OVERLAP = xu.OverlapRegridder - RELATIVEOVERLAP = xu.RelativeOverlapRegridder - +from imod.mf6.interfaces.ipointdatapackage import IPointDataPackage +from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.ipackage import IPackage +from imod.mf6.utilities.regridding_types import RegridderType +import copy +from xarray.core.utils import is_scalar class RegridderInstancesCollection: """ @@ -157,4 +148,145 @@ def regrid_like( """ return clip_by_grid(package, target_grid) +@typedispatch # type: ignore[no-redef] +def regrid_like( + package: IPointDataPackage, target_grid: GridDataArray, *_ +) -> IPointDataPackage: + """ + The regrid_like method is irrelevant for this package as it is + grid-agnostic, instead this method clips the package based on the grid + exterior. + """ + target_grid_2d = target_grid.isel(layer=0, drop=True, missing_dims="ignore") + return clip_by_grid(package, target_grid_2d) + +def _regrid_array( + package: IGridPackage, + varname: str, + regridder_collection: RegridderInstancesCollection, + regridder_name: str, + regridder_function: str, + target_grid: GridDataArray, + ) -> Optional[GridDataArray]: + """ + Regrids a data_array. The array is specified by its key in the dataset. + Each data-array can represent: + -a scalar value, valid for the whole grid + -an array of a different scalar per layer + -an array with a value per grid block + -None + """ + + # skip regridding for arrays with no valid values (such as "None") + if not package._valid(package.dataset[varname].values[()]): + return None + + # the dataarray might be a scalar. If it is, then it does not need regridding. + if is_scalar(package.dataset[varname]): + return package.dataset[varname].values[()] + + if isinstance(package.dataset[varname], xr.DataArray): + coords = package.dataset[varname].coords + # if it is an xr.DataArray it may be layer-based; then no regridding is needed + if not ("x" in coords and "y" in coords): + return package.dataset[varname] + + # if it is an xr.DataArray it needs the dx, dy coordinates for regridding, which are otherwise not mandatory + if not ("dx" in coords and "dy" in coords): + raise ValueError( + f"DataArray {varname} does not have both a dx and dy coordinates" + ) + + # obtain an instance of a regridder for the chosen method + regridder = regridder_collection.get_regridder( + regridder_name, + regridder_function, + ) + + # store original dtype of data + original_dtype = package.dataset[varname].dtype + + # regrid data array + regridded_array = regridder.regrid(package.dataset[varname]) + + # reconvert the result to the same dtype as the original + return regridded_array.astype(original_dtype) + +@typedispatch # type: ignore[no-redef] +def _regrid_like( + package: IGridPackage, + target_grid: GridDataArray, + regridder_types: Optional[dict[str, tuple[RegridderType, str]]] = None, +) -> IPackage: + """ + Creates a package of the same type as this package, based on another discretization. + It regrids all the arrays in this package to the desired discretization, and leaves the options + unmodified. At the moment only regridding to a different planar grid is supported, meaning + ``target_grid`` has different ``"x"`` and ``"y"`` or different ``cell2d`` coords. + + The regridding methods can be specified in the _regrid_method attribute of the package. These are the defaults + that specify how each array should be regridded. These defaults can be overridden using the input + parameters of this function. + + Examples + -------- + To regrid the npf package with a non-default method for the k-field, call regrid_like with these arguments: + + >>> new_npf = npf.regrid_like(like, {"k": (imod.RegridderType.OVERLAP, "mean")}) + + + Parameters + ---------- + target_grid: xr.DataArray or xu.UgridDataArray + a grid defined over the same discretization as the one we want to regrid the package to + regridder_types: dict(str->(regridder type,str)) + dictionary mapping arraynames (str) to a tuple of regrid type (a specialization class of BaseRegridder) and function name (str) + this dictionary can be used to override the default mapping method. + + Returns + ------- + a package with the same options as this package, and with all the data-arrays regridded to another discretization, + similar to the one used in input argument "target_grid" + """ + if not hasattr(package, "_regrid_method"): + raise NotImplementedError( + f"Package {type(package).__name__} does not support regridding" + ) + + regridder_collection = RegridderInstancesCollection( + package.dataset, target_grid=target_grid + ) + + regridder_settings = copy.deepcopy(package._regrid_method) + if regridder_types is not None: + regridder_settings.update(regridder_types) + + new_package_data = get_non_grid_data(package, list(regridder_settings.keys())) + + for ( + varname, + regridder_type_and_function, + ) in regridder_settings.items(): + regridder_name, regridder_function = regridder_type_and_function + + # skip variables that are not in this dataset + if varname not in package.dataset.keys(): + continue + + # regrid the variable + new_package_data[varname] = _regrid_array( + varname, + regridder_collection, + regridder_name, + regridder_function, + target_grid, + ) + # set dx and dy if present in target_grid + new_package_data[varname] = assign_coord_if_present( + "dx", target_grid, new_package_data[varname] + ) + new_package_data[varname] = assign_coord_if_present( + "dy", target_grid, new_package_data[varname] + ) + return package.__class__(**new_package_data) diff --git a/imod/mf6/utilities/regridding_types.py b/imod/mf6/utilities/regridding_types.py new file mode 100644 index 000000000..13a72e82a --- /dev/null +++ b/imod/mf6/utilities/regridding_types.py @@ -0,0 +1,15 @@ +from enum import Enum +import xugrid as xu + +class RegridderType(Enum): + """ + Enumerator referring to regridder types in ``xugrid``. + These can be used safely in scripts, remaining backwards compatible for + when it is decided to rename regridders in ``xugrid``. For an explanation + what each regridder type does, we refer to the `xugrid documentation `_ + """ + + CENTROIDLOCATOR = xu.CentroidLocatorRegridder + BARYCENTRIC = xu.BarycentricInterpolator + OVERLAP = xu.OverlapRegridder + RELATIVEOVERLAP = xu.RelativeOverlapRegridder \ No newline at end of file From 4c15f5d0498c3c7d2ae433c7b92684a04fe6eb3b Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 6 Mar 2024 16:49:44 +0100 Subject: [PATCH 03/29] fixed tests until now --- imod/mf6/interfaces/igridpackage.py | 3 ++- imod/mf6/package.py | 2 +- imod/mf6/utilities/regrid.py | 1 + 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/imod/mf6/interfaces/igridpackage.py b/imod/mf6/interfaces/igridpackage.py index 5fa46e412..f3ae5dd57 100644 --- a/imod/mf6/interfaces/igridpackage.py +++ b/imod/mf6/interfaces/igridpackage.py @@ -8,4 +8,5 @@ class IGridPackage(IPackage, abc.ABC): @property def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: if hasattr(self, "_regrid_method"): - return self._regrid_method \ No newline at end of file + return self._regrid_method + return None \ No newline at end of file diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 345f3826a..deba9ca28 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -635,7 +635,7 @@ def regrid_like( a package with the same options as this package, and with all the data-arrays regridded to another discretization, similar to the one used in input argument "target_grid" """ - return _regrid_like(target_grid,regridder_types) + return _regrid_like(self, target_grid,regridder_types) def _skip_masking_dataarray(self, array_name: str) -> bool: diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 61f998877..88acf3e1b 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -275,6 +275,7 @@ def _regrid_like( # regrid the variable new_package_data[varname] = _regrid_array( + package, varname, regridder_collection, regridder_name, From 21c756129cdb7651e953f4e9fb176a3e55198e08 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 6 Mar 2024 17:10:23 +0100 Subject: [PATCH 04/29] work in progress --- imod/mf6/chd.py | 4 ++-- imod/mf6/cnc.py | 4 ++-- imod/mf6/dis.py | 4 ++-- imod/mf6/disv.py | 4 ++-- imod/mf6/drn.py | 4 ++-- imod/mf6/dsp.py | 4 ++-- imod/mf6/evt.py | 4 ++-- imod/mf6/ghb.py | 4 ++-- imod/mf6/ic.py | 4 ++-- imod/mf6/interfaces/{igridpackage.py => iregridpackage.py} | 2 +- imod/mf6/ist.py | 4 ++-- imod/mf6/model.py | 4 ++-- imod/mf6/model_gwf.py | 4 ++-- imod/mf6/mst.py | 4 ++-- imod/mf6/npf.py | 4 ++-- imod/mf6/rch.py | 4 ++-- imod/mf6/riv.py | 4 ++-- imod/mf6/src.py | 4 ++-- imod/mf6/sto.py | 4 ++-- imod/mf6/utilities/regrid.py | 6 +++--- 20 files changed, 40 insertions(+), 40 deletions(-) rename imod/mf6/interfaces/{igridpackage.py => iregridpackage.py} (89%) diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py index 583f9dcdc..a56cd625d 100644 --- a/imod/mf6/chd.py +++ b/imod/mf6/chd.py @@ -13,10 +13,10 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class ConstantHead(BoundaryCondition, IGridPackage): +class ConstantHead(BoundaryCondition, IRegridPackage): """ Constant-Head package. Any number of CHD Packages can be specified for a single groundwater flow model; however, an error will occur if a CHD Package diff --git a/imod/mf6/cnc.py b/imod/mf6/cnc.py index b27cc19c5..8a58e407f 100644 --- a/imod/mf6/cnc.py +++ b/imod/mf6/cnc.py @@ -11,9 +11,9 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class ConstantConcentration(BoundaryCondition, IGridPackage): +class ConstantConcentration(BoundaryCondition, IRegridPackage): """ Constant Concentration package. diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index 0aa69549a..8d18f01e2 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -16,9 +16,9 @@ IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class StructuredDiscretization(Package, IGridPackage): +class StructuredDiscretization(Package, IRegridPackage): """ Discretization information for structered grids is specified using the file. (DIS6) Only one discretization input file (DISU6, DISV6 or DIS6) can be diff --git a/imod/mf6/disv.py b/imod/mf6/disv.py index a48c28bfc..eb64cf226 100644 --- a/imod/mf6/disv.py +++ b/imod/mf6/disv.py @@ -14,8 +14,8 @@ IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage -class VerticesDiscretization(Package, IGridPackage): +from imod.mf6.interfaces.iregridpackage import IRegridPackage +class VerticesDiscretization(Package, IRegridPackage): """ Discretization by Vertices (DISV). diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index 3aa506ceb..6a033c731 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -14,9 +14,9 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class Drainage(BoundaryCondition, IGridPackage): +class Drainage(BoundaryCondition, IRegridPackage): """ The Drain package is used to simulate head-dependent flux boundaries. https://water.usgs.gov/ogw/modflow/mf6io.pdf#page=67 diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index a224ef45c..28b45e90e 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -9,9 +9,9 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class Dispersion(Package, IGridPackage): +class Dispersion(Package, IRegridPackage): """ Molecular Diffusion and Dispersion. diff --git a/imod/mf6/evt.py b/imod/mf6/evt.py index 110e6685d..355e85ccb 100644 --- a/imod/mf6/evt.py +++ b/imod/mf6/evt.py @@ -31,9 +31,9 @@ | DimsSchema("segment", "{face_dim}") ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class Evapotranspiration(BoundaryCondition, IGridPackage): +class Evapotranspiration(BoundaryCondition, IRegridPackage): """ Evapotranspiration (EVT) Package. Any number of EVT Packages can be specified for a single groundwater flow diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index d7f4657b3..4c51db0a7 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -14,10 +14,10 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class GeneralHeadBoundary(BoundaryCondition,IGridPackage): +class GeneralHeadBoundary(BoundaryCondition,IRegridPackage): """ The General-Head Boundary package is used to simulate head-dependent flux boundaries. diff --git a/imod/mf6/ic.py b/imod/mf6/ic.py index 2714227ba..7eefa1eb8 100644 --- a/imod/mf6/ic.py +++ b/imod/mf6/ic.py @@ -8,9 +8,9 @@ from imod.schemata import DTypeSchema, IdentityNoDataSchema, IndexesSchema -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class InitialConditions(Package, IGridPackage): +class InitialConditions(Package, IRegridPackage): """ Initial Conditions (IC) Package information is read from the file that is specified by "IC6" as the file type. Only one IC Package can be specified diff --git a/imod/mf6/interfaces/igridpackage.py b/imod/mf6/interfaces/iregridpackage.py similarity index 89% rename from imod/mf6/interfaces/igridpackage.py rename to imod/mf6/interfaces/iregridpackage.py index f3ae5dd57..8c0550186 100644 --- a/imod/mf6/interfaces/igridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -4,7 +4,7 @@ from typing import Optional from imod.mf6.utilities.regridding_types import RegridderType -class IGridPackage(IPackage, abc.ABC): +class IRegridPackage(IPackage, abc.ABC): @property def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: if hasattr(self, "_regrid_method"): diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index ebc9867c9..4dc6e4994 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -11,9 +11,9 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class ImmobileStorageTransfer(Package, IGridPackage): +class ImmobileStorageTransfer(Package, IRegridPackage): """ The Immobile Storage and Transfer (IST) package represents an immobile fraction of groundwater. Any number of IST Packages can be specified for a diff --git a/imod/mf6/model.py b/imod/mf6/model.py index ab9b79172..35a1e9920 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -27,7 +27,7 @@ from imod.schemata import ValidationError from imod.typing import GridDataArray from imod.typing.grid import ones_like - +from imod.mf6.interfaces.iregridpackage import IRegridPackage class Modflow6Model(collections.UserDict, abc.ABC): _mandatory_packages: Tuple[str, ...] = () @@ -479,7 +479,7 @@ def regrid_like( new_model = self.__class__() for pkg_name, pkg in self.items(): - if pkg.is_regridding_supported(): + if isinstance(pkg, IRegridPackage): new_model[pkg_name] = pkg.regrid_like(target_grid) else: raise NotImplementedError( diff --git a/imod/mf6/model_gwf.py b/imod/mf6/model_gwf.py index bdb49f934..4023e0f71 100644 --- a/imod/mf6/model_gwf.py +++ b/imod/mf6/model_gwf.py @@ -11,7 +11,7 @@ from imod.mf6.model import Modflow6Model from imod.mf6.utilities.regrid import RegridderType from imod.typing import GridDataArray - +from imod.mf6.interfaces.iregridpackage import IRegridPackage class GroundwaterFlowModel(Modflow6Model): _mandatory_packages = ("npf", "ic", "oc", "sto") @@ -44,7 +44,7 @@ def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: """ methods: defaultdict = defaultdict(list) for pkg_name, pkg in self.items(): - if pkg.is_regridding_supported(): + if isinstance(pkg, IRegridPackage): pkg_methods = pkg.get_regrid_methods() for variable in pkg_methods: if ( diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index 1e51f7296..55f5f7c97 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -8,9 +8,9 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class MobileStorageTransfer(Package, IGridPackage): +class MobileStorageTransfer(Package, IRegridPackage): """ Mobile Storage. diff --git a/imod/mf6/npf.py b/imod/mf6/npf.py index c32a8da03..ec8bbc089 100644 --- a/imod/mf6/npf.py +++ b/imod/mf6/npf.py @@ -14,7 +14,7 @@ IndexesSchema, ) from imod.typing import GridDataArray -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: if griddataarray is None or griddataarray.values is None: @@ -33,7 +33,7 @@ def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: return griddataarray.values.item() -class NodePropertyFlow(Package, IGridPackage): +class NodePropertyFlow(Package, IRegridPackage): """ Node Property Flow package. diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index 01769c12d..fefb28680 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -14,8 +14,8 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage -class Recharge(BoundaryCondition, IGridPackage): +from imod.mf6.interfaces.iregridpackage import IRegridPackage +class Recharge(BoundaryCondition, IRegridPackage): """ Recharge Package. Any number of RCH Packages can be specified for a single groundwater flow diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index 68d25f4e7..7ed045397 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -14,9 +14,9 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage -class River(BoundaryCondition, IGridPackage): +class River(BoundaryCondition, IRegridPackage): """ River package. Any number of RIV Packages can be specified for a single groundwater flow diff --git a/imod/mf6/src.py b/imod/mf6/src.py index 2a83031de..34d89b54b 100644 --- a/imod/mf6/src.py +++ b/imod/mf6/src.py @@ -11,8 +11,8 @@ OtherCoordsSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage -class MassSourceLoading(BoundaryCondition, IGridPackage): +from imod.mf6.interfaces.iregridpackage import IRegridPackage +class MassSourceLoading(BoundaryCondition, IRegridPackage): """ Mass Source Loading (SRC) package for structured discretization (DIS) models. Any number of SRC Packages can be specified for a single diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index d9158ddf3..2308e617c 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -12,7 +12,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage class Storage(Package): _pkg_id = "sto_deprecated" @@ -23,7 +23,7 @@ def __init__(*args, **kwargs): ) -class StorageBase(Package, IGridPackage, abc.ABC): +class StorageBase(Package, IRegridPackage, abc.ABC): def get_options(self, d): # Skip both variables in grid_data and "transient". not_options = list(self._grid_data.keys()) diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 88acf3e1b..8cfaf2917 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -9,7 +9,7 @@ from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.interfaces.ilinedatapackage import ILineDataPackage from imod.mf6.interfaces.ipointdatapackage import IPointDataPackage -from imod.mf6.interfaces.igridpackage import IGridPackage +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.interfaces.ipackage import IPackage from imod.mf6.utilities.regridding_types import RegridderType import copy @@ -161,7 +161,7 @@ def regrid_like( return clip_by_grid(package, target_grid_2d) def _regrid_array( - package: IGridPackage, + package: IRegridPackage, varname: str, regridder_collection: RegridderInstancesCollection, regridder_name: str, @@ -214,7 +214,7 @@ def _regrid_array( @typedispatch # type: ignore[no-redef] def _regrid_like( - package: IGridPackage, + package: IRegridPackage, target_grid: GridDataArray, regridder_types: Optional[dict[str, tuple[RegridderType, str]]] = None, ) -> IPackage: From 10b43210243b11566d9c61e1c3373a623e1cb9a2 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 10:57:07 +0100 Subject: [PATCH 05/29] fixed tests --- imod/mf6/model.py | 2 +- imod/mf6/model_gwf.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 35a1e9920..5b31b4f81 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -479,7 +479,7 @@ def regrid_like( new_model = self.__class__() for pkg_name, pkg in self.items(): - if isinstance(pkg, IRegridPackage): + if pkg.is_regridding_supported(): new_model[pkg_name] = pkg.regrid_like(target_grid) else: raise NotImplementedError( diff --git a/imod/mf6/model_gwf.py b/imod/mf6/model_gwf.py index 4023e0f71..7f8af0065 100644 --- a/imod/mf6/model_gwf.py +++ b/imod/mf6/model_gwf.py @@ -44,7 +44,7 @@ def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: """ methods: defaultdict = defaultdict(list) for pkg_name, pkg in self.items(): - if isinstance(pkg, IRegridPackage): + if pkg.is_regridding_supported(): pkg_methods = pkg.get_regrid_methods() for variable in pkg_methods: if ( From c40e1c6ace1c840eeb7287ea490f138b6fcee1dc Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 11:59:43 +0100 Subject: [PATCH 06/29] starting on model regridding --- imod/mf6/interfaces/imodel.py | 8 ++++++++ imod/mf6/interfaces/ipackage.py | 5 +---- imod/mf6/model.py | 10 ++++------ imod/mf6/utilities/regrid.py | 25 +++++++++++++++++++++++++ 4 files changed, 38 insertions(+), 10 deletions(-) create mode 100644 imod/mf6/interfaces/imodel.py diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py new file mode 100644 index 000000000..6653fecdb --- /dev/null +++ b/imod/mf6/interfaces/imodel.py @@ -0,0 +1,8 @@ +from abc import abstractmethod +from typing import Optional + +from imod.typing import GridDataArray + + +class IModel: + pass \ No newline at end of file diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py index 9b6154d99..018e67b31 100644 --- a/imod/mf6/interfaces/ipackage.py +++ b/imod/mf6/interfaces/ipackage.py @@ -8,7 +8,4 @@ class IPackage(IPackageBase, metaclass=abc.ABCMeta): The base methods and attributes available in all packages """ - @property - @abc.abstractmethod - def auxiliary_data_fields(self) -> dict[str, str]: - raise NotImplementedError + pass \ No newline at end of file diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 5b31b4f81..db9cc3a5a 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -27,9 +27,10 @@ from imod.schemata import ValidationError from imod.typing import GridDataArray from imod.typing.grid import ones_like -from imod.mf6.interfaces.iregridpackage import IRegridPackage +from imod.mf6.interfaces.imodel import IModel +from imod.mf6.utilities.regrid import _get_unique_regridder_types -class Modflow6Model(collections.UserDict, abc.ABC): +class Modflow6Model(collections.UserDict, IModel, abc.ABC): _mandatory_packages: Tuple[str, ...] = () _model_id: Optional[str] = None _template: Template @@ -486,7 +487,7 @@ def regrid_like( f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" ) - methods = self._get_unique_regridder_types() + methods = _get_unique_regridder_types(self) output_domain = self._get_regridding_domain(target_grid, methods) new_model.mask_all_packages(output_domain) new_model.purge_empty_packages() @@ -585,6 +586,3 @@ def __repr__(self) -> str: def is_use_newton(self): return False - - def _get_unique_regridder_types(self): - raise NotImplementedError(f"Regridding not supported for {self}") diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 8cfaf2917..fc66920a3 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -14,6 +14,8 @@ from imod.mf6.utilities.regridding_types import RegridderType import copy from xarray.core.utils import is_scalar +from imod.mf6.interfaces.imodel import IModel +from collections import defaultdict class RegridderInstancesCollection: """ @@ -291,3 +293,26 @@ def _regrid_like( ) return package.__class__(**new_package_data) + +def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: + """ + This function loops over the packages and collects all regridder-types that are in use. + """ + methods: defaultdict = defaultdict(list) + for pkg_name, pkg in self.items(): + if pkg.is_regridding_supported(): + pkg_methods = pkg.get_regrid_methods() + for variable in pkg_methods: + if ( + variable in pkg.dataset.data_vars + and pkg.dataset[variable].values[()] is not None + ): + regriddertype = pkg_methods[variable][0] + functiontype = pkg_methods[variable][1] + if functiontype not in methods[regriddertype]: + methods[regriddertype].append(functiontype) + else: + raise NotImplementedError( + f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" + ) + return methods \ No newline at end of file From 749180d79e5c0815893d3415d09a2688c7b45a39 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 11:59:51 +0100 Subject: [PATCH 07/29] starting on model regridding --- imod/mf6/model_gwf.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/imod/mf6/model_gwf.py b/imod/mf6/model_gwf.py index 7f8af0065..1d502595d 100644 --- a/imod/mf6/model_gwf.py +++ b/imod/mf6/model_gwf.py @@ -37,30 +37,6 @@ def __init__( "under_relaxation": under_relaxation, } - - def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: - """ - This function loops over the packages and collects all regridder-types that are in use. - """ - methods: defaultdict = defaultdict(list) - for pkg_name, pkg in self.items(): - if pkg.is_regridding_supported(): - pkg_methods = pkg.get_regrid_methods() - for variable in pkg_methods: - if ( - variable in pkg.dataset.data_vars - and pkg.dataset[variable].values[()] is not None - ): - regriddertype = pkg_methods[variable][0] - functiontype = pkg_methods[variable][1] - if functiontype not in methods[regriddertype]: - methods[regriddertype].append(functiontype) - else: - raise NotImplementedError( - f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" - ) - return methods - def clip_box( self, time_min: Optional[cftime.datetime | np.datetime64 | str] = None, From 0abe7ea8b0f8c3010135b537a3320fcd26534c08 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 15:37:29 +0100 Subject: [PATCH 08/29] tests and examples pass --- imod/mf6/model.py | 23 ++------------- imod/mf6/utilities/regrid.py | 57 ++++++++++++++++++++++++++++++++++-- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 3671b9ad5..67f094f93 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -29,6 +29,7 @@ from imod.typing.grid import ones_like from imod.mf6.interfaces.imodel import IModel from imod.mf6.utilities.regrid import _get_unique_regridder_types +from imod.mf6.utilities.regrid import _regrid_like class Modflow6Model(collections.UserDict, IModel, abc.ABC): _mandatory_packages: Tuple[str, ...] = () @@ -477,27 +478,9 @@ def regrid_like( a model with similar packages to the input model, and with all the data-arrays regridded to another discretization, similar to the one used in input argument "target_grid" """ - new_model = self.__class__() - - for pkg_name, pkg in self.items(): - if pkg.is_regridding_supported(): - new_model[pkg_name] = pkg.regrid_like(target_grid) - else: - raise NotImplementedError( - f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" - ) - - methods = _get_unique_regridder_types(self) - output_domain = self._get_regridding_domain(target_grid, methods) - new_model.mask_all_packages(output_domain) - new_model.purge_empty_packages() - if validate: - status_info = NestedStatusInfo("Model validation status") - status_info.add(new_model._validate("Regridded model")) - if status_info.has_errors(): - raise ValidationError("\n" + status_info.to_string()) - return new_model + return _regrid_like(self, target_grid, validate) + def mask_all_packages( self, mask: GridDataArray, diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index fc66920a3..0b6dbd3d3 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -16,6 +16,9 @@ from xarray.core.utils import is_scalar from imod.mf6.interfaces.imodel import IModel from collections import defaultdict +from imod.mf6.statusinfo import NestedStatusInfo +from imod.schemata import ValidationError +from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables, remove_expanded_auxiliary_variables_from_dataset class RegridderInstancesCollection: """ @@ -254,6 +257,9 @@ def _regrid_like( raise NotImplementedError( f"Package {type(package).__name__} does not support regridding" ) + + if hasattr(package,"auxiliary_data_fields"): + remove_expanded_auxiliary_variables_from_dataset(package) regridder_collection = RegridderInstancesCollection( package.dataset, target_grid=target_grid @@ -291,15 +297,17 @@ def _regrid_like( new_package_data[varname] = assign_coord_if_present( "dy", target_grid, new_package_data[varname] ) + if hasattr(package,"auxiliary_data_fields"): + expand_transient_auxiliary_variables(package) return package.__class__(**new_package_data) -def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: +def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, list[str]]: """ This function loops over the packages and collects all regridder-types that are in use. """ methods: defaultdict = defaultdict(list) - for pkg_name, pkg in self.items(): + for pkg_name, pkg in model.items(): if pkg.is_regridding_supported(): pkg_methods = pkg.get_regrid_methods() for variable in pkg_methods: @@ -315,4 +323,47 @@ def _get_unique_regridder_types(self) -> defaultdict[RegridderType, list[str]]: raise NotImplementedError( f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" ) - return methods \ No newline at end of file + return methods + +@typedispatch # type: ignore[no-redef] +def _regrid_like( + model: IModel, target_grid: GridDataArray, validate: bool = True +) -> IModel: + """ + Creates a model by regridding the packages of this model to another discretization. + It regrids all the arrays in the package using the default regridding methods. + At the moment only regridding to a different planar grid is supported, meaning + ``target_grid`` has different ``"x"`` and ``"y"`` or different ``cell2d`` coords. + + Parameters + ---------- + target_grid: xr.DataArray or xu.UgridDataArray + a grid defined over the same discretization as the one we want to regrid the package to + validate: bool + set to true to validate the regridded packages + + Returns + ------- + a model with similar packages to the input model, and with all the data-arrays regridded to another discretization, + similar to the one used in input argument "target_grid" + """ + new_model = model.__class__() + + for pkg_name, pkg in model.items(): + if pkg.is_regridding_supported(): + new_model[pkg_name] = pkg.regrid_like(target_grid) + else: + raise NotImplementedError( + f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" + ) + + methods = _get_unique_regridder_types(model) + output_domain = model._get_regridding_domain(target_grid, methods) + new_model.mask_all_packages(output_domain) + new_model.purge_empty_packages() + if validate: + status_info = NestedStatusInfo("Model validation status") + status_info.add(new_model._validate("Regridded model")) + if status_info.has_errors(): + raise ValidationError("\n" + status_info.to_string()) + return new_model From 1085f78cc26f33f0887eab2bc2c376f34829cf0b Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 15:54:49 +0100 Subject: [PATCH 09/29] mypy and ruff fixes --- imod/mf6/chd.py | 2 +- imod/mf6/cnc.py | 3 ++- imod/mf6/dis.py | 2 +- imod/mf6/disv.py | 3 ++- imod/mf6/drn.py | 3 ++- imod/mf6/dsp.py | 3 ++- imod/mf6/evt.py | 1 + imod/mf6/ghb.py | 2 +- imod/mf6/ic.py | 3 +-- imod/mf6/interfaces/imodel.py | 3 --- imod/mf6/interfaces/iregridpackage.py | 4 +++- imod/mf6/ist.py | 3 ++- imod/mf6/model.py | 17 ++++++++------- imod/mf6/model_gwf.py | 2 +- imod/mf6/mst.py | 3 ++- imod/mf6/npf.py | 3 ++- imod/mf6/package.py | 4 ++-- imod/mf6/rch.py | 4 +++- imod/mf6/riv.py | 3 ++- imod/mf6/src.py | 3 ++- imod/mf6/sto.py | 3 ++- imod/mf6/utilities/regrid.py | 29 +++++++++++++++----------- imod/mf6/utilities/regridding_types.py | 2 ++ 23 files changed, 63 insertions(+), 42 deletions(-) diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py index a56cd625d..286121655 100644 --- a/imod/mf6/chd.py +++ b/imod/mf6/chd.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.schemata import ( @@ -13,7 +14,6 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage class ConstantHead(BoundaryCondition, IRegridPackage): diff --git a/imod/mf6/cnc.py b/imod/mf6/cnc.py index 8a58e407f..138d66ef2 100644 --- a/imod/mf6/cnc.py +++ b/imod/mf6/cnc.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA from imod.schemata import ( AllInsideNoDataSchema, @@ -11,7 +12,7 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class ConstantConcentration(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index 8d18f01e2..050d24612 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -3,6 +3,7 @@ import numpy as np import imod +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import DisBottomSchema @@ -16,7 +17,6 @@ IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage class StructuredDiscretization(Package, IRegridPackage): """ diff --git a/imod/mf6/disv.py b/imod/mf6/disv.py index eb64cf226..6307cdfb5 100644 --- a/imod/mf6/disv.py +++ b/imod/mf6/disv.py @@ -1,6 +1,7 @@ import numpy as np import pandas as pd +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import DisBottomSchema @@ -14,7 +15,7 @@ IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class VerticesDiscretization(Package, IRegridPackage): """ Discretization by Vertices (DISV). diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index 6a033c731..4858fa6f6 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.schemata import ( @@ -14,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class Drainage(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index 69f145fa3..eed8cfa6b 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -1,5 +1,6 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA @@ -10,7 +11,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class Dispersion(Package, IRegridPackage): """ diff --git a/imod/mf6/evt.py b/imod/mf6/evt.py index 355e85ccb..63ab9e4a3 100644 --- a/imod/mf6/evt.py +++ b/imod/mf6/evt.py @@ -33,6 +33,7 @@ from imod.mf6.interfaces.iregridpackage import IRegridPackage + class Evapotranspiration(BoundaryCondition, IRegridPackage): """ Evapotranspiration (EVT) Package. diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index 4c51db0a7..323216452 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.schemata import ( @@ -14,7 +15,6 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage class GeneralHeadBoundary(BoundaryCondition,IRegridPackage): diff --git a/imod/mf6/ic.py b/imod/mf6/ic.py index 7eefa1eb8..cbac5b306 100644 --- a/imod/mf6/ic.py +++ b/imod/mf6/ic.py @@ -2,14 +2,13 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import DTypeSchema, IdentityNoDataSchema, IndexesSchema -from imod.mf6.interfaces.iregridpackage import IRegridPackage - class InitialConditions(Package, IRegridPackage): """ Initial Conditions (IC) Package information is read from the file that is diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py index 6653fecdb..2320d702e 100644 --- a/imod/mf6/interfaces/imodel.py +++ b/imod/mf6/interfaces/imodel.py @@ -1,7 +1,4 @@ -from abc import abstractmethod -from typing import Optional -from imod.typing import GridDataArray class IModel: diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index 8c0550186..8bfd69891 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -1,9 +1,11 @@ -from imod.mf6.interfaces.ipackage import IPackage import abc from typing import Optional + +from imod.mf6.interfaces.ipackage import IPackage from imod.mf6.utilities.regridding_types import RegridderType + class IRegridPackage(IPackage, abc.ABC): @property def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index 4dc6e4994..70c6c6490 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -3,6 +3,7 @@ import numpy as np import xarray as xr +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import ( @@ -11,7 +12,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class ImmobileStorageTransfer(Package, IRegridPackage): """ diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 67f094f93..bb12b4a29 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -7,7 +7,7 @@ from collections import defaultdict from copy import deepcopy from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Optional, Union import cftime import jinja2 @@ -19,20 +19,23 @@ from jinja2 import Template import imod +from imod.mf6.interfaces.imodel import IModel from imod.mf6.package import Package from imod.mf6.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase -from imod.mf6.utilities.regrid import RegridderInstancesCollection, RegridderType +from imod.mf6.utilities.regrid import ( + RegridderInstancesCollection, + RegridderType, + _regrid_like, +) from imod.mf6.validation import pkg_errors_to_status_info from imod.mf6.write_context import WriteContext from imod.schemata import ValidationError from imod.typing import GridDataArray from imod.typing.grid import ones_like -from imod.mf6.interfaces.imodel import IModel -from imod.mf6.utilities.regrid import _get_unique_regridder_types -from imod.mf6.utilities.regrid import _regrid_like + class Modflow6Model(collections.UserDict, IModel, abc.ABC): - _mandatory_packages: Tuple[str, ...] = () + _mandatory_packages: tuple[str, ...] = () _model_id: Optional[str] = None _template: Template @@ -164,7 +167,7 @@ def _model_checks(self, modelkey: str): def __get_domain_geometry( self, - ) -> Tuple[ + ) -> tuple[ Union[xr.DataArray, xu.UgridDataArray], Union[xr.DataArray, xu.UgridDataArray], Union[xr.DataArray, xu.UgridDataArray], diff --git a/imod/mf6/model_gwf.py b/imod/mf6/model_gwf.py index 62130f443..81e4163cf 100644 --- a/imod/mf6/model_gwf.py +++ b/imod/mf6/model_gwf.py @@ -9,7 +9,7 @@ from imod.mf6.clipped_boundary_condition_creator import create_clipped_boundary from imod.mf6.model import Modflow6Model from imod.typing import GridDataArray -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class GroundwaterFlowModel(Modflow6Model): _mandatory_packages = ("npf", "ic", "oc", "sto") diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index 994f58486..c550fd706 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -1,5 +1,6 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA @@ -9,7 +10,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class MobileStorageTransfer(Package, IRegridPackage): """ diff --git a/imod/mf6/npf.py b/imod/mf6/npf.py index ec8bbc089..4981f4d1b 100644 --- a/imod/mf6/npf.py +++ b/imod/mf6/npf.py @@ -2,6 +2,7 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA @@ -14,7 +15,7 @@ IndexesSchema, ) from imod.typing import GridDataArray -from imod.mf6.interfaces.iregridpackage import IRegridPackage + def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: if griddataarray is None or griddataarray.values is None: diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 4c1957994..b623a9c73 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -1,7 +1,6 @@ from __future__ import annotations import abc -import copy import numbers import pathlib from collections import defaultdict @@ -28,6 +27,7 @@ ) from imod.mf6.utilities.regrid import ( RegridderType, + _regrid_like, ) from imod.mf6.utilities.schemata import filter_schemata_dict from imod.mf6.validation import validation_pkg_error_message @@ -39,7 +39,7 @@ ValidationError, ) from imod.typing import GridDataArray -from imod.mf6.utilities.regrid import _regrid_like + class Package(PackageBase, IPackage, abc.ABC): """ diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index 18ef6e7c9..c578a4501 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.schemata import ( @@ -14,7 +15,8 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + + class Recharge(BoundaryCondition, IRegridPackage): """ Recharge Package. diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index 7ed045397..c0211b0c5 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.schemata import ( @@ -14,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class River(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/src.py b/imod/mf6/src.py index 34d89b54b..b27ae51c3 100644 --- a/imod/mf6/src.py +++ b/imod/mf6/src.py @@ -1,6 +1,7 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA from imod.schemata import ( @@ -11,7 +12,7 @@ OtherCoordsSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class MassSourceLoading(BoundaryCondition, IRegridPackage): """ Mass Source Loading (SRC) package for structured discretization (DIS) diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 2308e617c..9895d4d8d 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -2,6 +2,7 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA @@ -12,7 +13,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class Storage(Package): _pkg_id = "sto_deprecated" diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 0b6dbd3d3..158d1c845 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -1,24 +1,29 @@ import abc - +import copy +from collections import defaultdict from typing import Any, Optional, Union + import xarray as xr import xugrid as xu -from xugrid.regrid.regridder import BaseRegridder from fastcore.dispatch import typedispatch -from imod.typing.grid import GridDataArray -from imod.mf6.utilities.clip import clip_by_grid +from xarray.core.utils import is_scalar +from xugrid.regrid.regridder import BaseRegridder + +from imod.mf6.auxiliary_variables import ( + expand_transient_auxiliary_variables, + remove_expanded_auxiliary_variables_from_dataset, +) from imod.mf6.interfaces.ilinedatapackage import ILineDataPackage +from imod.mf6.interfaces.imodel import IModel +from imod.mf6.interfaces.ipackage import IPackage from imod.mf6.interfaces.ipointdatapackage import IPointDataPackage from imod.mf6.interfaces.iregridpackage import IRegridPackage -from imod.mf6.interfaces.ipackage import IPackage -from imod.mf6.utilities.regridding_types import RegridderType -import copy -from xarray.core.utils import is_scalar -from imod.mf6.interfaces.imodel import IModel -from collections import defaultdict from imod.mf6.statusinfo import NestedStatusInfo +from imod.mf6.utilities.clip import clip_by_grid +from imod.mf6.utilities.regridding_types import RegridderType from imod.schemata import ValidationError -from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables, remove_expanded_auxiliary_variables_from_dataset +from imod.typing.grid import GridDataArray + class RegridderInstancesCollection: """ @@ -33,7 +38,7 @@ def __init__( target_grid: Union[xr.DataArray, xu.UgridDataArray], ) -> None: self.regridder_instances: dict[ - Tuple[type[BaseRegridder], Optional[str]], BaseRegridder + tuple[type[BaseRegridder], Optional[str]], BaseRegridder ] = {} self._source_grid = source_grid self._target_grid = target_grid diff --git a/imod/mf6/utilities/regridding_types.py b/imod/mf6/utilities/regridding_types.py index 13a72e82a..6e598e88a 100644 --- a/imod/mf6/utilities/regridding_types.py +++ b/imod/mf6/utilities/regridding_types.py @@ -1,6 +1,8 @@ from enum import Enum + import xugrid as xu + class RegridderType(Enum): """ Enumerator referring to regridder types in ``xugrid``. From c2b2ada874980dce880f3f619cbce09d9ceb1f53 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 15:54:57 +0100 Subject: [PATCH 10/29] mypy and ruff fixes --- imod/mf6/interfaces/ipackage.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py index 018e67b31..8b41ef517 100644 --- a/imod/mf6/interfaces/ipackage.py +++ b/imod/mf6/interfaces/ipackage.py @@ -8,4 +8,7 @@ class IPackage(IPackageBase, metaclass=abc.ABCMeta): The base methods and attributes available in all packages """ - pass \ No newline at end of file + @property + @abc.abstractmethod + def auxiliary_data_fields(self) -> dict[str, str]: + raise NotImplementedError \ No newline at end of file From 48d09e07e0d0ecf319a54011f84abf1b4cb94ce1 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 16:13:53 +0100 Subject: [PATCH 11/29] moved _get_regridding_domain --- imod/mf6/model.py | 26 -------------------------- imod/mf6/utilities/regrid.py | 30 ++++++++++++++++++++++++++++-- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index bb12b4a29..64c7daffe 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -527,32 +527,6 @@ def bottom(self): dis = self._get_diskey() return self[dis]["bottom"] - def _get_regridding_domain( - self, - target_grid: GridDataArray, - methods: defaultdict[RegridderType, list[str]], - ) -> GridDataArray: - """ - This method computes the output-domain for a regridding operation by regridding idomain with - all regridders. Each regridder may leave some cells inactive. The output domain for the model consists of those - cells that all regridders consider active. - """ - idomain = self.domain - regridder_collection = RegridderInstancesCollection( - idomain, target_grid=target_grid - ) - included_in_all = ones_like(target_grid) - regridders =[regridder_collection.get_regridder( regriddertype,function) for regriddertype, functionlist in methods.items() for function in functionlist] - for regridder in regridders: - regridded_idomain = regridder.regrid(idomain) - included_in_all = included_in_all.where(regridded_idomain.notnull()) - included_in_all = regridded_idomain.where(regridded_idomain <=0, other = included_in_all) - - new_idomain = included_in_all.where(included_in_all.notnull(), other=0) - new_idomain = new_idomain.astype(int) - - return new_idomain - def __repr__(self) -> str: INDENT = " " typename = type(self).__name__ diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 158d1c845..071440848 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -22,7 +22,7 @@ from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.utilities.regridding_types import RegridderType from imod.schemata import ValidationError -from imod.typing.grid import GridDataArray +from imod.typing.grid import GridDataArray, ones_like class RegridderInstancesCollection: @@ -363,7 +363,7 @@ def _regrid_like( ) methods = _get_unique_regridder_types(model) - output_domain = model._get_regridding_domain(target_grid, methods) + output_domain = _get_regridding_domain(model, target_grid, methods) new_model.mask_all_packages(output_domain) new_model.purge_empty_packages() if validate: @@ -372,3 +372,29 @@ def _regrid_like( if status_info.has_errors(): raise ValidationError("\n" + status_info.to_string()) return new_model + +def _get_regridding_domain( + model: IModel, + target_grid: GridDataArray, + methods: defaultdict[RegridderType, list[str]], +) -> GridDataArray: + """ + This method computes the output-domain for a regridding operation by regridding idomain with + all regridders. Each regridder may leave some cells inactive. The output domain for the model consists of those + cells that all regridders consider active. + """ + idomain = model.domain + regridder_collection = RegridderInstancesCollection( + idomain, target_grid=target_grid + ) + included_in_all = ones_like(target_grid) + regridders =[regridder_collection.get_regridder( regriddertype,function) for regriddertype, functionlist in methods.items() for function in functionlist] + for regridder in regridders: + regridded_idomain = regridder.regrid(idomain) + included_in_all = included_in_all.where(regridded_idomain.notnull()) + included_in_all = regridded_idomain.where(regridded_idomain <=0, other = included_in_all) + + new_idomain = included_in_all.where(included_in_all.notnull(), other=0) + new_idomain = new_idomain.astype(int) + + return new_idomain \ No newline at end of file From 99c737e7caa12fd8c39f5daa26f593e287e03dd1 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 16:14:03 +0100 Subject: [PATCH 12/29] mypy and ruff fixes --- imod/mf6/model.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 64c7daffe..a69eb407e 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -4,7 +4,6 @@ import collections import inspect import pathlib -from collections import defaultdict from copy import deepcopy from pathlib import Path from typing import Optional, Union @@ -23,15 +22,12 @@ from imod.mf6.package import Package from imod.mf6.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase from imod.mf6.utilities.regrid import ( - RegridderInstancesCollection, - RegridderType, _regrid_like, ) from imod.mf6.validation import pkg_errors_to_status_info from imod.mf6.write_context import WriteContext from imod.schemata import ValidationError from imod.typing import GridDataArray -from imod.typing.grid import ones_like class Modflow6Model(collections.UserDict, IModel, abc.ABC): From 5c75caf86d2af75e895751850810d7214804a378 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 16:20:24 +0100 Subject: [PATCH 13/29] removed type ignores --- imod/mf6/utilities/regrid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 071440848..aeda38351 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -222,7 +222,7 @@ def _regrid_array( # reconvert the result to the same dtype as the original return regridded_array.astype(original_dtype) -@typedispatch # type: ignore[no-redef] +@typedispatch def _regrid_like( package: IRegridPackage, target_grid: GridDataArray, @@ -330,7 +330,7 @@ def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, lis ) return methods -@typedispatch # type: ignore[no-redef] +@typedispatch def _regrid_like( model: IModel, target_grid: GridDataArray, validate: bool = True ) -> IModel: From 905a9f080136ea78640a645f9fdc4ef0730d4d4d Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 17:27:35 +0100 Subject: [PATCH 14/29] cleanup --- imod/mf6/interfaces/imodel.py | 23 +++++++++++++++++++++-- imod/mf6/interfaces/ipackage.py | 15 ++++++++++++++- imod/mf6/interfaces/ipointdatapackage.py | 8 ++++---- imod/mf6/interfaces/iregridpackage.py | 2 ++ imod/mf6/package.py | 21 +++++++++++++++++++++ imod/mf6/utilities/regrid.py | 21 --------------------- 6 files changed, 62 insertions(+), 28 deletions(-) diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py index 2320d702e..4f02633e7 100644 --- a/imod/mf6/interfaces/imodel.py +++ b/imod/mf6/interfaces/imodel.py @@ -1,5 +1,24 @@ +from imod.typing import GridDataArray +from typing import Optional +from imod.mf6.statusinfo import StatusInfoBase +class IModel: + @property + def items(self): + raise NotImplementedError + def mask_all_packages( self, mask: GridDataArray ): + raise NotImplementedError -class IModel: - pass \ No newline at end of file + def mask_all_packages( self, mask: GridDataArray ): + raise NotImplementedError + + def purge_empty_packages(self, model_name: Optional[str]) -> None: + raise NotImplementedError + + def _validate(self, model_name: str = "") -> StatusInfoBase: + raise NotImplementedError + + @property + def domain(self): + raise NotImplementedError \ No newline at end of file diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py index 8b41ef517..45dee1b33 100644 --- a/imod/mf6/interfaces/ipackage.py +++ b/imod/mf6/interfaces/ipackage.py @@ -1,12 +1,25 @@ import abc +from abc import abstractmethod from imod.mf6.interfaces.ipackagebase import IPackageBase - +from typing import Any +import xarray as xr class IPackage(IPackageBase, metaclass=abc.ABCMeta): """ The base methods and attributes available in all packages """ + @abstractmethod + def _valid(self, value): + raise NotImplementedError + + @abstractmethod + def __init__(self, *args, **kwargs) -> None: + raise NotImplementedError + + @abc.abstractmethod + def get_non_grid_data(self, grid_names: list[str]) -> dict[str, Any]: + raise NotImplementedError @property @abc.abstractmethod diff --git a/imod/mf6/interfaces/ipointdatapackage.py b/imod/mf6/interfaces/ipointdatapackage.py index 016a5e92a..09ee67b90 100644 --- a/imod/mf6/interfaces/ipointdatapackage.py +++ b/imod/mf6/interfaces/ipointdatapackage.py @@ -1,6 +1,6 @@ from abc import abstractmethod - -from numpy import ndarray +import numpy as np +from numpy.typing import NDArray from imod.mf6.interfaces.ipackagebase import IPackageBase @@ -12,10 +12,10 @@ class IPointDataPackage(IPackageBase): @property @abstractmethod - def x(self) -> ndarray[float]: + def x(self) -> NDArray[np.float64]: raise NotImplementedError @property @abstractmethod - def y(self) -> ndarray[float]: + def y(self) ->NDArray[np.float64]: raise NotImplementedError diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index 8bfd69891..7027e0f41 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -7,7 +7,9 @@ class IRegridPackage(IPackage, abc.ABC): + @property + @abc.abstractmethod def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: if hasattr(self, "_regrid_method"): return self._regrid_method diff --git a/imod/mf6/package.py b/imod/mf6/package.py index b623a9c73..9acfe7bdd 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -667,3 +667,24 @@ def auxiliary_data_fields(self) -> dict[str, str]: if hasattr(self, "_auxiliary_data"): return self._auxiliary_data return {} + + def get_non_grid_data(self, grid_names: list[str]) -> dict[str, Any]: + """ + This function copies the attributes of a dataset that are scalars, such as options. + + parameters + ---------- + grid_names: list of str + the names of the attribbutes of a dataset that are grids. + """ + result = {} + all_non_grid_data = list(self.dataset.keys()) + for name in grid_names: + if name in all_non_grid_data: + all_non_grid_data.remove(name) + for name in all_non_grid_data: + if "time" in self.dataset[name].coords: + result[name] = self.dataset[name] + else: + result[name] = self.dataset[name].values[()] + return result \ No newline at end of file diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index aeda38351..3305a887c 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -112,27 +112,6 @@ def get_regridder( return self.__get_existing_regridder(regridder_class, method) -def get_non_grid_data(package, grid_names: list[str]) -> dict[str, Any]: - """ - This function copies the attributes of a dataset that are scalars, such as options. - - parameters - ---------- - grid_names: list of str - the names of the attribbutes of a dataset that are grids. - """ - result = {} - all_non_grid_data = list(package.dataset.keys()) - for name in grid_names: - if name in all_non_grid_data: - all_non_grid_data.remove(name) - for name in all_non_grid_data: - if "time" in package.dataset[name].coords: - result[name] = package.dataset[name] - else: - result[name] = package.dataset[name].values[()] - return result - def assign_coord_if_present( coordname: str, target_grid: GridDataArray, maybe_has_coords_attr: Any From 223f74428c81dd60c6c4c6bf46507f7fb7b3fb05 Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 17:56:19 +0100 Subject: [PATCH 15/29] cleanup --- imod/mf6/interfaces/imodel.py | 10 +++------- imod/mf6/interfaces/iregridpackage.py | 1 - imod/mf6/model.py | 2 +- imod/mf6/utilities/regrid.py | 2 +- 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py index 4f02633e7..7951c5ebe 100644 --- a/imod/mf6/interfaces/imodel.py +++ b/imod/mf6/interfaces/imodel.py @@ -1,19 +1,15 @@ from imod.typing import GridDataArray from typing import Optional from imod.mf6.statusinfo import StatusInfoBase +import collections -class IModel: - @property - def items(self): - raise NotImplementedError +class IModel(collections.UserDict): def mask_all_packages( self, mask: GridDataArray ): raise NotImplementedError - def mask_all_packages( self, mask: GridDataArray ): - raise NotImplementedError - def purge_empty_packages(self, model_name: Optional[str]) -> None: + def purge_empty_packages(self, model_name: Optional[str] = "") -> None: raise NotImplementedError def _validate(self, model_name: str = "") -> StatusInfoBase: diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index 7027e0f41..cfd971187 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -9,7 +9,6 @@ class IRegridPackage(IPackage, abc.ABC): @property - @abc.abstractmethod def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: if hasattr(self, "_regrid_method"): return self._regrid_method diff --git a/imod/mf6/model.py b/imod/mf6/model.py index a69eb407e..cc6e75d93 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -30,7 +30,7 @@ from imod.typing import GridDataArray -class Modflow6Model(collections.UserDict, IModel, abc.ABC): +class Modflow6Model( IModel, abc.ABC): _mandatory_packages: tuple[str, ...] = () _model_id: Optional[str] = None _template: Template diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 3305a887c..3b1259ddf 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -253,7 +253,7 @@ def _regrid_like( if regridder_types is not None: regridder_settings.update(regridder_types) - new_package_data = get_non_grid_data(package, list(regridder_settings.keys())) + new_package_data = package.get_non_grid_data( list(regridder_settings.keys())) for ( varname, From d60f73a71141e24d910b1c0125ab9dcfbe0d725a Mon Sep 17 00:00:00 2001 From: luitjan Date: Thu, 7 Mar 2024 18:01:46 +0100 Subject: [PATCH 16/29] fix mypy errors --- imod/mf6/utilities/regrid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 3b1259ddf..e958d3a84 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -201,7 +201,7 @@ def _regrid_array( # reconvert the result to the same dtype as the original return regridded_array.astype(original_dtype) -@typedispatch +@typedispatch # type: ignore[no-redef] def _regrid_like( package: IRegridPackage, target_grid: GridDataArray, @@ -309,7 +309,7 @@ def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, lis ) return methods -@typedispatch +@typedispatch # type: ignore[no-redef] def _regrid_like( model: IModel, target_grid: GridDataArray, validate: bool = True ) -> IModel: From 672e7b38c43812669cbd9dd1ec8d6385d1082082 Mon Sep 17 00:00:00 2001 From: luitjan Date: Fri, 8 Mar 2024 14:14:52 +0100 Subject: [PATCH 17/29] cleanup and adding IDict --- imod/mf6/interfaces/idict.py | 60 +++++++++++++ imod/mf6/interfaces/imodel.py | 12 ++- imod/mf6/interfaces/ipackage.py | 8 +- imod/mf6/interfaces/ipackagebase.py | 2 +- imod/mf6/interfaces/ipointdatapackage.py | 1 + imod/mf6/interfaces/iregridpackage.py | 4 +- imod/mf6/interfaces/isimulation.py | 15 ++++ imod/mf6/model.py | 2 +- imod/mf6/simulation.py | 71 ++++++---------- imod/mf6/utilities/regrid.py | 98 +++++++++++++++++----- imod/tests/test_mf6/test_mf6_simulation.py | 2 +- 11 files changed, 194 insertions(+), 81 deletions(-) create mode 100644 imod/mf6/interfaces/idict.py create mode 100644 imod/mf6/interfaces/isimulation.py diff --git a/imod/mf6/interfaces/idict.py b/imod/mf6/interfaces/idict.py new file mode 100644 index 000000000..b7f0b9690 --- /dev/null +++ b/imod/mf6/interfaces/idict.py @@ -0,0 +1,60 @@ +import abc + + +class IDict(abc.ABC): + """ + Interface for collections.UserDict + """ + def __setitem__(self, key, item): + raise NotImplementedError + + def __getitem__(self, key): + raise NotImplementedError + + def __repr__(self): + raise NotImplementedError + + def __len__(self): + raise NotImplementedError + + def __delitem__(self, key): + raise NotImplementedError + + def clear(self): + raise NotImplementedError + + def copy(self): + raise NotImplementedError + + def has_key(self, k): + raise NotImplementedError + + def update(self, *args, **kwargs): + raise NotImplementedError + + def keys(self): + raise NotImplementedError + + def values(self): + raise NotImplementedError + + def items(self): + raise NotImplementedError + + def pop(self, *args): + raise NotImplementedError + + def __cmp__(self, dict_): + raise NotImplementedError + + def __contains__(self, item): + raise NotImplementedError + + def __iter__(self): + raise NotImplementedError + + def __unicode__(self): + raise NotImplementedError + + + diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py index 7951c5ebe..172a16d3b 100644 --- a/imod/mf6/interfaces/imodel.py +++ b/imod/mf6/interfaces/imodel.py @@ -1,10 +1,14 @@ -from imod.typing import GridDataArray from typing import Optional -from imod.mf6.statusinfo import StatusInfoBase -import collections -class IModel(collections.UserDict): +from imod.mf6.interfaces.idict import IDict +from imod.mf6.statusinfo import StatusInfoBase +from imod.typing import GridDataArray + +class IModel( IDict): + """ + Interface for imod.mf6.model.Modflow6Model + """ def mask_all_packages( self, mask: GridDataArray ): raise NotImplementedError diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py index 45dee1b33..aeb698b93 100644 --- a/imod/mf6/interfaces/ipackage.py +++ b/imod/mf6/interfaces/ipackage.py @@ -1,13 +1,13 @@ import abc -from abc import abstractmethod +from abc import abstractmethod +from typing import Any from imod.mf6.interfaces.ipackagebase import IPackageBase -from typing import Any -import xarray as xr + class IPackage(IPackageBase, metaclass=abc.ABCMeta): """ - The base methods and attributes available in all packages + Interface for imod.mf6.package.Package """ @abstractmethod def _valid(self, value): diff --git a/imod/mf6/interfaces/ipackagebase.py b/imod/mf6/interfaces/ipackagebase.py index fa9b4627b..038e20876 100644 --- a/imod/mf6/interfaces/ipackagebase.py +++ b/imod/mf6/interfaces/ipackagebase.py @@ -5,7 +5,7 @@ class IPackageBase(ABC): """ - The base methods and attributes available in all packages + Interface for imod.mf6.pkgbase.PackageBase """ @property diff --git a/imod/mf6/interfaces/ipointdatapackage.py b/imod/mf6/interfaces/ipointdatapackage.py index 09ee67b90..f20e728be 100644 --- a/imod/mf6/interfaces/ipointdatapackage.py +++ b/imod/mf6/interfaces/ipointdatapackage.py @@ -1,4 +1,5 @@ from abc import abstractmethod + import numpy as np from numpy.typing import NDArray diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index cfd971187..8fd26229b 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -7,7 +7,9 @@ class IRegridPackage(IPackage, abc.ABC): - + """ + Interface for packages that support regridding + """ @property def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: if hasattr(self, "_regrid_method"): diff --git a/imod/mf6/interfaces/isimulation.py b/imod/mf6/interfaces/isimulation.py new file mode 100644 index 000000000..404423854 --- /dev/null +++ b/imod/mf6/interfaces/isimulation.py @@ -0,0 +1,15 @@ + + +from imod.mf6.interfaces.idict import IDict + + +class ISimulation(IDict): + """ + Interface for imod.mf6.simulation.Modflow6Simulation + """ + def is_split(self) -> bool: + raise NotImplementedError + + + def has_one_flow_model(self) -> bool: + raise NotImplementedError \ No newline at end of file diff --git a/imod/mf6/model.py b/imod/mf6/model.py index cc6e75d93..c873d156f 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -30,7 +30,7 @@ from imod.typing import GridDataArray -class Modflow6Model( IModel, abc.ABC): +class Modflow6Model( collections.UserDict, IModel, abc.ABC): _mandatory_packages: tuple[str, ...] = () _model_id: Optional[str] = None _template: Template diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index ea9382121..c8ef05da2 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -1,7 +1,6 @@ from __future__ import annotations import collections -import copy import pathlib import subprocess import warnings @@ -25,6 +24,7 @@ from imod.mf6.gwfgwt import GWFGWT from imod.mf6.gwtgwt import GWTGWT from imod.mf6.ims import Solution +from imod.mf6.interfaces.isimulation import ISimulation from imod.mf6.model import Modflow6Model from imod.mf6.model_gwf import GroundwaterFlowModel from imod.mf6.model_gwt import GroundwaterTransportModel @@ -37,6 +37,7 @@ from imod.mf6.package import Package from imod.mf6.ssm import SourceSinkMixing from imod.mf6.statusinfo import NestedStatusInfo +from imod.mf6.utilities.regrid import _regrid_like from imod.mf6.write_context import WriteContext from imod.schemata import ValidationError from imod.typing import GridDataArray, GridDataset @@ -75,16 +76,7 @@ def get_packages(simulation: Modflow6Simulation) -> dict[str, Package]: if isinstance(pkg, Package) } - -def is_split(simulation: Modflow6Simulation) -> bool: - return "split_exchanges" in simulation.keys() - - -def has_one_flow_model(simulation: Modflow6Simulation) -> bool: - flow_models = simulation.get_models_of_type("gwf6") - return len(flow_models) == 1 - -class Modflow6Simulation(collections.UserDict): +class Modflow6Simulation(collections.UserDict, ISimulation): def _initialize_template(self): loader = jinja2.PackageLoader("imod", "templates/mf6") env = jinja2.Environment(loader=loader, keep_trailing_newline=True) @@ -252,7 +244,7 @@ def write( """ # create write context write_context = WriteContext(directory, binary, use_absolute_paths) - if is_split(self): + if self.is_split(): write_context.is_partitioned = True # Check models for required content @@ -606,7 +598,7 @@ def _open_single_output( elif len(modelnames) == 1: modelname = next(iter(modelnames)) return self._open_single_output_single_model(modelname, output, **settings) - elif is_split(self): + elif self.is_split(): if "budget" in output: return self._merge_budgets(modelnames, output, **settings) else: @@ -705,7 +697,7 @@ def _concat_species( # [[Ta, Tb]] -> [[Ta], [Tb]] tpt_names_per_species = list(zip(*all_tpt_names)) - if is_split(self): + if self.is_split(): # [[Ta_1, Tb_1], [Ta_2, Tb_2]] -> [Ta, Tb] unpartitioned_modelnames = [ tpt_name.rpartition("_")[0] for tpt_name in all_tpt_names[0] @@ -901,7 +893,7 @@ def get_exchange_relationships(self): result.append(exchange.get_specification()) # exchange for splitting models - if is_split(self): + if self.is_split(): for exchange in self["split_exchanges"]: result.append(exchange.get_specification()) return result @@ -954,12 +946,12 @@ def clip_box( clipped : Simulation """ - if is_split(self): + if self.is_split(): raise RuntimeError( "Unable to clip simulation. Clipping can only be done on simulations that haven't been split." + "Therefore clipping should be done before splitting the simulation." ) - if not has_one_flow_model(self): + if not self.has_one_flow_model(): raise ValueError( "Unable to clip simulation. Clipping can only be done on simulations that have a single flow model ." ) @@ -993,12 +985,12 @@ def split(self, submodel_labels: GridDataArray) -> Modflow6Simulation: The method return a new simulation containing all the split models and packages """ - if is_split(self): + if self.is_split(): raise RuntimeError( "Unable to split simulation. Splitting can only be done on simulations that haven't been split." ) - if not has_one_flow_model(self): + if not self.has_one_flow_model(): raise ValueError( "splitting of simulations with more (or less) than 1 flow model currently not supported." ) @@ -1060,7 +1052,7 @@ def split(self, submodel_labels: GridDataArray) -> Modflow6Simulation: def regrid_like( self, regridded_simulation_name: str, - target_grid: Union[xr.DataArray, xu.UgridDataArray], + target_grid: GridDataArray, validate: bool = True, ) -> "Modflow6Simulation": """ @@ -1082,32 +1074,10 @@ def regrid_like( a new simulation object with regridded models """ - if is_split(self): - raise RuntimeError( - "Unable to regrid simulation. Regridding can only be done on simulations that haven't been split." + - " Therefore regridding should be done before splitting the simulation." - ) - if not has_one_flow_model(self) : - raise ValueError( - "Unable to regrid simulation. Regridding can only be done on simulations that have a single flow model." - ) - result = self.__class__(regridded_simulation_name) - for key, item in self.items(): - if isinstance(item, Modflow6Model): - result[key] = item.regrid_like(target_grid, validate) - elif isinstance(item, imod.mf6.Solution) or isinstance( - item, imod.mf6.TimeDiscretization - ): - result[key] = copy.deepcopy(item) - elif key == "gwtgwf_exchanges": - pass - else: - raise NotImplementedError(f"regridding not supported for {key}") - - return result + return _regrid_like(self, regridded_simulation_name,target_grid, validate ) def _add_modelsplit_exchanges(self, exchanges_list: list[GWFGWF]) -> None: - if not is_split(self): + if not self.is_split(): self["split_exchanges"] = [] self["split_exchanges"].extend(exchanges_list) @@ -1239,7 +1209,7 @@ def _update_ssm_packages(self) -> None: "auxiliary_variable_name" ].values[0] ssm_package = SourceSinkMixing.from_flow_model( - flow_model, state_variable_name, is_split=is_split(self) + flow_model, state_variable_name, is_split=self.is_split() ) if ssm_package is not None: tpt_model[ssm_key] = ssm_package @@ -1248,4 +1218,13 @@ def _update_buoyancy_packages(self)->None: flow_transport_mapping = self._get_transport_models_per_flow_model() for flow_name, tpt_models_of_flow_model in flow_transport_mapping.items(): flow_model = self[flow_name] - flow_model.update_buoyancy_package(tpt_models_of_flow_model) \ No newline at end of file + flow_model.update_buoyancy_package(tpt_models_of_flow_model) + + + def is_split(self) -> bool: + return "split_exchanges" in self.keys() + + + def has_one_flow_model(self) -> bool: + flow_models = self.get_models_of_type("gwf6") + return len(flow_models) == 1 \ No newline at end of file diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index e958d3a84..e48813aa1 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -18,6 +18,7 @@ from imod.mf6.interfaces.ipackage import IPackage from imod.mf6.interfaces.ipointdatapackage import IPointDataPackage from imod.mf6.interfaces.iregridpackage import IRegridPackage +from imod.mf6.interfaces.isimulation import ISimulation from imod.mf6.statusinfo import NestedStatusInfo from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.utilities.regridding_types import RegridderType @@ -201,6 +202,30 @@ def _regrid_array( # reconvert the result to the same dtype as the original return regridded_array.astype(original_dtype) +def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, list[str]]: + """ + This function loops over the packages and collects all regridder-types that are in use. + """ + methods: defaultdict = defaultdict(list) + for pkg_name, pkg in model.items(): + if pkg.is_regridding_supported(): + pkg_methods = pkg.get_regrid_methods() + for variable in pkg_methods: + if ( + variable in pkg.dataset.data_vars + and pkg.dataset[variable].values[()] is not None + ): + regriddertype = pkg_methods[variable][0] + functiontype = pkg_methods[variable][1] + if functiontype not in methods[regriddertype]: + methods[regriddertype].append(functiontype) + else: + raise NotImplementedError( + f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" + ) + return methods + + @typedispatch # type: ignore[no-redef] def _regrid_like( package: IRegridPackage, @@ -286,29 +311,6 @@ def _regrid_like( return package.__class__(**new_package_data) -def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, list[str]]: - """ - This function loops over the packages and collects all regridder-types that are in use. - """ - methods: defaultdict = defaultdict(list) - for pkg_name, pkg in model.items(): - if pkg.is_regridding_supported(): - pkg_methods = pkg.get_regrid_methods() - for variable in pkg_methods: - if ( - variable in pkg.dataset.data_vars - and pkg.dataset[variable].values[()] is not None - ): - regriddertype = pkg_methods[variable][0] - functiontype = pkg_methods[variable][1] - if functiontype not in methods[regriddertype]: - methods[regriddertype].append(functiontype) - else: - raise NotImplementedError( - f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" - ) - return methods - @typedispatch # type: ignore[no-redef] def _regrid_like( model: IModel, target_grid: GridDataArray, validate: bool = True @@ -352,6 +354,56 @@ def _regrid_like( raise ValidationError("\n" + status_info.to_string()) return new_model +@typedispatch # type: ignore[no-redef] +def _regrid_like( + simulation: ISimulation, + regridded_simulation_name: str, + target_grid: GridDataArray, + validate: bool = True, +) -> ISimulation: + """ + This method creates a new simulation object. The models contained in the new simulation are regridded versions + of the models in the input object (this). + Time discretization and solver settings are copied. + + Parameters + ---------- + regridded_simulation_name: str + name given to the output simulation + target_grid: xr.DataArray or xu.UgridDataArray + discretization onto which the models in this simulation will be regridded + validate: bool + set to true to validate the regridded packages + + Returns + ------- + a new simulation object with regridded models + """ + + if simulation.is_split(): + raise RuntimeError( + "Unable to regrid simulation. Regridding can only be done on simulations that haven't been split." + + " Therefore regridding should be done before splitting the simulation." + ) + if not simulation.has_one_flow_model() : + raise ValueError( + "Unable to regrid simulation. Regridding can only be done on simulations that have a single flow model." + ) + result = simulation.__class__(regridded_simulation_name) + for key, item in simulation.items(): + if isinstance(item, IModel): + result[key] = item.regrid_like(target_grid, validate) + elif key == "gwtgwf_exchanges": + pass + elif isinstance(item, IPackage) and not isinstance(item, IRegridPackage): + result[key] = copy.deepcopy(item) + + else: + raise NotImplementedError(f"regridding not supported for {key}") + + return result + + def _get_regridding_domain( model: IModel, target_grid: GridDataArray, diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py index 8218c70fe..78a5b3b5d 100644 --- a/imod/tests/test_mf6/test_mf6_simulation.py +++ b/imod/tests/test_mf6/test_mf6_simulation.py @@ -415,7 +415,7 @@ def test_prevent_regrid_like_after_split( # Act/Assert with pytest.raises(RuntimeError): - _ = split_simulation.regrid_like(None, None) + _ = split_simulation.regrid_like("new_simulation", split_transient_twri_model["GWF_1_2"].domain) def compare_submodel_partition_info(first: PartitionInfo, second: PartitionInfo): From 6162f1822f914145038a21b6f710345372a8b8ba Mon Sep 17 00:00:00 2001 From: luitjan Date: Fri, 8 Mar 2024 14:29:46 +0100 Subject: [PATCH 18/29] cleanup --- imod/mf6/hfb.py | 11 ---------- imod/mf6/utilities/regrid.py | 42 ++++++++++++++++++------------------ imod/mf6/wel.py | 9 -------- 3 files changed, 21 insertions(+), 41 deletions(-) diff --git a/imod/mf6/hfb.py b/imod/mf6/hfb.py index 0f8389972..e0aa0f226 100644 --- a/imod/mf6/hfb.py +++ b/imod/mf6/hfb.py @@ -16,7 +16,6 @@ from imod.mf6.interfaces.ilinedatapackage import ILineDataPackage from imod.mf6.mf6_hfb_adapter import Mf6HorizontalFlowBarrier from imod.mf6.package import Package -from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.utilities.grid import broadcast_to_full_domain from imod.mf6.utilities.regrid import RegridderType from imod.schemata import EmptyIndexesSchema @@ -592,16 +591,6 @@ def clip_box( new.dataset = copy.deepcopy(self.dataset) return new - def regrid_like( - self, target_grid: GridDataArray, *_ - ) -> "HorizontalFlowBarrierBase": - """ - The regrid_like method is irrelevant for this package as it is - grid-agnostic, instead this method clips the package based on the grid - exterior. - """ - return clip_by_grid(self, target_grid) - def mask(self, _) -> Package: """ The mask method is irrelevant for this package as it is diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index e48813aa1..821716583 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -128,27 +128,6 @@ def assign_coord_if_present( ) return maybe_has_coords_attr -@typedispatch # type: ignore[no-redef] -def regrid_like( - package: ILineDataPackage, target_grid: GridDataArray, *_) -> ILineDataPackage: - """ - The regrid_like method is irrelevant for this package as it is - grid-agnostic, instead this method clips the package based on the grid - exterior. - """ - return clip_by_grid(package, target_grid) - -@typedispatch # type: ignore[no-redef] -def regrid_like( - package: IPointDataPackage, target_grid: GridDataArray, *_ -) -> IPointDataPackage: - """ - The regrid_like method is irrelevant for this package as it is - grid-agnostic, instead this method clips the package based on the grid - exterior. - """ - target_grid_2d = target_grid.isel(layer=0, drop=True, missing_dims="ignore") - return clip_by_grid(package, target_grid_2d) def _regrid_array( package: IRegridPackage, @@ -403,6 +382,27 @@ def _regrid_like( return result +@typedispatch # type: ignore[no-redef] +def _regrid_like( + package: ILineDataPackage, target_grid: GridDataArray, *_) -> ILineDataPackage: + """ + The regrid_like method is irrelevant for this package as it is + grid-agnostic, instead this method clips the package based on the grid + exterior. + """ + return clip_by_grid(package, target_grid) + +@typedispatch # type: ignore[no-redef] +def _regrid_like( + package: IPointDataPackage, target_grid: GridDataArray, *_ +) -> IPointDataPackage: + """ + he regrid_like method is irrelevant for this package as it is + grid-agnostic, instead this method clips the package based on the grid + exterior. + """ + target_grid_2d = target_grid.isel(layer=0, drop=True, missing_dims="ignore") + return clip_by_grid(package, target_grid_2d) def _get_regridding_domain( model: IModel, diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index 6e03c804a..d1fa7e06c 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -19,7 +19,6 @@ from imod.mf6.interfaces.ipointdatapackage import IPointDataPackage from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.package import Package -from imod.mf6.utilities.clip import clip_by_grid from imod.mf6.utilities.dataset import remove_inactive from imod.mf6.utilities.grid import create_layered_top from imod.mf6.utilities.regrid import RegridderType @@ -596,14 +595,6 @@ def to_mf6_pkg( return Mf6Wel(**ds.data_vars) - def regrid_like(self, target_grid: GridDataArray, *_) -> Well: - """ - The regrid_like method is irrelevant for this package as it is - grid-agnostic, instead this method clips the package based on the grid - exterior. - """ - target_grid_2d = target_grid.isel(layer=0, drop=True, missing_dims="ignore") - return clip_by_grid(self, target_grid_2d) def mask(self, domain: GridDataArray) -> Well: """ From dfeff0aa6c7cd5ae0a791a3bc500478104255a76 Mon Sep 17 00:00:00 2001 From: luitjan Date: Fri, 8 Mar 2024 16:39:01 +0100 Subject: [PATCH 19/29] review comments (_validate public, and if/loop combining) --- imod/mf6/interfaces/imodel.py | 2 +- imod/mf6/model.py | 6 +++--- imod/mf6/package.py | 5 ++--- imod/mf6/utilities/regrid.py | 2 +- imod/tests/test_mf6/test_mf6_regrid_model.py | 10 +++++----- 5 files changed, 12 insertions(+), 13 deletions(-) diff --git a/imod/mf6/interfaces/imodel.py b/imod/mf6/interfaces/imodel.py index 172a16d3b..08e62cad7 100644 --- a/imod/mf6/interfaces/imodel.py +++ b/imod/mf6/interfaces/imodel.py @@ -16,7 +16,7 @@ def mask_all_packages( self, mask: GridDataArray ): def purge_empty_packages(self, model_name: Optional[str] = "") -> None: raise NotImplementedError - def _validate(self, model_name: str = "") -> StatusInfoBase: + def validate(self, model_name: str = "") -> StatusInfoBase: raise NotImplementedError @property diff --git a/imod/mf6/model.py b/imod/mf6/model.py index c873d156f..ab5a2425e 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -185,7 +185,7 @@ def __get_k(self): k = npf["k"] return k - def _validate(self, model_name: str = "") -> StatusInfoBase: + def validate(self, model_name: str = "") -> StatusInfoBase: try: diskey = self._get_diskey() except Exception as e: @@ -236,7 +236,7 @@ def write( modeldirectory = workdir / modelname Path(modeldirectory).mkdir(exist_ok=True, parents=True) if validate: - model_status_info = self._validate(modelname) + model_status_info = self.validate(modelname) if model_status_info.has_errors(): return model_status_info @@ -295,7 +295,7 @@ def dump( modeldirectory = pathlib.Path(directory) / modelname modeldirectory.mkdir(exist_ok=True, parents=True) if validate: - statusinfo = self._validate() + statusinfo = self.validate() if statusinfo.has_errors(): raise ValidationError(statusinfo.to_string()) diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 9acfe7bdd..3364ac456 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -679,9 +679,8 @@ def get_non_grid_data(self, grid_names: list[str]) -> dict[str, Any]: """ result = {} all_non_grid_data = list(self.dataset.keys()) - for name in grid_names: - if name in all_non_grid_data: - all_non_grid_data.remove(name) + for name in (gridname for gridname in grid_names if gridname in all_non_grid_data): + all_non_grid_data.remove(name) for name in all_non_grid_data: if "time" in self.dataset[name].coords: result[name] = self.dataset[name] diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 821716583..b97248e33 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -328,7 +328,7 @@ def _regrid_like( new_model.purge_empty_packages() if validate: status_info = NestedStatusInfo("Model validation status") - status_info.add(new_model._validate("Regridded model")) + status_info.add(new_model.validate("Regridded model")) if status_info.has_errors(): raise ValidationError("\n" + status_info.to_string()) return new_model diff --git a/imod/tests/test_mf6/test_mf6_regrid_model.py b/imod/tests/test_mf6/test_mf6_regrid_model.py index 2f0cf9e0d..f75d60143 100644 --- a/imod/tests/test_mf6/test_mf6_regrid_model.py +++ b/imod/tests/test_mf6/test_mf6_regrid_model.py @@ -19,7 +19,7 @@ def test_regrid_structured_model_to_structured_model( new_gwf_model = structured_flow_model.regrid_like(finer_idomain) assert len(new_gwf_model.items()) == len(structured_flow_model.items()) - validation_result = new_gwf_model._validate("test_model") + validation_result = new_gwf_model.validate("test_model") assert not validation_result.has_errors() @@ -40,7 +40,7 @@ def test_regrid_structured_model_with_wells_to_structured_model( new_gwf_model = structured_flow_model.regrid_like(finer_idomain) assert len(new_gwf_model.items()) == len(structured_flow_model.items()) - validation_result = new_gwf_model._validate("test_model") + validation_result = new_gwf_model.validate("test_model") assert not validation_result.has_errors() @@ -61,7 +61,7 @@ def test_regrid_unstructured_model_with_wells_to_unstructured_model( new_gwf_model = unstructured_flow_model.regrid_like(finer_idomain) assert len(new_gwf_model.items()) == len(unstructured_flow_model.items()) - validation_result = new_gwf_model._validate("test_model") + validation_result = new_gwf_model.validate("test_model") assert not validation_result.has_errors() @@ -73,7 +73,7 @@ def test_regrid_unstructured_model_to_unstructured_model( new_gwf_model = unstructured_flow_model.regrid_like(finer_idomain) assert len(new_gwf_model.items()) == len(unstructured_flow_model.items()) - validation_result = new_gwf_model._validate("test_model") + validation_result = new_gwf_model.validate("test_model") assert not validation_result.has_errors() @@ -113,7 +113,7 @@ def test_regrid_unstructured_model_with_inactive_cells( new_gwf_model = unstructured_flow_model.regrid_like(finer_idomain) assert len(new_gwf_model.items()) == len(unstructured_flow_model.items()) - validation_result = new_gwf_model._validate("test_model") + validation_result = new_gwf_model.validate("test_model") assert not validation_result.has_errors() new_idomain = new_gwf_model.domain assert ( From 6071110d7845cbc3558eff9bde7e2ffb8402acf9 Mon Sep 17 00:00:00 2001 From: luitjan Date: Mon, 11 Mar 2024 11:50:30 +0100 Subject: [PATCH 20/29] review comments --- imod/mf6/interfaces/iregridpackage.py | 6 +----- imod/tests/test_mf6/test_mf6_model_masking.py | 6 +++--- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index 8fd26229b..abaecbfa5 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -10,8 +10,4 @@ class IRegridPackage(IPackage, abc.ABC): """ Interface for packages that support regridding """ - @property - def regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: - if hasattr(self, "_regrid_method"): - return self._regrid_method - return None \ No newline at end of file + pass \ No newline at end of file diff --git a/imod/tests/test_mf6/test_mf6_model_masking.py b/imod/tests/test_mf6/test_mf6_model_masking.py index a8e4eb293..1a46fec0a 100644 --- a/imod/tests/test_mf6/test_mf6_model_masking.py +++ b/imod/tests/test_mf6/test_mf6_model_masking.py @@ -25,7 +25,7 @@ def test_masked_model_validation_inactive_cell_pillar( unstructured_flow_model.mask_all_packages(mask) # test output validity - errors = unstructured_flow_model._validate("model") + errors = unstructured_flow_model.validate("model") assert len(errors.errors) == 0 assert_model_can_run(unstructured_flow_model, "disv", tmp_path) @@ -47,7 +47,7 @@ def test_masked_model_validation_one_inactive_cell( unstructured_flow_model.mask_all_packages(mask) # test output validity - errors = unstructured_flow_model._validate("model") + errors = unstructured_flow_model.validate("model") assert len(errors.errors) == 0 assert_model_can_run(unstructured_flow_model, "disv", tmp_path) @@ -88,7 +88,7 @@ def test_masked_model_layered_and_scalar_package_input( unstructured_flow_model.mask_all_packages(mask) # Test output validity - errors = unstructured_flow_model._validate("model") + errors = unstructured_flow_model.validate("model") assert len(errors.errors) == 0 assert_model_can_run(unstructured_flow_model, "disv", tmp_path) From d450fa39ceb31519fbe1039ad442c80f3e3eea22 Mon Sep 17 00:00:00 2001 From: luitjan Date: Mon, 11 Mar 2024 17:51:43 +0100 Subject: [PATCH 21/29] review comments --- imod/mf6/interfaces/iregridpackage.py | 3 ++- imod/mf6/package.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/imod/mf6/interfaces/iregridpackage.py b/imod/mf6/interfaces/iregridpackage.py index abaecbfa5..698061cb3 100644 --- a/imod/mf6/interfaces/iregridpackage.py +++ b/imod/mf6/interfaces/iregridpackage.py @@ -10,4 +10,5 @@ class IRegridPackage(IPackage, abc.ABC): """ Interface for packages that support regridding """ - pass \ No newline at end of file + def get_regrid_methods(self) -> Optional[dict[str, tuple[RegridderType, str]]]: + raise NotImplementedError \ No newline at end of file diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 3364ac456..246421ca9 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -686,4 +686,4 @@ def get_non_grid_data(self, grid_names: list[str]) -> dict[str, Any]: result[name] = self.dataset[name] else: result[name] = self.dataset[name].values[()] - return result \ No newline at end of file + return result From 59d0e104cfb7548eb64bd29431ac5a360619225c Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 12:57:09 +0100 Subject: [PATCH 22/29] tests pass --- imod/mf6/adv.py | 6 ++++-- imod/mf6/chd.py | 5 ++++- imod/mf6/dis.py | 6 +++++- imod/mf6/disv.py | 4 ++++ imod/mf6/drn.py | 5 ++++- imod/mf6/dsp.py | 7 +++++-- imod/mf6/evt.py | 5 ++++- imod/mf6/ghb.py | 5 ++++- imod/mf6/hfb.py | 4 +++- imod/mf6/ic.py | 5 ++++- imod/mf6/ist.py | 6 +++++- imod/mf6/mst.py | 5 ++++- imod/mf6/npf.py | 7 +++++-- imod/mf6/oc.py | 4 ++++ imod/mf6/package.py | 6 ------ imod/mf6/rch.py | 5 ++++- imod/mf6/riv.py | 5 ++++- imod/mf6/src.py | 2 +- imod/mf6/ssm.py | 4 +++- imod/mf6/sto.py | 5 +++-- imod/mf6/utilities/regrid.py | 2 +- imod/mf6/wel.py | 3 +++ 22 files changed, 78 insertions(+), 28 deletions(-) diff --git a/imod/mf6/adv.py b/imod/mf6/adv.py index 2103556cc..d50977f85 100644 --- a/imod/mf6/adv.py +++ b/imod/mf6/adv.py @@ -13,7 +13,7 @@ from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType - +from typing import Optional, Tuple class Advection(Package): _pkg_id = "adv" @@ -34,7 +34,9 @@ def mask(self, _) -> Package: retuns a copy of itself. """ return deepcopy(self) - + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method class AdvectionUpstream(Advection): """ diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py index 286121655..73e687758 100644 --- a/imod/mf6/chd.py +++ b/imod/mf6/chd.py @@ -14,7 +14,7 @@ IndexesSchema, OtherCoordsSchema, ) - +from typing import Optional, Tuple class ConstantHead(BoundaryCondition, IRegridPackage): """ @@ -145,3 +145,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index 050d24612..3504bbc96 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -1,7 +1,7 @@ import pathlib import numpy as np - +from typing import Tuple, Optional import imod from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package @@ -144,4 +144,8 @@ def _validate(self, schemata, **kwargs): kwargs["bottom"] = self["bottom"] errors = super()._validate(schemata, **kwargs) + return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/disv.py b/imod/mf6/disv.py index 6307cdfb5..23d0ae683 100644 --- a/imod/mf6/disv.py +++ b/imod/mf6/disv.py @@ -1,5 +1,6 @@ import numpy as np import pandas as pd +from typing import Tuple, Optional from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package @@ -152,3 +153,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index 4858fa6f6..9e6ea6eff 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -15,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) - +from typing import Optional, Tuple class Drainage(BoundaryCondition, IRegridPackage): """ @@ -148,3 +148,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index eed8cfa6b..5ba680f49 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -11,7 +11,7 @@ IdentityNoDataSchema, IndexesSchema, ) - +from typing import Optional, Tuple class Dispersion(Package, IRegridPackage): """ @@ -187,4 +187,7 @@ def _validate(self, schemata, **kwargs): kwargs["xt3d_off"] = self["xt3d_off"] errors = super()._validate(schemata, **kwargs) - return errors \ No newline at end of file + return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/evt.py b/imod/mf6/evt.py index 63ab9e4a3..30d2fc8a6 100644 --- a/imod/mf6/evt.py +++ b/imod/mf6/evt.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, Tuple import numpy as np @@ -254,3 +254,6 @@ def _get_bin_ds(self): bin_ds = unstack_dim_into_variable(bin_ds, "segment") return bin_ds + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index 323216452..2c7c9c2da 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -15,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) - +from typing import Optional, Tuple class GeneralHeadBoundary(BoundaryCondition,IRegridPackage): """ @@ -153,3 +153,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/hfb.py b/imod/mf6/hfb.py index e0aa0f226..1574eb39a 100644 --- a/imod/mf6/hfb.py +++ b/imod/mf6/hfb.py @@ -295,7 +295,9 @@ def __init__( super().__init__(dict_dataset) self.line_data = geometry - + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method + def _get_variable_names_for_gdf(self) -> list[str]: return [ self._get_variable_name(), diff --git a/imod/mf6/ic.py b/imod/mf6/ic.py index cbac5b306..8dacf3f4f 100644 --- a/imod/mf6/ic.py +++ b/imod/mf6/ic.py @@ -7,7 +7,7 @@ from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import DTypeSchema, IdentityNoDataSchema, IndexesSchema - +from typing import Optional, Tuple class InitialConditions(Package, IRegridPackage): """ @@ -92,3 +92,6 @@ def render(self, directory, pkgname, globaltimes, binary): self["start"], icdirectory, "strt", binary=binary ) return self._template.render(d) + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index 70c6c6490..f8adc15dd 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -12,7 +12,8 @@ IdentityNoDataSchema, IndexesSchema, ) - +from typing import Optional, Tuple +from imod.mf6.utilities.regrid import RegridderType class ImmobileStorageTransfer(Package, IRegridPackage): """ @@ -253,3 +254,6 @@ def __init__( } super().__init__(dict_dataset) self._validate_init_schemata(validate) + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index c550fd706..1cad8d754 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -1,5 +1,5 @@ import numpy as np - +from typing import Optional, Tuple from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType @@ -142,3 +142,6 @@ def __init__( } super().__init__(dict_dataset) self._validate_init_schemata(validate) + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/npf.py b/imod/mf6/npf.py index 4981f4d1b..b7cec1aa6 100644 --- a/imod/mf6/npf.py +++ b/imod/mf6/npf.py @@ -15,7 +15,7 @@ IndexesSchema, ) from imod.typing import GridDataArray - +from typing import Tuple, Optional def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: if griddataarray is None or griddataarray.values is None: @@ -450,4 +450,7 @@ def _validate(self, schemata, **kwargs): kwargs["xt3d_option"] = self["xt3d_option"] errors = super()._validate(schemata, **kwargs) - return errors \ No newline at end of file + return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/oc.py b/imod/mf6/oc.py index 6676b93fb..8d9449b0b 100644 --- a/imod/mf6/oc.py +++ b/imod/mf6/oc.py @@ -10,6 +10,7 @@ from imod.mf6.utilities.regrid import RegridderType from imod.mf6.write_context import WriteContext from imod.schemata import DTypeSchema +from typing import Optional, Tuple OUTPUT_EXT_MAPPING = { "head": "hds", @@ -206,3 +207,6 @@ def write( @property def is_budget_output(self) -> bool: return self.dataset["save_budget"].values[()] is not None + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 246421ca9..b6e3aa87a 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -600,12 +600,6 @@ def is_regridding_supported(self) -> bool: """ return hasattr(self, "_regrid_method") - def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: - if hasattr(self, "_regrid_method"): - return self._regrid_method - return None - - def regrid_like( self, diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index c578a4501..fce68b8e7 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -15,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) - +from typing import Optional, Tuple class Recharge(BoundaryCondition, IRegridPackage): """ @@ -140,3 +140,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index c0211b0c5..75b88a940 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -15,7 +15,7 @@ IndexesSchema, OtherCoordsSchema, ) - +from typing import Optional, Tuple class River(BoundaryCondition, IRegridPackage): """ @@ -167,3 +167,6 @@ def _validate(self, schemata, **kwargs): errors = super()._validate(schemata, **kwargs) return errors + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/src.py b/imod/mf6/src.py index b27ae51c3..684688934 100644 --- a/imod/mf6/src.py +++ b/imod/mf6/src.py @@ -13,7 +13,7 @@ ) -class MassSourceLoading(BoundaryCondition, IRegridPackage): +class MassSourceLoading(BoundaryCondition): """ Mass Source Loading (SRC) package for structured discretization (DIS) models. Any number of SRC Packages can be specified for a single diff --git a/imod/mf6/ssm.py b/imod/mf6/ssm.py index 62276d3eb..3766c2723 100644 --- a/imod/mf6/ssm.py +++ b/imod/mf6/ssm.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Tuple, Optional import numpy as np @@ -169,3 +169,5 @@ def from_flow_model( save_flows=save_flows, validate=validate, ) + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 9895d4d8d..0693c55a1 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -13,7 +13,7 @@ IdentityNoDataSchema, IndexesSchema, ) - +from typing import Optional, Tuple class Storage(Package): _pkg_id = "sto_deprecated" @@ -63,7 +63,8 @@ def _render_dict(self, directory, pkgname, globaltimes, binary): d = self.get_options(d) return d - + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method class SpecificStorage(StorageBase): """ diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index b97248e33..0bc86688b 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -253,7 +253,7 @@ def _regrid_like( package.dataset, target_grid=target_grid ) - regridder_settings = copy.deepcopy(package._regrid_method) + regridder_settings = package.get_regrid_methods() if regridder_types is not None: regridder_settings.update(regridder_types) diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index d1fa7e06c..7845b773a 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -611,6 +611,9 @@ def mask(self, domain: GridDataArray) -> Well: "layer", errors="ignore" ) return mask_2D(self, domain_2d) + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method class WellDisStructured(DisStructuredBoundaryCondition): From ef1280e3e5cda4d45aa2604674f743c1eb78453c Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 12:58:09 +0100 Subject: [PATCH 23/29] ruff --- imod/mf6/adv.py | 3 ++- imod/mf6/chd.py | 4 +++- imod/mf6/dis.py | 3 ++- imod/mf6/disv.py | 3 ++- imod/mf6/drn.py | 4 +++- imod/mf6/dsp.py | 4 +++- imod/mf6/ghb.py | 4 +++- imod/mf6/ic.py | 3 ++- imod/mf6/ist.py | 6 +++--- imod/mf6/mst.py | 4 +++- imod/mf6/npf.py | 3 ++- imod/mf6/oc.py | 3 +-- imod/mf6/rch.py | 4 +++- imod/mf6/riv.py | 4 +++- imod/mf6/src.py | 1 - imod/mf6/ssm.py | 2 +- imod/mf6/sto.py | 3 ++- 17 files changed, 38 insertions(+), 20 deletions(-) diff --git a/imod/mf6/adv.py b/imod/mf6/adv.py index d50977f85..799d753da 100644 --- a/imod/mf6/adv.py +++ b/imod/mf6/adv.py @@ -10,10 +10,11 @@ """ from copy import deepcopy +from typing import Optional, Tuple from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType -from typing import Optional, Tuple + class Advection(Package): _pkg_id = "adv" diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py index 73e687758..14616d9bd 100644 --- a/imod/mf6/chd.py +++ b/imod/mf6/chd.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.boundary_condition import BoundaryCondition @@ -14,7 +16,7 @@ IndexesSchema, OtherCoordsSchema, ) -from typing import Optional, Tuple + class ConstantHead(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/dis.py b/imod/mf6/dis.py index 3504bbc96..c050e00b7 100644 --- a/imod/mf6/dis.py +++ b/imod/mf6/dis.py @@ -1,7 +1,8 @@ import pathlib +from typing import Optional, Tuple import numpy as np -from typing import Tuple, Optional + import imod from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package diff --git a/imod/mf6/disv.py b/imod/mf6/disv.py index 23d0ae683..79538f209 100644 --- a/imod/mf6/disv.py +++ b/imod/mf6/disv.py @@ -1,6 +1,7 @@ +from typing import Optional, Tuple + import numpy as np import pandas as pd -from typing import Tuple, Optional from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index 9e6ea6eff..ad3d6db38 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.boundary_condition import BoundaryCondition @@ -15,7 +17,7 @@ IndexesSchema, OtherCoordsSchema, ) -from typing import Optional, Tuple + class Drainage(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index 5ba680f49..75fce6af2 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.interfaces.iregridpackage import IRegridPackage @@ -11,7 +13,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from typing import Optional, Tuple + class Dispersion(Package, IRegridPackage): """ diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index 2c7c9c2da..604be0f5b 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.boundary_condition import BoundaryCondition @@ -15,7 +17,7 @@ IndexesSchema, OtherCoordsSchema, ) -from typing import Optional, Tuple + class GeneralHeadBoundary(BoundaryCondition,IRegridPackage): """ diff --git a/imod/mf6/ic.py b/imod/mf6/ic.py index 8dacf3f4f..c3d3daab9 100644 --- a/imod/mf6/ic.py +++ b/imod/mf6/ic.py @@ -1,4 +1,5 @@ import warnings +from typing import Optional, Tuple import numpy as np @@ -7,7 +8,7 @@ from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import DTypeSchema, IdentityNoDataSchema, IndexesSchema -from typing import Optional, Tuple + class InitialConditions(Package, IRegridPackage): """ diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index f8adc15dd..94ad7d0bd 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -1,10 +1,11 @@ -from typing import Optional +from typing import Optional, Tuple import numpy as np import xarray as xr from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package +from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import ( AllValueSchema, @@ -12,8 +13,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from typing import Optional, Tuple -from imod.mf6.utilities.regrid import RegridderType + class ImmobileStorageTransfer(Package, IRegridPackage): """ diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index 1cad8d754..d2e44f948 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -1,5 +1,7 @@ -import numpy as np from typing import Optional, Tuple + +import numpy as np + from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType diff --git a/imod/mf6/npf.py b/imod/mf6/npf.py index b7cec1aa6..47432ffec 100644 --- a/imod/mf6/npf.py +++ b/imod/mf6/npf.py @@ -1,4 +1,5 @@ import warnings +from typing import Optional, Tuple import numpy as np @@ -15,7 +16,7 @@ IndexesSchema, ) from imod.typing import GridDataArray -from typing import Tuple, Optional + def _dataarray_to_bool(griddataarray: GridDataArray) -> bool: if griddataarray is None or griddataarray.values is None: diff --git a/imod/mf6/oc.py b/imod/mf6/oc.py index 8d9449b0b..075e18c06 100644 --- a/imod/mf6/oc.py +++ b/imod/mf6/oc.py @@ -1,7 +1,7 @@ import collections import os from pathlib import Path -from typing import Tuple, Union +from typing import Optional, Tuple, Union import numpy as np @@ -10,7 +10,6 @@ from imod.mf6.utilities.regrid import RegridderType from imod.mf6.write_context import WriteContext from imod.schemata import DTypeSchema -from typing import Optional, Tuple OUTPUT_EXT_MAPPING = { "head": "hds", diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index fce68b8e7..163a33dd2 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.boundary_condition import BoundaryCondition @@ -15,7 +17,7 @@ IndexesSchema, OtherCoordsSchema, ) -from typing import Optional, Tuple + class Recharge(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index 75b88a940..340a4402a 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from imod.mf6.boundary_condition import BoundaryCondition @@ -15,7 +17,7 @@ IndexesSchema, OtherCoordsSchema, ) -from typing import Optional, Tuple + class River(BoundaryCondition, IRegridPackage): """ diff --git a/imod/mf6/src.py b/imod/mf6/src.py index 684688934..195de09ab 100644 --- a/imod/mf6/src.py +++ b/imod/mf6/src.py @@ -1,7 +1,6 @@ import numpy as np from imod.mf6.boundary_condition import BoundaryCondition -from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA from imod.schemata import ( diff --git a/imod/mf6/ssm.py b/imod/mf6/ssm.py index 3766c2723..c12c82f3b 100644 --- a/imod/mf6/ssm.py +++ b/imod/mf6/ssm.py @@ -1,4 +1,4 @@ -from typing import Tuple, Optional +from typing import Optional, Tuple import numpy as np diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 0693c55a1..4e0cfc59e 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -1,4 +1,5 @@ import abc +from typing import Optional, Tuple import numpy as np @@ -13,7 +14,7 @@ IdentityNoDataSchema, IndexesSchema, ) -from typing import Optional, Tuple + class Storage(Package): _pkg_id = "sto_deprecated" From f35c5266ca6975cc5903fb2d0667f6c41f8a8ed6 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 13:13:07 +0100 Subject: [PATCH 24/29] advection is regrid package now --- imod/mf6/adv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/mf6/adv.py b/imod/mf6/adv.py index 799d753da..2ed449231 100644 --- a/imod/mf6/adv.py +++ b/imod/mf6/adv.py @@ -14,9 +14,9 @@ from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType +from imod.mf6.interfaces.iregridpackage import IRegridPackage - -class Advection(Package): +class Advection(Package, IRegridPackage): _pkg_id = "adv" _template = Package._initialize_template(_pkg_id) _regrid_method: dict[str, tuple[RegridderType, str]] = {} From bba0a74b0b4b04b9ba21b2307c85f31ae8f2ba60 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 13:31:50 +0100 Subject: [PATCH 25/29] fixed OutputControl and Adv --- imod/mf6/adv.py | 3 ++- imod/mf6/oc.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/imod/mf6/adv.py b/imod/mf6/adv.py index 2ed449231..448e4cc64 100644 --- a/imod/mf6/adv.py +++ b/imod/mf6/adv.py @@ -12,9 +12,10 @@ from copy import deepcopy from typing import Optional, Tuple +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.regrid import RegridderType -from imod.mf6.interfaces.iregridpackage import IRegridPackage + class Advection(Package, IRegridPackage): _pkg_id = "adv" diff --git a/imod/mf6/oc.py b/imod/mf6/oc.py index 075e18c06..111c7ef2b 100644 --- a/imod/mf6/oc.py +++ b/imod/mf6/oc.py @@ -10,6 +10,7 @@ from imod.mf6.utilities.regrid import RegridderType from imod.mf6.write_context import WriteContext from imod.schemata import DTypeSchema +from imod.mf6.interfaces.iregridpackage import IRegridPackage OUTPUT_EXT_MAPPING = { "head": "hds", @@ -18,7 +19,7 @@ } -class OutputControl(Package): +class OutputControl(Package, IRegridPackage): """ The Output Control Option determines how and when heads, budgets and/or concentrations are printed to the listing file and/or written to a separate From 1fb28ef70146fc79d6d4c86968a29a9ff97e496b Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 14:55:42 +0100 Subject: [PATCH 26/29] removed is_regridding_supported --- imod/mf6/oc.py | 2 +- imod/mf6/package.py | 6 ------ imod/mf6/ssm.py | 3 ++- imod/mf6/utilities/regrid.py | 10 +++------- 4 files changed, 6 insertions(+), 15 deletions(-) diff --git a/imod/mf6/oc.py b/imod/mf6/oc.py index 111c7ef2b..362d3ed10 100644 --- a/imod/mf6/oc.py +++ b/imod/mf6/oc.py @@ -5,12 +5,12 @@ import numpy as np +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package from imod.mf6.utilities.dataset import is_dataarray_none from imod.mf6.utilities.regrid import RegridderType from imod.mf6.write_context import WriteContext from imod.schemata import DTypeSchema -from imod.mf6.interfaces.iregridpackage import IRegridPackage OUTPUT_EXT_MAPPING = { "head": "hds", diff --git a/imod/mf6/package.py b/imod/mf6/package.py index b6e3aa87a..324302364 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -594,12 +594,6 @@ def _adjust_mask_for_unlayered_data(self, da: GridDataArray, mask: GridDataArray return array_mask - def is_regridding_supported(self) -> bool: - """ - returns true if package supports regridding. - """ - return hasattr(self, "_regrid_method") - def regrid_like( self, diff --git a/imod/mf6/ssm.py b/imod/mf6/ssm.py index c12c82f3b..e5f64edac 100644 --- a/imod/mf6/ssm.py +++ b/imod/mf6/ssm.py @@ -5,6 +5,7 @@ from imod.logging import logger from imod.mf6 import GroundwaterFlowModel from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.utilities.regrid import RegridderType from imod.schemata import DTypeSchema @@ -17,7 +18,7 @@ def with_index_dim(array_like): return ("index", arr1d) -class SourceSinkMixing(BoundaryCondition): +class SourceSinkMixing(BoundaryCondition, IRegridPackage): """ Parameters ---------- diff --git a/imod/mf6/utilities/regrid.py b/imod/mf6/utilities/regrid.py index 0bc86688b..c1c4fea0d 100644 --- a/imod/mf6/utilities/regrid.py +++ b/imod/mf6/utilities/regrid.py @@ -186,8 +186,8 @@ def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, lis This function loops over the packages and collects all regridder-types that are in use. """ methods: defaultdict = defaultdict(list) - for pkg_name, pkg in model.items(): - if pkg.is_regridding_supported(): + for _, pkg in model.items(): + if isinstance(pkg, IRegridPackage): pkg_methods = pkg.get_regrid_methods() for variable in pkg_methods: if ( @@ -198,10 +198,6 @@ def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, lis functiontype = pkg_methods[variable][1] if functiontype not in methods[regriddertype]: methods[regriddertype].append(functiontype) - else: - raise NotImplementedError( - f"regridding is not implemented for package {pkg_name} of type {type(pkg)}" - ) return methods @@ -315,7 +311,7 @@ def _regrid_like( new_model = model.__class__() for pkg_name, pkg in model.items(): - if pkg.is_regridding_supported(): + if isinstance(pkg, (IRegridPackage, ILineDataPackage, IPointDataPackage)): new_model[pkg_name] = pkg.regrid_like(target_grid) else: raise NotImplementedError( From e65710589b232c27f2b0842c214d0aef608e2493 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 15:09:25 +0100 Subject: [PATCH 27/29] removed test --- imod/tests/test_mf6/test_mf6_gwfgwf.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/imod/tests/test_mf6/test_mf6_gwfgwf.py b/imod/tests/test_mf6/test_mf6_gwfgwf.py index 96bde4925..e8d3fcefa 100644 --- a/imod/tests/test_mf6/test_mf6_gwfgwf.py +++ b/imod/tests/test_mf6/test_mf6_gwfgwf.py @@ -139,11 +139,6 @@ def test_error_clip(self, sample_gwfgwf_structured: imod.mf6.GWFGWF): with pytest.raises(NotImplementedError): sample_gwfgwf_structured.clip_box(0, 100, 1, 12, 0, 100, 0, 100) - @pytest.mark.usefixtures("sample_gwfgwf_structured") - def test_error_regrid(self, sample_gwfgwf_structured: imod.mf6.GWFGWF): - assert not sample_gwfgwf_structured.is_regridding_supported() - - @pytest.mark.parametrize("newton_option", [False, True]) def test_option_newton_propagated(circle_model, newton_option, tmp_path): # set newton option on original model From d11f5a5f8a8aba1f77cc13f9232255ecc5373220 Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 15:29:02 +0100 Subject: [PATCH 28/29] mypy --- imod/mf6/sto.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/imod/mf6/sto.py b/imod/mf6/sto.py index 4e0cfc59e..f6d382fa7 100644 --- a/imod/mf6/sto.py +++ b/imod/mf6/sto.py @@ -64,8 +64,7 @@ def _render_dict(self, directory, pkgname, globaltimes, binary): d = self.get_options(d) return d - def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: - return self._regrid_method + class SpecificStorage(StorageBase): """ @@ -189,7 +188,9 @@ def render(self, directory, pkgname, globaltimes, binary): d = self._render_dict(directory, pkgname, globaltimes, binary) return self._template.render(d) - + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method + class StorageCoefficient(StorageBase): """ Storage Package with a storage coefficient. Be careful, @@ -323,3 +324,6 @@ def render(self, directory, pkgname, globaltimes, binary): d = self._render_dict(directory, pkgname, globaltimes, binary) d["storagecoefficient"] = True return self._template.render(d) + + def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: + return self._regrid_method \ No newline at end of file From 9264e28f6e97515ba3b7fd05cdb00650ae2ae56b Mon Sep 17 00:00:00 2001 From: luitjan Date: Wed, 13 Mar 2024 15:32:41 +0100 Subject: [PATCH 29/29] ist no longer regriddable --- imod/mf6/ist.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/imod/mf6/ist.py b/imod/mf6/ist.py index 94ad7d0bd..a43678c7c 100644 --- a/imod/mf6/ist.py +++ b/imod/mf6/ist.py @@ -1,11 +1,9 @@ -from typing import Optional, Tuple +from typing import Optional import numpy as np import xarray as xr -from imod.mf6.interfaces.iregridpackage import IRegridPackage from imod.mf6.package import Package -from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import PKG_DIMS_SCHEMA from imod.schemata import ( AllValueSchema, @@ -15,7 +13,7 @@ ) -class ImmobileStorageTransfer(Package, IRegridPackage): +class ImmobileStorageTransfer(Package): """ The Immobile Storage and Transfer (IST) package represents an immobile fraction of groundwater. Any number of IST Packages can be specified for a @@ -253,7 +251,4 @@ def __init__( "format": format, } super().__init__(dict_dataset) - self._validate_init_schemata(validate) - - def get_regrid_methods(self) -> Optional[dict[str, Tuple[RegridderType, str]]]: - return self._regrid_method \ No newline at end of file + self._validate_init_schemata(validate) \ No newline at end of file