Deltares · luitjansl · Feb 5, 2024 · Jan 23, 2024 · Jan 24, 2024 · Jan 24, 2024
diff --git a/imod/mf6/exchangebase.py b/imod/mf6/exchangebase.py
@@ -2,7 +2,7 @@
 
 from imod.mf6.package import Package
 
-_pkg_id_to_type = {"gwfgwf": "GWF6-GWF6", "gwfgwt": "GWF6-GWT6"}
+_pkg_id_to_type = {"gwfgwf": "GWF6-GWF6", "gwfgwt": "GWF6-GWT6", "gwtgwt": "GWT6-GWT6"}
 
 
 class ExchangeBase(Package):

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
@@ -0,0 +1,92 @@
+from typing import Optional
+
+import cftime
+import numpy as np
+import xarray as xr
+
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
+from imod.mf6.exchangebase import ExchangeBase
+from imod.mf6.package import Package
+from imod.typing import GridDataArray
+
+
+class GWTGWT(ExchangeBase):
+    """
+    This package is for writing an exchange file, used for splitting up a model
+    into different submodels (that can be solved in parallel). It (usually)
+    is not instantiated by users, but created by the "split" method of the
+    simulation class.
+    """
+
+    _auxiliary_data = {"auxiliary_data": "variable"}
+    _pkg_id = "gwtgwt"
+    _template = Package._initialize_template(_pkg_id)
+
+    def __init__(
+        self,
+        transport_model_id1: str,
+        transport_model_id2: str,
+        flow_model_id1: str,
+        flow_model_id2: str,
+        cell_id1: xr.DataArray,
+        cell_id2: xr.DataArray,
+        layer: xr.DataArray,
+        cl1: xr.DataArray,
+        cl2: xr.DataArray,
+        hwva: xr.DataArray,
+        angldegx: Optional[xr.DataArray] = None,
+        cdist: Optional[xr.DataArray] = None,
+    ):
+        dict_dataset = {
+            "cell_id1": cell_id1,
+            "cell_id2": cell_id2,
+            "layer": layer,
+            "model_name_1": transport_model_id1,
+            "model_name_2": transport_model_id2,
+            "flow_model_name_1": flow_model_id1,
+            "flow_model_name_2": flow_model_id2,
+            "ihc": xr.DataArray(np.ones_like(cl1, dtype=int)),
+            "cl1": cl1,
+            "cl2": cl2,
+            "hwva": hwva,
+        }
+        super().__init__(dict_dataset)
+
+        auxiliary_variables = [var for var in [angldegx, cdist] if var is not None]
+        if auxiliary_variables:
+            self.dataset["auxiliary_data"] = xr.merge(auxiliary_variables).to_array(
+                name="auxiliary_data"
+            )
+            expand_transient_auxiliary_variables(self)
+
+    def set_options(
+        self,
+        print_input: Optional[bool] = None,
+        print_flows: Optional[bool] = None,
+        save_flows: Optional[bool] = None,
+        adv_scheme: Optional[str] = None,
+        dsp_xt3d_off: Optional[bool] = None,
+        dsp_xt3d_rhs: Optional[bool] = None,
+    ):
+        self.dataset["print_input"] = print_input
+        self.dataset["print_flows"] = print_flows
+        self.dataset["save_flows"] = save_flows
+        self.dataset["adv_scheme"] = adv_scheme
+        self.dataset["dsp_xt3d_off"] = dsp_xt3d_off
+        self.dataset["dsp_xt3d_rhs"] = dsp_xt3d_rhs
+
+    def clip_box(
+        self,
+        time_min: Optional[cftime.datetime | np.datetime64 | str] = None,
+        time_max: Optional[cftime.datetime | np.datetime64 | str] = None,
+        layer_min: Optional[int] = None,
+        layer_max: Optional[int] = None,
+        x_min: Optional[float] = None,
+        x_max: Optional[float] = None,
+        y_min: Optional[float] = None,
+        y_max: Optional[float] = None,
+        top: Optional[GridDataArray] = None,
+        bottom: Optional[GridDataArray] = None,
+        state_for_boundary: Optional[GridDataArray] = None,
+    ) -> Package:
+        raise NotImplementedError("this package cannot be clipped")
diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
@@ -6,6 +6,7 @@
 import xarray as xr
 
 from imod.mf6.gwfgwf import GWFGWF
+from imod.mf6.gwtgwt import GWTGWT
 from imod.mf6.multimodel.modelsplitter import PartitionInfo
 from imod.mf6.utilities.grid import get_active_domain_slice, to_cell_idx
 from imod.typing import GridDataArray
@@ -144,50 +145,93 @@ def create_gwfgwf_exchanges(
             for model_id2, connected_domain_pair in grouped_connected_models.groupby(
                 "cell_label2"
             ):
-                model_id1 = int(model_id1)
-                model_id2 = int(model_id2)
-                mapping1 = (
-                    self._global_to_local_mapping[model_id1]
-                    .drop(columns=["local_idx"])
-                    .rename(
-                        columns={"global_idx": "cell_idx1", "local_cell_id": "cell_id1"}
+                connected_cells_dataset = (
+                    self._collect_geometric_constants_connected_cells(
+                        model_id1, model_id2, connected_domain_pair, layers
                     )
                 )
 
-                mapping2 = (
-                    self._global_to_local_mapping[model_id2]
-                    .drop(columns=["local_idx"])
-                    .rename(
-                        columns={"global_idx": "cell_idx2", "local_cell_id": "cell_id2"}
+                exchanges.append(
+                    GWFGWF(
+                        f"{model_name}_{model_id1}",
+                        f"{model_name}_{model_id2}",
+                        **connected_cells_dataset,
                     )
                 )
 
-                connected_cells = (
-                    connected_domain_pair.merge(mapping1)
-                    .merge(mapping2)
-                    .filter(
-                        [
-                            "cell_id1",
-                            "cell_id2",
-                            "cl1",
-                            "cl2",
-                            "hwva",
-                            "angldegx",
-                            "cdist",
-                        ]
-                    )
-                )
+        return exchanges
+
+    def _collect_geometric_constants_connected_cells(
+        self,
+        model_id1: int,
+        model_id2: int,
+        connected_domain_pair: pd.DataFrame,
+        layers: GridDataArray,
+    ) -> xr.Dataset:
+        mapping1 = (
+            self._global_to_local_mapping[model_id1]
+            .drop(columns=["local_idx"])
+            .rename(columns={"global_idx": "cell_idx1", "local_cell_id": "cell_id1"})
+        )
+
+        mapping2 = (
+            self._global_to_local_mapping[model_id2]
+            .drop(columns=["local_idx"])
+            .rename(columns={"global_idx": "cell_idx2", "local_cell_id": "cell_id2"})
+        )
 
-                connected_cells = pd.merge(layers, connected_cells, how="cross")
+        connected_cells = (
+            connected_domain_pair.merge(mapping1)
+            .merge(mapping2)
+            .filter(
+                [
+                    "cell_id1",
+                    "cell_id2",
+                    "cl1",
+                    "cl2",
+                    "hwva",
+                    "angldegx",
+                    "cdist",
+                ]
+            )
+        )
 
-                connected_cells_dataset = self._to_xarray(connected_cells)
+        connected_cells = pd.merge(layers, connected_cells, how="cross")
 
-                _adjust_gridblock_indexing(connected_cells_dataset)
+        connected_cells_dataset = self._to_xarray(connected_cells)
+
+        _adjust_gridblock_indexing(connected_cells_dataset)
+
+        return connected_cells_dataset
+
+    def create_gwtgwt_exchanges(
+        self, transport_model_name: str, flow_model_name: str, layers: GridDataArray
+    ) -> list[GWTGWT]:
+        layers = layers.to_dataframe().filter(["layer"])
 
+        connected_cells_with_geometric_info = pd.merge(
+            self._connected_cells, self._geometric_information
+        )
+
+        exchanges = []
+        for (
+            model_id1,
+            grouped_connected_models,
+        ) in connected_cells_with_geometric_info.groupby("cell_label1"):
+            for model_id2, connected_domain_pair in grouped_connected_models.groupby(
+                "cell_label2"
+            ):
+                connected_cells_dataset = (
+                    self._collect_geometric_constants_connected_cells(
+                        model_id1, model_id2, connected_domain_pair, layers
+                    )
+                )
                 exchanges.append(
-                    GWFGWF(
-                        f"{model_name}_{model_id1}",
-                        f"{model_name}_{model_id2}",
+                    GWTGWT(
+                        f"{transport_model_name}_{model_id1}",
+                        f"{transport_model_name}_{model_id2}",
+                        f"{flow_model_name}_{model_id1}",
+                        f"{flow_model_name}_{model_id2}",
                         **connected_cells_dataset,
                     )
                 )

diff --git a/imod/mf6/pkgbase.py b/imod/mf6/pkgbase.py
@@ -12,7 +12,7 @@
 from imod.typing.grid import GridDataArray, GridDataset, merge_with_dictionary
 
 TRANSPORT_PACKAGES = ("adv", "dsp", "ssm", "mst", "ist", "src")
-EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt")
+EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt", "gwtgwt")
 
 
 class PackageBase(IPackageBase, abc.ABC):

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
@@ -21,6 +21,7 @@
 import imod.mf6.exchangebase
 from imod.mf6.gwfgwf import GWFGWF
 from imod.mf6.gwfgwt import GWFGWT
+from imod.mf6.gwtgwt import GWTGWT
 from imod.mf6.ims import Solution
 from imod.mf6.model import Modflow6Model
 from imod.mf6.model_gwf import GroundwaterFlowModel
@@ -822,7 +823,6 @@ def get_exchange_relationships(self):
         if is_split(self):
             for exchange in self["split_exchanges"]:
                 result.append(exchange.get_specification())
-
         return result
 
     def get_models_of_type(self, modeltype):
@@ -912,6 +912,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 "Unable to split simulation. Splitting can only be done on simulations that haven't been split."
             )
 
+        flow_models = self.get_models_of_type("gwf6")
+        transport_models = self.get_models_of_type("gwt6")
+        if any(transport_models) and len(flow_models) != 1:
+            raise ValueError(
+                "splitting of simulations with more (or less) than 1 flow model currently not supported, if a transport model is present"
+            )
+
+        if not any(flow_models) and not any(transport_models):
+            raise ValueError("a simulation without any models cannot be split.")
+
         original_models = get_models(self)
         original_packages = get_packages(self)
 
@@ -941,12 +951,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 new_simulation[solution_name].add_model_to_solution(new_model_name)
 
         exchanges = []
-        for model_name, model in original_models.items():
-            if isinstance(model, GroundwaterFlowModel):
-                exchanges += exchange_creator.create_gwfgwf_exchanges(
-                    model_name, model.domain.layer
-                )
 
+        for flow_model_name, flow_model in flow_models.items():
+            exchanges += exchange_creator.create_gwfgwf_exchanges(
+                flow_model_name, flow_model.domain.layer
+            )
+        if any(transport_models):
+            for tpt_model_name in transport_models:
+                exchanges += exchange_creator.create_gwtgwt_exchanges(
+                    tpt_model_name, flow_model_name, model.domain.layer
+                )
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()
 
@@ -1014,6 +1028,9 @@ def _set_exchange_options(self):
                     xt3d=model_1["npf"].get_xt3d_option(),
                     newton=model_1.is_use_newton(),
                 )
+            elif isinstance(exchange, GWTGWT):
+                # TODO: issue #747
+                continue
 
     def _filter_inactive_cells_from_exchanges(self) -> None:
         for ex in self["split_exchanges"]:

diff --git a/imod/templates/mf6/exg-gwtgwt.j2 b/imod/templates/mf6/exg-gwtgwt.j2
@@ -0,0 +1,37 @@
+# this file contains the exchanges from model {{model_name_1}} to model {{model_name_2}}
+
+begin options
+  gwfmodelname1 {{flow_model_name_1}}
+  gwfmodelname2 {{flow_model_name_2}}
+{% if auxiliary is defined %}  auxiliary {{auxiliary|join(" ")}}
+{% endif -%}  
+{%- if print_input is defined -%}  print_input
+{% endif -%}
+{%- if print_flows is defined -%}  print_flows
+{% endif -%}
+{%- if save_flows is defined -%}  save_flows
+{% endif -%}
+{%- if adv_scheme is defined -%}  adv_scheme {{adv_scheme}}
+{% endif -%}
+{%- if dsp_xt3d_off is defined -%}  dsp_xt3d_off 
+{% endif -%}
+{%- if dsp_xt3d_rhs is defined -%}  dsp_xt3d_rhs
+{% endif -%}
+{%- if mvt is defined -%}  mvt6 filein {{mvt}}
+{% endif -%}
+{%- if obs is defined -%}  obs6 filein {{obs}}
+{% endif -%}
+end options
+{% set nexg = layer | length %}
+begin dimensions
+  nexg {{nexg}}
+end dimensions
+
+begin exchangedata
+# first 3 (structured) or 2 (unstructured) columns are the exchange boundary cell indices in the numbering local to {{model_name_1}}
+# second 3 (structured) or 2 (unstructured) columns are the exchange boundary cell indices in the numbering local to {{model_name_2}}
+# followed by columns ihc, cl1, cl2, hwva and auxiliary variables, if any
+{%- for i in range(nexg) %}
+  {{layer[i]}} {{cell_id1[i]|join(" ")}} {{layer[i]}} {{cell_id2[i]|join(" ")}} {{ihc[i]}} {{cl1[i]}} {{cl2[i]}} {{hwva[i]}} {%- if auxiliary_data is defined %} {{auxiliary_data.T[i]|join(" ")}}{% endif %}
+{%- endfor %}
+end exchangedata
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -196,16 +196,16 @@ def flow_transport_simulation():
         print_option="summary",
         csv_output=False,
         no_ptc=True,
-        outer_dvclose=1.0e-4,
+        outer_dvclose=1.0e-6,
         outer_maximum=500,
         under_relaxation=None,
-        inner_dvclose=1.0e-4,
-        inner_rclose=0.001,
-        inner_maximum=100,
+        inner_dvclose=1.0e-6,
+        inner_rclose=0.0001,
+        inner_maximum=200,
         linear_acceleration="bicgstab",
         scaling_method=None,
         reordering_method=None,
-        relaxation_factor=0.97,
+        relaxation_factor=0.9,
     )
 
     duration = pd.to_timedelta("2000d")