diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index b1d9f6c8a..57cb2c2ea 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -43,6 +43,7 @@ Added - :meth:`imod.mf6.Modflow6Simulation.open_concentration` and :meth:`imod.mf6.Modflow6Simulation.open_transport_budget` support opening split multi-species simulations. + :meth:`imod.mf6.Modflow6Simulation.regrid_like` can now regrid simulations that have 1 or more transport models. Changed ~~~~~~~ diff --git a/imod/mf6/adv.py b/imod/mf6/adv.py index e693330eb..2103556cc 100644 --- a/imod/mf6/adv.py +++ b/imod/mf6/adv.py @@ -12,13 +12,15 @@ from copy import deepcopy from imod.mf6.package import Package +from imod.mf6.utilities.regrid import RegridderType class Advection(Package): _pkg_id = "adv" _template = Package._initialize_template(_pkg_id) + _regrid_method: dict[str, tuple[RegridderType, str]] = {} - def __init__(self, scheme): + def __init__(self, scheme: str): dict_dataset = {"scheme": scheme} super().__init__(dict_dataset) @@ -42,7 +44,9 @@ class AdvectionUpstream(Advection): Note: all constructor arguments will be ignored """ - def __init__(self): + def __init__(self, scheme: str = "upstream"): + if not scheme =="upstream": + raise ValueError("error in scheme parameter. Should be 'upstream' if present.") super().__init__(scheme="upstream") @@ -58,7 +62,9 @@ class AdvectionCentral(Advection): Note: all constructor arguments will be ignored """ - def __init__(self): + def __init__(self, scheme: str = "central"): + if not scheme =="central": + raise ValueError("error in scheme parameter. Should be 'central' if present.") super().__init__(scheme="central") @@ -69,5 +75,7 @@ class AdvectionTVD(Advection): Note: all constructor arguments will be ignored """ - def __init__(self): + def __init__(self, scheme: str = "TVD"): + if not scheme =="TVD": + raise ValueError("error in scheme parameter. Should be 'TVD' if present.") super().__init__(scheme="TVD") diff --git a/imod/mf6/dsp.py b/imod/mf6/dsp.py index 5370a3ae1..6ce4c2a22 100644 --- a/imod/mf6/dsp.py +++ b/imod/mf6/dsp.py @@ -1,6 +1,7 @@ import numpy as np 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 ( CompatibleSettingsSchema, @@ -139,6 +140,22 @@ class Dispersion(Package): ), } + _regrid_method = { + "diffusion_coefficient": (RegridderType.OVERLAP, "mean"), + "longitudinal_horizontal": (RegridderType.OVERLAP, "mean"), + "transversal_horizontal1": ( + RegridderType.OVERLAP, + "mean", + ), + "longitudinal_vertical": ( + RegridderType.OVERLAP, + "mean", + ), + "transversal_horizontal2": (RegridderType.OVERLAP, "mean"), + "transversal_vertical": (RegridderType.OVERLAP, "mean"), + } + + def __init__( self, diffusion_coefficient, diff --git a/imod/mf6/exchangebase.py b/imod/mf6/exchangebase.py index 65f4f5337..66d2990a5 100644 --- a/imod/mf6/exchangebase.py +++ b/imod/mf6/exchangebase.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Tuple, Union +from typing import Union import numpy as np import pandas as pd @@ -32,7 +32,7 @@ def model_name2(self) -> str: def package_name(self) -> str: return f"{self.model_name1}_{self.model_name2}" - def get_specification(self) -> Tuple[str, str, str, str]: + def get_specification(self) -> tuple[str, str, str, str]: """ Returns a tuple containing the exchange type, the exchange file name, and the model names. This can be used to write the exchange information in the simulation .nam input file diff --git a/imod/mf6/model.py b/imod/mf6/model.py index ab9b79172..5fdbe2fdc 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -586,5 +586,25 @@ 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}") + 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 diff --git a/imod/mf6/model_gwf.py b/imod/mf6/model_gwf.py index bdb49f934..6a73bbbee 100644 --- a/imod/mf6/model_gwf.py +++ b/imod/mf6/model_gwf.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections import defaultdict from typing import Optional import cftime @@ -9,7 +8,6 @@ from imod.mf6 import ConstantHead from imod.mf6.clipped_boundary_condition_creator import create_clipped_boundary from imod.mf6.model import Modflow6Model -from imod.mf6.utilities.regrid import RegridderType from imod.typing import GridDataArray @@ -38,29 +36,6 @@ def __init__( } - 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, diff --git a/imod/mf6/mst.py b/imod/mf6/mst.py index f46a233d6..2b076f04e 100644 --- a/imod/mf6/mst.py +++ b/imod/mf6/mst.py @@ -1,6 +1,7 @@ import numpy as np 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, @@ -95,6 +96,19 @@ class MobileStorageTransfer(Package): "sp2": (IdentityNoDataSchema(other="idomain", is_other_notnull=(">", 0)),), } + + _regrid_method = { + "porosity": (RegridderType.OVERLAP, "mean"), + "decay": (RegridderType.OVERLAP, "mean"), + "decay_sorbed": ( + RegridderType.OVERLAP, + "mean", + ), + "bulk_density": (RegridderType.OVERLAP, "mean"), + "distcoef": (RegridderType.OVERLAP, "mean"), + "sp2": (RegridderType.OVERLAP, "mean"), + } + def __init__( self, porosity, diff --git a/imod/mf6/package.py b/imod/mf6/package.py index 8e8990724..76df92860 100644 --- a/imod/mf6/package.py +++ b/imod/mf6/package.py @@ -15,7 +15,11 @@ from xarray.core.utils import is_scalar import imod -from imod.mf6.auxiliary_variables import get_variable_names +from imod.mf6.auxiliary_variables import ( + expand_transient_auxiliary_variables, + get_variable_names, + remove_expanded_auxiliary_variables_from_dataset, +) from imod.mf6.interfaces.ipackage import IPackage from imod.mf6.pkgbase import ( EXCHANGE_PACKAGES, @@ -540,12 +544,15 @@ def mask(self, mask: GridDataArray) -> Any: """ masked = {} + if hasattr(self,"auxiliary_data_fields"): + remove_expanded_auxiliary_variables_from_dataset(self) for var in self.dataset.data_vars.keys(): if self._skip_masking_variable(var, self.dataset[var]): masked[var] = self.dataset[var] else: masked[var] = self._mask_spatial_var(var, mask) - + if hasattr(self,"auxiliary_data_fields"): + expand_transient_auxiliary_variables(self) return type(self)(**masked) def _skip_masking_variable(self, var: str, da: GridDataArray)->bool: @@ -692,6 +699,9 @@ def regrid_like( raise NotImplementedError( f"Package {type(self).__name__} does not support regridding" ) + + if hasattr(self,"auxiliary_data_fields"): + remove_expanded_auxiliary_variables_from_dataset(self) regridder_collection = RegridderInstancesCollection( self.dataset, target_grid=target_grid @@ -728,6 +738,8 @@ def regrid_like( new_package_data[varname] = assign_coord_if_present( "dy", target_grid, new_package_data[varname] ) + if hasattr(self,"auxiliary_data_fields"): + expand_transient_auxiliary_variables(self) return self.__class__(**new_package_data) diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index 3727d5c48..6236f23f0 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -105,6 +105,7 @@ class Recharge(BoundaryCondition): _regrid_method = { "rate": (RegridderType.OVERLAP, "mean"), + "concentration": (RegridderType.OVERLAP, "mean"), } def __init__( diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index 895eaadcc..ea9382121 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -1093,12 +1093,14 @@ def regrid_like( ) result = self.__class__(regridded_simulation_name) for key, item in self.items(): - if isinstance(item, GroundwaterFlowModel): + 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}") diff --git a/imod/mf6/ssm.py b/imod/mf6/ssm.py index 35cb8a539..62276d3eb 100644 --- a/imod/mf6/ssm.py +++ b/imod/mf6/ssm.py @@ -1,8 +1,11 @@ +from typing import Tuple + import numpy as np from imod.logging import logger from imod.mf6 import GroundwaterFlowModel from imod.mf6.boundary_condition import BoundaryCondition +from imod.mf6.utilities.regrid import RegridderType from imod.schemata import DTypeSchema @@ -39,6 +42,8 @@ class SourceSinkMixing(BoundaryCondition): _write_schemata = {} + _regrid_method: dict[str, Tuple[RegridderType, str]] = {} + def __init__( self, package_names, diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py index 807453fba..242a0b333 100644 --- a/imod/tests/fixtures/flow_transport_simulation_fixture.py +++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py @@ -71,9 +71,15 @@ def create_transport_model(flow_model, species_name, dispersivity, retardation, # %% -# Create the spatial discretization. @pytest.fixture(scope="function") def flow_transport_simulation(): + """ + This fixture is a variation on the model also present in + examples/mf6/example_1d_transport.py. To make that model more useful for + testing eg partitioning or regridding, some boundary conditions were added + (2 wells, one extractor and one injector which injects with a nonzero + concentration) as well as a recharge zone with a nonzero concentration. + """ nlay = 1 nrow = 2 ncol = 101 @@ -85,7 +91,7 @@ def flow_transport_simulation(): x = np.arange(xmin, xmax, dx) + 0.5 * dx grid_dims = ("layer", "y", "x") - grid_coords = {"layer": layer, "y": y, "x": x} + grid_coords = {"layer": layer, "y": y, "x": x, "dx": dx, "dy": 1.} grid_shape = (nlay, nrow, ncol) grid = xr.DataArray( np.ones(grid_shape, dtype=int), coords=grid_coords, dims=grid_dims @@ -116,9 +122,6 @@ def flow_transport_simulation(): gwf_model["npf"] = imod.mf6.NodePropertyFlow( icelltype=1, k=xr.full_like(grid, 1.0, dtype=float), - variable_vertical_conductance=True, - dewatered=True, - perched=True, ) gwf_model["dis"] = imod.mf6.StructuredDiscretization( top=0.0, @@ -138,7 +141,7 @@ def flow_transport_simulation(): species=["species_a", "species_b", "species_c", "species_d"] ) recharge_rate = xr.full_like(grid, np.nan, dtype=float) - recharge_rate[..., 20:60] = 0.001 + recharge_rate[..., 20:60] = 0.0001 gwf_model["rch"] = imod.mf6.Recharge(recharge_rate, recharge_conc, "AUX") # %% # Create the simulation. @@ -157,7 +160,7 @@ def flow_transport_simulation(): concentration_boundary_type="Aux", screen_top=[0.0, 0.0], screen_bottom=[-1.0, -1.0], - rate=[1.0, -2.0], + rate=[0.1, -0.2], minimum_k=0.0001, concentration=injection_concentration, ) diff --git a/imod/tests/test_mf6/test_mf6_regrid_simulation.py b/imod/tests/test_mf6/test_mf6_regrid_simulation.py index 3861a7448..95f403316 100644 --- a/imod/tests/test_mf6/test_mf6_regrid_simulation.py +++ b/imod/tests/test_mf6/test_mf6_regrid_simulation.py @@ -1,7 +1,6 @@ from pathlib import Path import numpy as np -import pytest import imod from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run @@ -51,27 +50,4 @@ def test_regridded_simulation_has_required_packages( assert isinstance( new_simulation["time_discretization"], imod.mf6.TimeDiscretization ) - assert isinstance(new_simulation["flow"], imod.mf6.GroundwaterFlowModel) - - -def test_simulation_gives_exception_if_transport_model_is_regridded( - unstructured_flow_simulation: imod.mf6.Modflow6Simulation, -): - transport_model = imod.mf6.GroundwaterTransportModel() - transport_model["adv"] = imod.mf6.AdvectionUpstream() - transport_model["dsp"] = imod.mf6.Dispersion( - diffusion_coefficient=0.0, - longitudinal_horizontal=2.3, - transversal_horizontal1=0.0, - xt3d_off=False, - xt3d_rhs=False, - ) - flow_and_transport_simulation = unstructured_flow_simulation - flow_and_transport_simulation["transport"] = transport_model - - finer_idomain = grid_data_unstructured(np.int32, 1, 0.4) - - with pytest.raises(NotImplementedError): - _ = unstructured_flow_simulation.regrid_like( - "regridded_simulation", finer_idomain - ) + assert isinstance(new_simulation["flow"], imod.mf6.GroundwaterFlowModel) \ No newline at end of file diff --git a/imod/tests/test_mf6/test_mf6_regrid_transport.py b/imod/tests/test_mf6/test_mf6_regrid_transport.py new file mode 100644 index 000000000..318420121 --- /dev/null +++ b/imod/tests/test_mf6/test_mf6_regrid_transport.py @@ -0,0 +1,82 @@ +from pathlib import Path +from typing import Tuple + +import numpy as np +import pytest +import xarray as xr + +import imod + + +def create_regridding_idomain(domain: xr.DataArray , ncol: int, nrow: int): + #Create an equidistant structured grid with the same horizontal and vertical + #extent as the input grid, and the same number of layer, but a different + #amount of rows and columns + dx = domain.coords["dx"].values[()] + dy = domain.coords["dy"].values[()] + x_min = domain.coords["x"].values[0] - dx/2 + x_max = domain.coords["x"].values[-1] + dx/2 + y_min = domain.coords["y"].values[-1] - dy/2 + y_max = domain.coords["y"].values[0] + dy/2 + nlayer = domain.coords["layer"][-1] + + column_size = (x_max - x_min) / ncol + row_size = (y_max -y_min) / nrow + + x = np.arange(x_min, x_max, column_size) + column_size/2 + y = np.arange(y_max, y_min, -row_size) - row_size/2 + + new_idomain = xr.DataArray(dims=["layer", "y", "x"], coords={"layer": np.arange(nlayer)+1, "y": y, "x": x, "dx": column_size, "dy": row_size}) + new_idomain.values[:,:,:] =1 + + return new_idomain + +@pytest.mark.parametrize("col_row_dimension",[(101,4 ), (55,3 ), (155,4 )]) +def test_regrid_transport( + tmp_path: Path, + flow_transport_simulation: imod.mf6.Modflow6Simulation, + col_row_dimension: Tuple[int, int] +): + # Run the original simulation + flow_transport_simulation.write( tmp_path/"original", binary=False) + flow_transport_simulation.run() + + #Set up the regridded domain. The original domain is 101 columns * 2 rows * 1 layer + domain = flow_transport_simulation["flow"].domain + other_idomain = create_regridding_idomain(domain, col_row_dimension[0], col_row_dimension[1]) + + new_simulation = flow_transport_simulation.regrid_like( + "regridded_simulation", other_idomain + ) + + # Test that the newly regridded simulation can run + new_simulation.write( tmp_path/"regridded", binary=False) + new_simulation.run() + + # simulation results + conc = flow_transport_simulation.open_concentration(["species_a", "species_b", "species_c", "species_d"]) + regridded_conc = new_simulation.open_concentration(["species_a", "species_b", "species_c", "species_d"]) + + dx = domain.coords["dx"].values[()] + dy = domain.coords["dy"].values[()] + cell_volume = dx * dy * 1 + + new_dx = other_idomain.coords["dx"].values[()] + new_dy = other_idomain.coords["dy"].values[()] + regridded_cell_volume = new_dx * new_dy + + conc = conc.where(conc > -1e29, 0) + regridded_conc = regridded_conc.where(regridded_conc > -1e29, 0) + + #Compute and compare the total mass of the original model and the regridded mode. + #Constants like porosity are not included. + for species in ["species_a", "species_b", "species_c", "species_d"]: + original_mass = np.sum(conc.sel(time=2000, species =species).values * cell_volume) + regridded_mass = np.sum(regridded_conc.sel(time=2000, species =species ).values * regridded_cell_volume) + assert abs((original_mass - regridded_mass))/original_mass < 3e-2 + + + + + +