Issue #863 transport regridding

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
 ~~~~~~~

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")

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,

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

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

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,

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,

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)
 

imod/mf6/rch.py

@@ -105,6 +105,7 @@ class Recharge(BoundaryCondition):
 
     _regrid_method = {
         "rate": (RegridderType.OVERLAP, "mean"),
+        "concentration": (RegridderType.OVERLAP, "mean"),
     }
 
     def __init__(

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}")
 

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,

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,
     )

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)