From fdade8694db5d86fe633aac12013325c34677d34 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 16:29:53 +0100
Subject: [PATCH 01/47] work started

---
 imod/mf6/auxiliary_variables.py               | 13 ++++++++
 imod/mf6/multimodel/modelsplitter.py          |  8 +++++
 .../flow_transport_simulation_fixture.py      | 28 ++++++++++++++++-
 imod/tests/fixtures/mf6_modelrun_fixture.py   | 31 +++++++++----------
 .../test_multimodel/test_mf6_modelsplitter.py | 31 ++++++++++++++++++-
 5 files changed, 93 insertions(+), 18 deletions(-)

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index 50bc05524..f9a24ed6d 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -33,3 +33,16 @@ def add_periodic_auxiliary_variable(package):
                 package.dataset[s] = package.dataset[aux_var_name].sel(
                     {aux_var_dimensions: s}
                 )
+
+def remove_periodic_auxiliary_variable(package):
+    if "species" in package.dataset.coords:
+        for species in package.dataset.coords["species"].values:
+            if species in list(package.dataset.keys()):
+                package.dataset = package.dataset.drop_vars(species)
+            
+def has_auxiliary_variable(package):
+    if hasattr(package, "_auxiliary_data"):
+        for aux_var_name, _ in package._auxiliary_data.items():
+            if aux_var_name in package.dataset.keys():
+                return True
+    return False
\ No newline at end of file
diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py
index 830d492c0..c5d138ac8 100644
--- a/imod/mf6/multimodel/modelsplitter.py
+++ b/imod/mf6/multimodel/modelsplitter.py
@@ -9,6 +9,8 @@
 from imod.mf6.wel import Well
 from imod.typing import GridDataArray
 from imod.typing.grid import is_unstructured, ones_like
+from imod.mf6.boundary_condition import BoundaryCondition
+from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable, remove_periodic_auxiliary_variable, has_auxiliary_variable
 
 HIGH_LEVEL_PKGS = (HorizontalFlowBarrierBase, Well)
 
@@ -85,6 +87,10 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
         new_idomain = model.domain.sel(coords).where(sliced_domain_2D, other=0)
 
     for pkg_name, package in model.items():
+        if issubclass(type(package), BoundaryCondition):
+            if has_auxiliary_variable(package):
+                remove_periodic_auxiliary_variable(package)
+
         sliced_package = clip_by_grid(package, partition_info.active_domain)
 
         sliced_package = sliced_package.mask(new_idomain)
@@ -98,4 +104,6 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
                 f"package {pkg_name} removed in partition {partition_info.id}, because all empty"
             )
 
+        if issubclass(type(package), BoundaryCondition):
+           add_periodic_auxiliary_variable( package )
     return new_model
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index 7a2309805..7d19f7c25 100644
--- a/imod/tests/fixtures/flow_transport_simulation_fixture.py
+++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -132,9 +132,35 @@ def flow_transport_simulation():
         transient=False,
         convertible=0,
     )
+    recharge_conc = xr.full_like(grid, np.nan, dtype=float)
+    recharge_conc[..., 20:60] = 0.001
+    recharge_conc = recharge_conc.expand_dims(
+        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
+    gwf_model["rch"] = imod.mf6.Recharge(recharge_rate,recharge_conc,"AUX" )
     # %%
     # Create the simulation.
 
+    rate = [1.0, 1.0]
+    injection_concentration = xr.DataArray(
+        [[0.2, 0.23],[ 0.5, 0.2], [0.2, 0.23], [0.5, 0.2]],
+        coords={"species":["species_a", "species_b", "species_c", "species_d"], "index":[0,1]},
+        dims=("species","index"),
+    )
+
+    gwf_model["well"] = imod.mf6.Well(
+        x=[20., 80.],
+        y=[0.6, 1.2],
+        concentration_boundary_type = "Aux",
+        screen_top=[0., 0.],
+        screen_bottom=[-1., -1.],
+        rate=[1.0, -2.0],
+        minimum_k=0.0001,
+        concentration =injection_concentration
+    )
+
     simulation = imod.mf6.Modflow6Simulation("1d_tpt_benchmark")
     simulation["flow"] = gwf_model
 
@@ -148,7 +174,7 @@ def flow_transport_simulation():
         gwf_model, "species_d", 10.0, 5.0, 0.002
     )
 
-    simulation["flow_solver"] = imod.mf6.Solution(
+    simulation["solver"] = imod.mf6.Solution(
         modelnames=["flow"],
         print_option="summary",
         csv_output=False,
diff --git a/imod/tests/fixtures/mf6_modelrun_fixture.py b/imod/tests/fixtures/mf6_modelrun_fixture.py
index ef15549e4..599213232 100644
--- a/imod/tests/fixtures/mf6_modelrun_fixture.py
+++ b/imod/tests/fixtures/mf6_modelrun_fixture.py
@@ -25,25 +25,24 @@ def assert_simulation_can_run(
     flow_model_names = [
         m for m in simulation if type(simulation[m]) is imod.mf6.GroundwaterFlowModel
     ]
-    assert len(flow_model_names) == 1
-    flowmodel = flow_model_names[0]
+    for flowmodel in flow_model_names:
 
-    # read the output  generated by modflow
-    dis_outputfile = (
-        modeldir / flowmodel / f"{discretization_name}.{discretization_name}.grb"
-    )
-    head = imod.mf6.out.open_hds(
-        modeldir / flowmodel / f"{flowmodel}.hds",
-        dis_outputfile,
-    )
+        # read the output  generated by modflow
+        dis_outputfile = (
+            modeldir / flowmodel / f"{discretization_name}.{discretization_name}.grb"
+        )
+        head = imod.mf6.out.open_hds(
+            modeldir / flowmodel / f"{flowmodel}.hds",
+            dis_outputfile,
+        )
 
-    # filter output on idomain
-    idomain = simulation[flowmodel].domain
-    head = head.reindex_like(idomain, "nearest", 1e-5)
-    head = head.where(idomain == 1, other=0)
+        # filter output on idomain
+        idomain = simulation[flowmodel].domain
+        head = head.reindex_like(idomain, "nearest", 1e-5)
+        head = head.where(idomain == 1, other=0)
 
-    # Test that heads are not nan
-    assert not np.any(isnull(head.values))
+        # Test that heads are not nan
+        assert not np.any(isnull(head.values))
 
 
 def assert_model_can_run(
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 7522812f7..13ddca3f3 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -3,7 +3,7 @@
 
 from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
 from imod.typing.grid import zeros_like
-
+from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 
 @pytest.mark.usefixtures("twri_model")
 def test_slice_model_structured(twri_model):
@@ -71,3 +71,32 @@ def test_slice_model_unstructured(circle_model):
     for submodel_label in unique_labels:
         active_count = new_models[submodel_label].domain.sel({"layer": 1}).count()
         assert label_counts[submodel_label] == active_count.values
+
+@pytest.mark.usefixtures("flow_transport_simulation")
+def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulation):
+    
+    flow_simulation = flow_transport_simulation
+    flow_simulation.pop("tpt_a")
+    flow_simulation.pop("tpt_b")
+    flow_simulation.pop("tpt_c")
+    flow_simulation.pop("tpt_d")    
+    flow_simulation.pop("transport_solver")   
+
+    flow_model = flow_simulation["flow"]   
+    active = flow_model.domain
+
+    submodel_labels = zeros_like(active)
+    submodel_labels = submodel_labels.drop_vars("layer")    
+    submodel_labels.values[:,:,50:] = 1
+    submodel_labels = submodel_labels.sel(layer=0, drop=True)
+
+    flow_simulation.write(tmp_path, binary=False)
+    split_simulation = flow_simulation.split(submodel_labels)
+    split_simulation.write(tmp_path, binary=False)
+    assert_simulation_can_run(  split_simulation, "dis", tmp_path)
+    pass
+
+
+
+    
+

From fb61e4d152fa199274925f5a33c6c90cef593586 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 16:37:36 +0100
Subject: [PATCH 02/47] now works with well

---
 imod/tests/fixtures/flow_transport_simulation_fixture.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index 7d19f7c25..0b8609c5e 100644
--- a/imod/tests/fixtures/flow_transport_simulation_fixture.py
+++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -151,7 +151,7 @@ def flow_transport_simulation():
     )
 
     gwf_model["well"] = imod.mf6.Well(
-        x=[20., 80.],
+        x=[20., 580.],
         y=[0.6, 1.2],
         concentration_boundary_type = "Aux",
         screen_top=[0., 0.],

From c70e696d29d50f6b79984605b36fd3ffab2c8310 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 17:16:42 +0100
Subject: [PATCH 03/47] formatting

---
 imod/mf6/auxiliary_variables.py               | 39 ++++++++++++++-----
 .../flow_transport_simulation_fixture.py      | 21 +++++-----
 .../test_multimodel/test_mf6_modelsplitter.py | 25 ++++++------
 3 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index f9a24ed6d..7e00326fc 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -1,3 +1,6 @@
+from imod.mf6.package import Package
+
+
 def get_variable_names(package):
     auxiliaries = _get_auxiliary_data_variable_names_mapping(
         package
@@ -18,14 +21,19 @@ def get_variable_names(package):
     return []
 
 
-def _get_auxiliary_data_variable_names_mapping(package):
+def _get_auxiliary_data_variable_names_mapping(package: Package):
     result = {}
     if hasattr(package, "_auxiliary_data"):
         result.update(package._auxiliary_data)
     return result
 
 
-def add_periodic_auxiliary_variable(package):
+def add_periodic_auxiliary_variable(package: Package) -> None:
+    """
+    splits an auxiliary dataarray (with one or more auxiliary variable dimension) into dataarrays per
+    auxiliary variable dimension. For example a concentration auxiliary variable in a flow package
+    will have a species dimension, and will be split in several dataarrays- one for each species.
+    """
     if hasattr(package, "_auxiliary_data"):
         for aux_var_name, aux_var_dimensions in package._auxiliary_data.items():
             aux_coords = package.dataset[aux_var_name].coords[aux_var_dimensions].values
@@ -34,15 +42,26 @@ def add_periodic_auxiliary_variable(package):
                     {aux_var_dimensions: s}
                 )
 
-def remove_periodic_auxiliary_variable(package):
-    if "species" in package.dataset.coords:
-        for species in package.dataset.coords["species"].values:
-            if species in list(package.dataset.keys()):
-                package.dataset = package.dataset.drop_vars(species)
-            
-def has_auxiliary_variable(package):
+
+def remove_periodic_auxiliary_variable(package: Package) -> None:
+    """
+    removes the data arrays created by add_periodic_auxiliary_variable(...) but does not
+    remove the auxiliary dataarray used as source for add_periodic_auxiliary_variable(...)
+    """
+    if hasattr(package, "_auxiliary_data"):
+        for aux_var_name, aux_var_dimensions in package._auxiliary_data.items():
+            if aux_var_dimensions in package.dataset.coords:
+                for species in package.dataset.coords[aux_var_dimensions].values:
+                    if species in list(package.dataset.keys()):
+                        package.dataset = package.dataset.drop_vars(species)
+
+
+def has_auxiliary_variable(package: Package) -> bool:
+    """
+    returns True if a package contains auxiliary data
+    """
     if hasattr(package, "_auxiliary_data"):
         for aux_var_name, _ in package._auxiliary_data.items():
             if aux_var_name in package.dataset.keys():
                 return True
-    return False
\ No newline at end of file
+    return False
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index 0b8609c5e..422a9cd87 100644
--- a/imod/tests/fixtures/flow_transport_simulation_fixture.py
+++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -139,26 +139,29 @@ def flow_transport_simulation():
     )
     recharge_rate = xr.full_like(grid, np.nan, dtype=float)
     recharge_rate[..., 20:60] = 0.001
-    gwf_model["rch"] = imod.mf6.Recharge(recharge_rate,recharge_conc,"AUX" )
+    gwf_model["rch"] = imod.mf6.Recharge(recharge_rate, recharge_conc, "AUX")
     # %%
     # Create the simulation.
 
     rate = [1.0, 1.0]
     injection_concentration = xr.DataArray(
-        [[0.2, 0.23],[ 0.5, 0.2], [0.2, 0.23], [0.5, 0.2]],
-        coords={"species":["species_a", "species_b", "species_c", "species_d"], "index":[0,1]},
-        dims=("species","index"),
+        [[0.2, 0.23], [0.5, 0.2], [0.2, 0.23], [0.5, 0.2]],
+        coords={
+            "species": ["species_a", "species_b", "species_c", "species_d"],
+            "index": [0, 1],
+        },
+        dims=("species", "index"),
     )
 
     gwf_model["well"] = imod.mf6.Well(
-        x=[20., 580.],
+        x=[20.0, 580.0],
         y=[0.6, 1.2],
-        concentration_boundary_type = "Aux",
-        screen_top=[0., 0.],
-        screen_bottom=[-1., -1.],
+        concentration_boundary_type="Aux",
+        screen_top=[0.0, 0.0],
+        screen_bottom=[-1.0, -1.0],
         rate=[1.0, -2.0],
         minimum_k=0.0001,
-        concentration =injection_concentration
+        concentration=injection_concentration,
     )
 
     simulation = imod.mf6.Modflow6Simulation("1d_tpt_benchmark")
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 13ddca3f3..6721a5f38 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -1,9 +1,11 @@
 import numpy as np
 import pytest
 
-from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
-from imod.typing.grid import zeros_like
+from imod.mf6.multimodel.modelsplitter import (create_partition_info,
+                                               slice_model)
 from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
+from imod.typing.grid import zeros_like
+
 
 @pytest.mark.usefixtures("twri_model")
 def test_slice_model_structured(twri_model):
@@ -72,31 +74,26 @@ def test_slice_model_unstructured(circle_model):
         active_count = new_models[submodel_label].domain.sel({"layer": 1}).count()
         assert label_counts[submodel_label] == active_count.values
 
+
 @pytest.mark.usefixtures("flow_transport_simulation")
 def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulation):
-    
     flow_simulation = flow_transport_simulation
     flow_simulation.pop("tpt_a")
     flow_simulation.pop("tpt_b")
     flow_simulation.pop("tpt_c")
-    flow_simulation.pop("tpt_d")    
-    flow_simulation.pop("transport_solver")   
+    flow_simulation.pop("tpt_d")
+    flow_simulation.pop("transport_solver")
 
-    flow_model = flow_simulation["flow"]   
+    flow_model = flow_simulation["flow"]
     active = flow_model.domain
 
     submodel_labels = zeros_like(active)
-    submodel_labels = submodel_labels.drop_vars("layer")    
-    submodel_labels.values[:,:,50:] = 1
+    submodel_labels = submodel_labels.drop_vars("layer")
+    submodel_labels.values[:, :, 50:] = 1
     submodel_labels = submodel_labels.sel(layer=0, drop=True)
 
     flow_simulation.write(tmp_path, binary=False)
     split_simulation = flow_simulation.split(submodel_labels)
     split_simulation.write(tmp_path, binary=False)
-    assert_simulation_can_run(  split_simulation, "dis", tmp_path)
+    assert_simulation_can_run(split_simulation, "dis", tmp_path)
     pass
-
-
-
-    
-

From 2193d3c77972505cd9b4b54ce18c20cfa608c397 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 17:17:40 +0100
Subject: [PATCH 04/47] formatting

---
 imod/mf6/multimodel/modelsplitter.py                      | 8 ++++++--
 imod/tests/fixtures/mf6_modelrun_fixture.py               | 1 -
 .../test_mf6/test_multimodel/test_mf6_modelsplitter.py    | 3 +--
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py
index c5d138ac8..9b0b94afe 100644
--- a/imod/mf6/multimodel/modelsplitter.py
+++ b/imod/mf6/multimodel/modelsplitter.py
@@ -10,7 +10,11 @@
 from imod.typing import GridDataArray
 from imod.typing.grid import is_unstructured, ones_like
 from imod.mf6.boundary_condition import BoundaryCondition
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable, remove_periodic_auxiliary_variable, has_auxiliary_variable
+from imod.mf6.auxiliary_variables import (
+    add_periodic_auxiliary_variable,
+    remove_periodic_auxiliary_variable,
+    has_auxiliary_variable,
+)
 
 HIGH_LEVEL_PKGS = (HorizontalFlowBarrierBase, Well)
 
@@ -105,5 +109,5 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
             )
 
         if issubclass(type(package), BoundaryCondition):
-           add_periodic_auxiliary_variable( package )
+            add_periodic_auxiliary_variable(package)
     return new_model
diff --git a/imod/tests/fixtures/mf6_modelrun_fixture.py b/imod/tests/fixtures/mf6_modelrun_fixture.py
index 599213232..acaced6b4 100644
--- a/imod/tests/fixtures/mf6_modelrun_fixture.py
+++ b/imod/tests/fixtures/mf6_modelrun_fixture.py
@@ -26,7 +26,6 @@ def assert_simulation_can_run(
         m for m in simulation if type(simulation[m]) is imod.mf6.GroundwaterFlowModel
     ]
     for flowmodel in flow_model_names:
-
         # read the output  generated by modflow
         dis_outputfile = (
             modeldir / flowmodel / f"{discretization_name}.{discretization_name}.grb"
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 6721a5f38..87dab8066 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -1,8 +1,7 @@
 import numpy as np
 import pytest
 
-from imod.mf6.multimodel.modelsplitter import (create_partition_info,
-                                               slice_model)
+from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
 from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 from imod.typing.grid import zeros_like
 

From feef11aacf4d7072a98eb1989010110875c78b39 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 17:45:15 +0100
Subject: [PATCH 05/47] formatting and fixing tests

---
 imod/mf6/auxiliary_variables.py      | 13 +++++--------
 imod/mf6/multimodel/modelsplitter.py | 11 ++++++-----
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index 7e00326fc..fc79d8015 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -1,6 +1,3 @@
-from imod.mf6.package import Package
-
-
 def get_variable_names(package):
     auxiliaries = _get_auxiliary_data_variable_names_mapping(
         package
@@ -21,14 +18,14 @@ def get_variable_names(package):
     return []
 
 
-def _get_auxiliary_data_variable_names_mapping(package: Package):
+def _get_auxiliary_data_variable_names_mapping(package):
     result = {}
     if hasattr(package, "_auxiliary_data"):
         result.update(package._auxiliary_data)
     return result
 
 
-def add_periodic_auxiliary_variable(package: Package) -> None:
+def add_periodic_auxiliary_variable(package) -> None:
     """
     splits an auxiliary dataarray (with one or more auxiliary variable dimension) into dataarrays per
     auxiliary variable dimension. For example a concentration auxiliary variable in a flow package
@@ -43,7 +40,7 @@ def add_periodic_auxiliary_variable(package: Package) -> None:
                 )
 
 
-def remove_periodic_auxiliary_variable(package: Package) -> None:
+def remove_periodic_auxiliary_variable(package) -> None:
     """
     removes the data arrays created by add_periodic_auxiliary_variable(...) but does not
     remove the auxiliary dataarray used as source for add_periodic_auxiliary_variable(...)
@@ -56,9 +53,9 @@ def remove_periodic_auxiliary_variable(package: Package) -> None:
                         package.dataset = package.dataset.drop_vars(species)
 
 
-def has_auxiliary_variable(package: Package) -> bool:
+def has_auxiliary_variable_source_array(package) -> bool:
     """
-    returns True if a package contains auxiliary data
+    returns True if a package contains an auxiliary data source data array
     """
     if hasattr(package, "_auxiliary_data"):
         for aux_var_name, _ in package._auxiliary_data.items():
diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py
index cd2737e42..9f2bd09ec 100644
--- a/imod/mf6/multimodel/modelsplitter.py
+++ b/imod/mf6/multimodel/modelsplitter.py
@@ -14,7 +14,7 @@
 from imod.mf6.auxiliary_variables import (
     add_periodic_auxiliary_variable,
     remove_periodic_auxiliary_variable,
-    has_auxiliary_variable,
+    has_auxiliary_variable_source_array,
 )
 
 HIGH_LEVEL_PKGS = (HorizontalFlowBarrierBase, Well)
@@ -92,8 +92,8 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
         new_idomain = model.domain.sel(coords).where(sliced_domain_2D, other=0)
 
     for pkg_name, package in model.items():
-        if issubclass(type(package), BoundaryCondition):
-            if has_auxiliary_variable(package):
+        if isinstance(package, BoundaryCondition):
+            if has_auxiliary_variable_source_array(package):
                 remove_periodic_auxiliary_variable(package)
 
         sliced_package = clip_by_grid(package, partition_info.active_domain)
@@ -109,6 +109,7 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
                 f"package {pkg_name} removed in partition {partition_info.id}, because all empty"
             )
 
-        if issubclass(type(package), BoundaryCondition):
-            add_periodic_auxiliary_variable(package)
+        if isinstance(package, BoundaryCondition):
+            if has_auxiliary_variable_source_array(package):
+                add_periodic_auxiliary_variable(package)
     return new_model

From 385f312b9bdb196112cd804a4cbbfb3c5338f48c Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 18:17:41 +0100
Subject: [PATCH 06/47] expanded test

---
 .../test_mf6/test_multimodel/test_mf6_modelsplitter.py      | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 87dab8066..0f8c656a6 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -95,4 +95,10 @@ def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulatio
     split_simulation = flow_simulation.split(submodel_labels)
     split_simulation.write(tmp_path, binary=False)
     assert_simulation_can_run(split_simulation, "dis", tmp_path)
+    assert "species_d" in list(split_simulation["flow_0"]["chd"].dataset.keys())
+    assert "species_d" in list(split_simulation["flow_0"]["rch"].dataset.keys())
+    assert "concentration" in list(split_simulation["flow_0"]["well"].dataset.keys())    
+    assert "species_d" in list(split_simulation["flow_1"]["chd"].dataset.keys())
+    assert "species_d" in list(split_simulation["flow_1"]["rch"].dataset.keys())
+    assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())    
     pass

From 707aad826ac32dcdc7f56e97a77f81147c644ff3 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 24 Jan 2024 18:18:02 +0100
Subject: [PATCH 07/47] formatting

---
 imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 0f8c656a6..987c02b30 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -97,8 +97,8 @@ def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulatio
     assert_simulation_can_run(split_simulation, "dis", tmp_path)
     assert "species_d" in list(split_simulation["flow_0"]["chd"].dataset.keys())
     assert "species_d" in list(split_simulation["flow_0"]["rch"].dataset.keys())
-    assert "concentration" in list(split_simulation["flow_0"]["well"].dataset.keys())    
+    assert "concentration" in list(split_simulation["flow_0"]["well"].dataset.keys())
     assert "species_d" in list(split_simulation["flow_1"]["chd"].dataset.keys())
     assert "species_d" in list(split_simulation["flow_1"]["rch"].dataset.keys())
-    assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())    
+    assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())
     pass

From 1327e0fd8917b2b4f99aaedccf7853a3d1592303 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 25 Jan 2024 10:42:27 +0100
Subject: [PATCH 08/47] ruff checks

---
 imod/mf6/multimodel/modelsplitter.py                 | 12 ++++++------
 .../fixtures/flow_transport_simulation_fixture.py    |  2 --
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py
index 9f2bd09ec..121599df9 100644
--- a/imod/mf6/multimodel/modelsplitter.py
+++ b/imod/mf6/multimodel/modelsplitter.py
@@ -2,6 +2,12 @@
 
 import numpy as np
 
+from imod.mf6.auxiliary_variables import (
+    add_periodic_auxiliary_variable,
+    has_auxiliary_variable_source_array,
+    remove_periodic_auxiliary_variable,
+)
+from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.hfb import HorizontalFlowBarrierBase
 from imod.mf6.model import Modflow6Model
 from imod.mf6.model_gwf import GroundwaterFlowModel
@@ -10,12 +16,6 @@
 from imod.mf6.wel import Well
 from imod.typing import GridDataArray
 from imod.typing.grid import is_unstructured, ones_like
-from imod.mf6.boundary_condition import BoundaryCondition
-from imod.mf6.auxiliary_variables import (
-    add_periodic_auxiliary_variable,
-    remove_periodic_auxiliary_variable,
-    has_auxiliary_variable_source_array,
-)
 
 HIGH_LEVEL_PKGS = (HorizontalFlowBarrierBase, Well)
 
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index 422a9cd87..d4d58a41d 100644
--- a/imod/tests/fixtures/flow_transport_simulation_fixture.py
+++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -142,8 +142,6 @@ def flow_transport_simulation():
     gwf_model["rch"] = imod.mf6.Recharge(recharge_rate, recharge_conc, "AUX")
     # %%
     # Create the simulation.
-
-    rate = [1.0, 1.0]
     injection_concentration = xr.DataArray(
         [[0.2, 0.23], [0.5, 0.2], [0.2, 0.23], [0.5, 0.2]],
         coords={

From 70df416d2e944c52d087fe0fc8ff9924eb8b8d26 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 25 Jan 2024 15:01:14 +0100
Subject: [PATCH 09/47] review comments

---
 imod/mf6/auxiliary_variables.py | 52 ++++++++++++++++-----------------
 imod/mf6/interfaces/ipackage.py | 18 ++++++++++++
 imod/mf6/package.py             |  9 +++++-
 3 files changed, 52 insertions(+), 27 deletions(-)
 create mode 100644 imod/mf6/interfaces/ipackage.py

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index fc79d8015..d0a3e1866 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -1,4 +1,6 @@
-def get_variable_names(package):
+from imod.mf6.interfaces.ipackage import IPackage
+
+def get_variable_names(package: IPackage):
     auxiliaries = _get_auxiliary_data_variable_names_mapping(
         package
     )  # returns something like {"concentration": "species"}
@@ -18,47 +20,45 @@ def get_variable_names(package):
     return []
 
 
-def _get_auxiliary_data_variable_names_mapping(package):
+def _get_auxiliary_data_variable_names_mapping(package: IPackage):
     result = {}
     if hasattr(package, "_auxiliary_data"):
         result.update(package._auxiliary_data)
     return result
 
 
-def add_periodic_auxiliary_variable(package) -> None:
+def add_periodic_auxiliary_variable(package: IPackage) -> None:
     """
-    splits an auxiliary dataarray (with one or more auxiliary variable dimension) into dataarrays per
+    Splits an auxiliary dataarray (with one or more auxiliary variable dimension) into dataarrays per
     auxiliary variable dimension. For example a concentration auxiliary variable in a flow package
     will have a species dimension, and will be split in several dataarrays- one for each species.
     """
-    if hasattr(package, "_auxiliary_data"):
-        for aux_var_name, aux_var_dimensions in package._auxiliary_data.items():
-            aux_coords = package.dataset[aux_var_name].coords[aux_var_dimensions].values
-            for s in aux_coords:
-                package.dataset[s] = package.dataset[aux_var_name].sel(
-                    {aux_var_dimensions: s}
-                )
 
+    for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields().items():
+        aux_coords = package.dataset[aux_var_name].coords[aux_var_dimensions].values
+        for s in aux_coords:
+            package.dataset[s] = package.dataset[aux_var_name].sel(
+                {aux_var_dimensions: s}
+            )
 
-def remove_periodic_auxiliary_variable(package) -> None:
+
+def remove_periodic_auxiliary_variable(package: IPackage) -> None:
     """
-    removes the data arrays created by add_periodic_auxiliary_variable(...) but does not
-    remove the auxiliary dataarray used as source for add_periodic_auxiliary_variable(...)
+    Removes the data arrays created by :meth:add_periodic_auxiliary_variable(...) but does not
+    remove the auxiliary dataarray used as source for :meth:add_periodic_auxiliary_variable(...)
     """
-    if hasattr(package, "_auxiliary_data"):
-        for aux_var_name, aux_var_dimensions in package._auxiliary_data.items():
-            if aux_var_dimensions in package.dataset.coords:
-                for species in package.dataset.coords[aux_var_dimensions].values:
-                    if species in list(package.dataset.keys()):
-                        package.dataset = package.dataset.drop_vars(species)
+    for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields().items():
+        if aux_var_dimensions in package.dataset.coords:
+            for species in package.dataset.coords[aux_var_dimensions].values:
+                if species in list(package.dataset.keys()):
+                    package.dataset = package.dataset.drop_vars(species)
 
 
-def has_auxiliary_variable_source_array(package) -> bool:
+def has_auxiliary_variable_source_array(package: IPackage) -> bool:
     """
-    returns True if a package contains an auxiliary data source data array
+    Returns True if a package contains an auxiliary data source data array
     """
-    if hasattr(package, "_auxiliary_data"):
-        for aux_var_name, _ in package._auxiliary_data.items():
-            if aux_var_name in package.dataset.keys():
-                return True
+    for aux_var_name, _ in package.auxiliary_data_fields().items():
+        if aux_var_name in package.dataset.keys():
+            return True
     return False
diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py
new file mode 100644
index 000000000..3f11f9d20
--- /dev/null
+++ b/imod/mf6/interfaces/ipackage.py
@@ -0,0 +1,18 @@
+from imod.mf6.interfaces.ipackagebase import IPackageBase
+
+import abc
+from typing import Dict
+class IPackage(IPackageBase, metaclass=abc.ABCMeta):
+    
+
+
+    """
+    The base methods and attributes available in all packages
+    """
+
+    dataset: None
+
+    @property
+    @abc.abstractmethod
+    def auxiliary_data_fields(self) -> Dict[str, str]:
+        raise NotImplementedError
\ No newline at end of file
diff --git a/imod/mf6/package.py b/imod/mf6/package.py
index 6d1342cf8..c6f72daf1 100644
--- a/imod/mf6/package.py
+++ b/imod/mf6/package.py
@@ -32,9 +32,10 @@
     ValidationError,
 )
 from imod.typing import GridDataArray
+from imod.mf6.interfaces.ipackage import IPackage
 
 
-class Package(PackageBase, abc.ABC):
+class Package(PackageBase, IPackage, abc.ABC):
     """
     Package is used to share methods for specific packages with no time
     component.
@@ -729,3 +730,9 @@ def __repr__(self) -> str:
     def _repr_html_(self) -> str:
         typename = type(self).__name__
         return f"<div>{typename}</div>{self.dataset._repr_html_()}"
+
+    def auxiliary_data_fields(self) -> Dict[str, str]:
+        if (hasattr(self, "_auxiliary_data")):
+            return self._auxiliary_data
+        else: 
+            return {}
\ No newline at end of file

From 40e4f7ae6bc41f49d77984a5a8a2a847e77d0fa6 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 25 Jan 2024 15:32:18 +0100
Subject: [PATCH 10/47] review comments

---
 imod/mf6/auxiliary_variables.py |  1 +
 imod/mf6/interfaces/ipackage.py | 11 +++++------
 imod/mf6/package.py             |  7 +++----
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index d0a3e1866..6312741b6 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -1,5 +1,6 @@
 from imod.mf6.interfaces.ipackage import IPackage
 
+
 def get_variable_names(package: IPackage):
     auxiliaries = _get_auxiliary_data_variable_names_mapping(
         package
diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py
index 3f11f9d20..327892ca1 100644
--- a/imod/mf6/interfaces/ipackage.py
+++ b/imod/mf6/interfaces/ipackage.py
@@ -1,18 +1,17 @@
-from imod.mf6.interfaces.ipackagebase import IPackageBase
-
 import abc
 from typing import Dict
-class IPackage(IPackageBase, metaclass=abc.ABCMeta):
-    
 
+from imod.mf6.interfaces.ipackagebase import IPackageBase
+
+
+class IPackage(IPackageBase, metaclass=abc.ABCMeta):
 
     """
     The base methods and attributes available in all packages
     """
 
-    dataset: None
 
     @property
     @abc.abstractmethod
     def auxiliary_data_fields(self) -> Dict[str, str]:
-        raise NotImplementedError
\ No newline at end of file
+        raise NotImplementedError
diff --git a/imod/mf6/package.py b/imod/mf6/package.py
index c6f72daf1..ed0ddcd6d 100644
--- a/imod/mf6/package.py
+++ b/imod/mf6/package.py
@@ -16,6 +16,7 @@
 
 import imod
 from imod.mf6.auxiliary_variables import get_variable_names
+from imod.mf6.interfaces.ipackage import IPackage
 from imod.mf6.pkgbase import EXCHANGE_PACKAGES, TRANSPORT_PACKAGES, PackageBase
 from imod.mf6.regridding_utils import (
     RegridderInstancesCollection,
@@ -32,7 +33,6 @@
     ValidationError,
 )
 from imod.typing import GridDataArray
-from imod.mf6.interfaces.ipackage import IPackage
 
 
 class Package(PackageBase, IPackage, abc.ABC):
@@ -732,7 +732,6 @@ def _repr_html_(self) -> str:
         return f"<div>{typename}</div>{self.dataset._repr_html_()}"
 
     def auxiliary_data_fields(self) -> Dict[str, str]:
-        if (hasattr(self, "_auxiliary_data")):
+        if hasattr(self, "_auxiliary_data"):
             return self._auxiliary_data
-        else: 
-            return {}
\ No newline at end of file
+        return {}

From 4bd7f4bb902194cd6a857cc413b9163b7185f4e2 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 25 Jan 2024 15:32:42 +0100
Subject: [PATCH 11/47] formatting

---
 imod/mf6/interfaces/ipackage.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/imod/mf6/interfaces/ipackage.py b/imod/mf6/interfaces/ipackage.py
index 327892ca1..c2004557b 100644
--- a/imod/mf6/interfaces/ipackage.py
+++ b/imod/mf6/interfaces/ipackage.py
@@ -10,7 +10,6 @@ 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]:

From 0827789e96d4adbde97985259d80b9081dab367b Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 25 Jan 2024 16:47:18 +0100
Subject: [PATCH 12/47] review comments

---
 imod/mf6/auxiliary_variables.py               | 42 +++++++++----------
 imod/mf6/chd.py                               |  4 +-
 imod/mf6/drn.py                               |  4 +-
 imod/mf6/evt.py                               |  4 +-
 imod/mf6/ghb.py                               |  4 +-
 imod/mf6/gwfgwf.py                            |  4 +-
 imod/mf6/mf6_wel_adapter.py                   |  4 +-
 imod/mf6/multimodel/modelsplitter.py          | 11 ++---
 imod/mf6/rch.py                               |  4 +-
 imod/mf6/riv.py                               |  4 +-
 imod/mf6/wel.py                               |  6 +--
 .../flow_transport_simulation_fixture.py      |  2 +-
 12 files changed, 43 insertions(+), 50 deletions(-)

diff --git a/imod/mf6/auxiliary_variables.py b/imod/mf6/auxiliary_variables.py
index 6312741b6..f076d62d0 100644
--- a/imod/mf6/auxiliary_variables.py
+++ b/imod/mf6/auxiliary_variables.py
@@ -28,38 +28,34 @@ def _get_auxiliary_data_variable_names_mapping(package: IPackage):
     return result
 
 
-def add_periodic_auxiliary_variable(package: IPackage) -> None:
+def expand_transient_auxiliary_variables(package: IPackage) -> None:
     """
-    Splits an auxiliary dataarray (with one or more auxiliary variable dimension) into dataarrays per
-    auxiliary variable dimension. For example a concentration auxiliary variable in a flow package
-    will have a species dimension, and will be split in several dataarrays- one for each species.
+    Splits an auxiliary dataarray of the kind that could potentially be
+    time-dependent (with one or more auxiliary variable dimension) into
+    dataarrays per auxiliary variable dimension. For example a concentration
+    auxiliary variable in a flow package will have a species dimension, and will
+    be split in several dataarrays- one for each species.
     """
 
-    for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields().items():
-        aux_coords = package.dataset[aux_var_name].coords[aux_var_dimensions].values
-        for s in aux_coords:
-            package.dataset[s] = package.dataset[aux_var_name].sel(
-                {aux_var_dimensions: s}
-            )
+    if len(package.auxiliary_data_fields()) > 0:
+        for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields().items():
+            if aux_var_name in list(package.dataset.keys()):
+                aux_coords = (
+                    package.dataset[aux_var_name].coords[aux_var_dimensions].values
+                )
+                for s in aux_coords:
+                    package.dataset[s] = package.dataset[aux_var_name].sel(
+                        {aux_var_dimensions: s}
+                    )
 
 
-def remove_periodic_auxiliary_variable(package: IPackage) -> None:
+def remove_expanded_auxiliary_variables_from_dataset(package: IPackage) -> None:
     """
-    Removes the data arrays created by :meth:add_periodic_auxiliary_variable(...) but does not
-    remove the auxiliary dataarray used as source for :meth:add_periodic_auxiliary_variable(...)
+    Removes the data arrays created by :meth:expand_transient_auxiliary_variables(...) but does not
+    remove the auxiliary dataarray used as source for :meth:expand_transient_auxiliary_variables(...)
     """
     for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields().items():
         if aux_var_dimensions in package.dataset.coords:
             for species in package.dataset.coords[aux_var_dimensions].values:
                 if species in list(package.dataset.keys()):
                     package.dataset = package.dataset.drop_vars(species)
-
-
-def has_auxiliary_variable_source_array(package: IPackage) -> bool:
-    """
-    Returns True if a package contains an auxiliary data source data array
-    """
-    for aux_var_name, _ in package.auxiliary_data_fields().items():
-        if aux_var_name in package.dataset.keys():
-            return True
-    return False
diff --git a/imod/mf6/chd.py b/imod/mf6/chd.py
index d76279a0f..ad39a54fb 100644
--- a/imod/mf6/chd.py
+++ b/imod/mf6/chd.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -131,7 +131,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
         self.dataset["save_flows"] = save_flows
diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py
index 9f199dafe..6e904e104 100644
--- a/imod/mf6/drn.py
+++ b/imod/mf6/drn.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
@@ -121,7 +121,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
         self.dataset["save_flows"] = save_flows
diff --git a/imod/mf6/evt.py b/imod/mf6/evt.py
index 0644bac7a..dc92025e5 100644
--- a/imod/mf6/evt.py
+++ b/imod/mf6/evt.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
@@ -207,7 +207,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["fixed_cell"] = fixed_cell
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py
index d47d60732..7f6a0caff 100644
--- a/imod/mf6/ghb.py
+++ b/imod/mf6/ghb.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -139,7 +139,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
         self.dataset["save_flows"] = save_flows
diff --git a/imod/mf6/gwfgwf.py b/imod/mf6/gwfgwf.py
index 8f1b8ae5a..a1d6df261 100644
--- a/imod/mf6/gwfgwf.py
+++ b/imod/mf6/gwfgwf.py
@@ -4,7 +4,7 @@
 import numpy as np
 import xarray as xr
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+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
@@ -50,7 +50,7 @@ def __init__(
             self.dataset["auxiliary_data"] = xr.merge(auxiliary_variables).to_array(
                 name="auxiliary_data"
             )
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
 
     def set_options(
         self,
diff --git a/imod/mf6/mf6_wel_adapter.py b/imod/mf6/mf6_wel_adapter.py
index 683efda65..89dc5fd54 100644
--- a/imod/mf6/mf6_wel_adapter.py
+++ b/imod/mf6/mf6_wel_adapter.py
@@ -14,7 +14,7 @@
 
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.schemata import DTypeSchema
 
@@ -62,7 +62,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self._validate_init_schemata(validate)
 
     def _ds_to_arrdict(self, ds):
diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py
index 121599df9..1ff583f68 100644
--- a/imod/mf6/multimodel/modelsplitter.py
+++ b/imod/mf6/multimodel/modelsplitter.py
@@ -3,9 +3,8 @@
 import numpy as np
 
 from imod.mf6.auxiliary_variables import (
-    add_periodic_auxiliary_variable,
-    has_auxiliary_variable_source_array,
-    remove_periodic_auxiliary_variable,
+    expand_transient_auxiliary_variables,
+    remove_expanded_auxiliary_variables_from_dataset,
 )
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.hfb import HorizontalFlowBarrierBase
@@ -93,8 +92,7 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
 
     for pkg_name, package in model.items():
         if isinstance(package, BoundaryCondition):
-            if has_auxiliary_variable_source_array(package):
-                remove_periodic_auxiliary_variable(package)
+            remove_expanded_auxiliary_variables_from_dataset(package)
 
         sliced_package = clip_by_grid(package, partition_info.active_domain)
 
@@ -110,6 +108,5 @@ def slice_model(partition_info: PartitionInfo, model: Modflow6Model) -> Modflow6
             )
 
         if isinstance(package, BoundaryCondition):
-            if has_auxiliary_variable_source_array(package):
-                add_periodic_auxiliary_variable(package)
+            expand_transient_auxiliary_variables(package)
     return new_model
diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py
index ed5f62647..2fcca90c0 100644
--- a/imod/mf6/rch.py
+++ b/imod/mf6/rch.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -125,7 +125,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
         self.dataset["save_flows"] = save_flows
diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py
index 5a27d554a..50ca419e2 100644
--- a/imod/mf6/riv.py
+++ b/imod/mf6/riv.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -151,7 +151,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
         self.dataset["print_input"] = print_input
         self.dataset["print_flows"] = print_flows
         self.dataset["save_flows"] = save_flows
diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py
index c86f91f13..5d640b435 100644
--- a/imod/mf6/wel.py
+++ b/imod/mf6/wel.py
@@ -11,7 +11,7 @@
 import xugrid as xu
 
 import imod
-from imod.mf6.auxiliary_variables import add_periodic_auxiliary_variable
+from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import (
     BoundaryCondition,
     DisStructuredBoundaryCondition,
@@ -678,7 +678,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
 
         self._validate_init_schemata(validate)
 
@@ -835,7 +835,7 @@ def __init__(
         if concentration is not None:
             self.dataset["concentration"] = concentration
             self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            add_periodic_auxiliary_variable(self)
+            expand_transient_auxiliary_variables(self)
 
         self._validate_init_schemata(validate)
 
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index d4d58a41d..346ba321f 100644
--- a/imod/tests/fixtures/flow_transport_simulation_fixture.py
+++ b/imod/tests/fixtures/flow_transport_simulation_fixture.py
@@ -72,7 +72,7 @@ def create_transport_model(flow_model, species_name, dispersivity, retardation,
 
 # %%
 # Create the spatial discretization.
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")
 def flow_transport_simulation():
     nlay = 1
     nrow = 2

From f2731dcb5f458b0aef634896958d52b3c990ecdf Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 26 Jan 2024 13:02:51 +0100
Subject: [PATCH 13/47] started adding test

---
 .../test_mf6_partitioning_structured_transport.py     | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py
new file mode 100644
index 000000000..6bdc0a85a
--- /dev/null
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py
@@ -0,0 +1,11 @@
+from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
+from imod.typing.grid import zeros_like
+
+
+def test_partition_structured(flow_transport_simulation):
+    # Arrange.
+    transport_model = flow_transport_simulation["tpt_a"]
+    submodel_labels = zeros_like(transport_model.domain)
+    submodel_labels[:, :, 30:] = 1
+
+    partition_info = create_partition_info(submodel_labels)
\ No newline at end of file

From f5f682a553e13fa98ede4be60e1b9b7355a50927 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 26 Jan 2024 14:57:09 +0100
Subject: [PATCH 14/47] started adding test

---
 .../test_mf6_partitioning_structured_transport.py           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py
index 6bdc0a85a..fd360f9b5 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_structured_transport.py
@@ -5,7 +5,7 @@
 def test_partition_structured(flow_transport_simulation):
     # Arrange.
     transport_model = flow_transport_simulation["tpt_a"]
-    submodel_labels = zeros_like(transport_model.domain)
-    submodel_labels[:, :, 30:] = 1
+    submodel_labels = zeros_like(transport_model.domain.isel(layer=0))
+    submodel_labels[:, 30:] = 1
 
-    partition_info = create_partition_info(submodel_labels)
\ No newline at end of file
+    split_simulation = flow_transport_simulation.split(submodel_labels)
\ No newline at end of file

From 48fcb58608d6d0d283069a43b83d9003319f360f Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 26 Jan 2024 15:11:29 +0100
Subject: [PATCH 15/47] split method now doesn't crash on transport model

---
 imod/mf6/simulation.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 4407f36fe..99c20b51a 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -991,13 +991,14 @@ def _set_exchange_options(self):
         for exchange in self["split_exchanges"]:
             model_name_1 = exchange.dataset["model_name_1"].values[()]
             model_1 = self[model_name_1]
-            exchange.set_options(
-                save_flows=model_1["oc"].is_budget_output,
-                dewatered=model_1["npf"].is_dewatered,
-                variablecv=model_1["npf"].is_variable_vertical_conductance,
-                xt3d=model_1["npf"].get_xt3d_option(),
-                newton=model_1.is_use_newton(),
-            )
+            if isinstance (model_1, GroundwaterFlowModel):
+                exchange.set_options(
+                    save_flows=model_1["oc"].is_budget_output,
+                    dewatered=model_1["npf"].is_dewatered,
+                    variablecv=model_1["npf"].is_variable_vertical_conductance,
+                    xt3d=model_1["npf"].get_xt3d_option(),
+                    newton=model_1.is_use_newton(),
+                )
 
     def _filter_inactive_cells_from_exchanges(self) -> None:
         for ex in self["split_exchanges"]:

From 866f6a2efde7522cc9afbb5c4ddce01603c8f197 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 26 Jan 2024 17:44:17 +0100
Subject: [PATCH 16/47] adding package and template for gwtgwt

---
 imod/mf6/gwtgwt.py              | 39 +++++++++++++++++++++++++++++++++
 imod/templates/mf6/ex-gwtgwt.j2 | 36 ++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 imod/mf6/gwtgwt.py
 create mode 100644 imod/templates/mf6/ex-gwtgwt.j2

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
new file mode 100644
index 000000000..d620354aa
--- /dev/null
+++ b/imod/mf6/gwtgwt.py
@@ -0,0 +1,39 @@
+
+
+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."""
+
+    def __init__(
+        self,
+        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,
+    ):
+        super().__init__(locals())
+        self.dataset["cell_id1"] = cell_id1
+        self.dataset["cell_id2"] = cell_id2
+        self.dataset["layer"] = layer
+        self.dataset["model_name_1"] = model_id1
+        self.dataset["model_name_2"] = model_id2
+        self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
+        self.dataset["cl1"] = cl1
+        self.dataset["cl2"] = cl2
+        self.dataset["hwva"] = hwva
+
+        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)    
\ No newline at end of file
diff --git a/imod/templates/mf6/ex-gwtgwt.j2 b/imod/templates/mf6/ex-gwtgwt.j2
new file mode 100644
index 000000000..c78606ef6
--- /dev/null
+++ b/imod/templates/mf6/ex-gwtgwt.j2
@@ -0,0 +1,36 @@
+# this file contains the exchanges from a transport model associated to flow model{{flow_model_name_1}}
+# to a transport model associated to flow model{{flow_model_name_2}}
+
+begin options
+{{flow_model_name_1}}
+{{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
+{% endif -%}
+{%- if dsp_xt3d_off is defined -%} dsp_xt3d_off
+{% endif -%}
+{%- if dsp_xt3d_rhs is defined -%}  dsp_xt3d_rhs 
+{% 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

From b007adab15f3e1376d84929988e3a47fde65d9e7 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Mon, 29 Jan 2024 11:58:32 +0100
Subject: [PATCH 17/47] work in progress

---
 imod/mf6/gwtgwt.py                      | 25 ++++++++++++++++++++++---
 imod/mf6/multimodel/exchange_creator.py |  7 ++++++-
 imod/mf6/simulation.py                  | 16 +++++++++-------
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
index d620354aa..9226b8646 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -11,6 +11,7 @@ def __init__(
         self,
         flow_model_id1: str,
         flow_model_id2: str,
+        
         cell_id1: xr.DataArray,
         cell_id2: xr.DataArray,
         layer: xr.DataArray,
@@ -24,8 +25,10 @@ def __init__(
         self.dataset["cell_id1"] = cell_id1
         self.dataset["cell_id2"] = cell_id2
         self.dataset["layer"] = layer
-        self.dataset["model_name_1"] = model_id1
-        self.dataset["model_name_2"] = model_id2
+        self.dataset["model_name_1"] = flow_model_id1
+
+        self.dataset["model_name_2"] = flow_model_id1
+
         self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
         self.dataset["cl1"] = cl1
         self.dataset["cl2"] = cl2
@@ -36,4 +39,20 @@ def __init__(
             self.dataset["auxiliary_data"] = xr.merge(auxiliary_variables).to_array(
                 name="auxiliary_data"
             )
-            expand_transient_auxiliary_variables(self)    
\ No newline at end of file
+            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
diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 7c7084de6..bf32110ca 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -10,6 +10,7 @@
 from imod.mf6.utilities.grid import get_active_domain_slice, to_cell_idx
 from imod.typing import GridDataArray
 
+from imod.mf6.gwfgwt import GWFGWT
 
 def _adjust_gridblock_indexing(connected_cells: xr.Dataset) -> xr.Dataset:
     """
@@ -110,7 +111,7 @@ def __init__(
 
         self._geometric_information = self._compute_geometric_information()
 
-    def create_exchanges(self, model_name: str, layers: GridDataArray) -> List[GWFGWF]:
+    def create_gwfgwf_exchanges(self, model_name: str, layers: GridDataArray) -> List[GWFGWF]:
         """
         Create GroundWaterFlow-GroundWaterFlow exchanges based on the submodel_labels array provided in the class
         constructor. The layer parameter is used to extrude the cell connection through all the layers. An exchange
@@ -252,3 +253,7 @@ def rearrange_connected_cells(self):
         df.loc[label_decreasing, colnames_reversed] = decreasing_connections
 
         self._connected_cells = df
+
+    def create_gwfgwt_exchanges(self, model_name: str, layers: GridDataArray) -> List[GWFGWT]:
+        pass
+        
\ No newline at end of file
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 99c20b51a..c7a4b872a 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -923,14 +923,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 )
 
         exchanges = []
+        flowmodels = self.get_models_of_type("gwf6")
         for model_name, model in original_models.items():
-            exchanges += exchange_creator.create_exchanges(
-                model_name, model.domain.layer
-            )
-
-        new_simulation["solver"]["modelnames"] = xr.DataArray(
-            list(get_models(new_simulation).keys())
-        )
+            if isinstance(model, GroundwaterFlowModel):
+                exchanges += exchange_creator.create_gwfgwf_exchanges(
+                    model_name, model.domain.layer
+                )
+            if isinstance(model, GroundwaterTransportModel):
+                exchanges += exchange_creator.create_gwfgwt_exchanges(
+                    model_name,model, flowmodels, model.domain.layer
+                )
 
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()

From 9bfa03ed330f75bedcd630ee786142e5a873a6c1 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Mon, 29 Jan 2024 13:44:04 +0100
Subject: [PATCH 18/47] work in progress

---
 imod/mf6/multimodel/exchange_creator.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index bf32110ca..73b32c236 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -255,5 +255,10 @@ def rearrange_connected_cells(self):
         self._connected_cells = df
 
     def create_gwfgwt_exchanges(self, model_name: str, layers: GridDataArray) -> List[GWFGWT]:
-        pass
+        
+        layers = layers.to_dataframe().filter(["layer"])
+
+        connected_cells_with_geometric_info = pd.merge(
+            self._connected_cells, self._geometric_information
+        )
         
\ No newline at end of file

From eef6081f16379701d9636fab77f797d502777a00 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Mon, 29 Jan 2024 17:08:36 +0100
Subject: [PATCH 19/47] adding gwfgwt packages

---
 imod/mf6/simulation.py                        | 34 +++++++++++++------
 .../test_multimodel/test_mf6_modelsplitter.py | 20 ++++++++++-
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index c7a4b872a..4baa17b38 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -929,10 +929,10 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 exchanges += exchange_creator.create_gwfgwf_exchanges(
                     model_name, model.domain.layer
                 )
-            if isinstance(model, GroundwaterTransportModel):
-                exchanges += exchange_creator.create_gwfgwt_exchanges(
-                    model_name,model, flowmodels, model.domain.layer
-                )
+
+        new_simulation["solver"]["modelnames"] = xr.DataArray(
+            list(get_models(new_simulation).keys())
+        )                
 
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()
@@ -991,8 +991,9 @@ def _add_modelsplit_exchanges(self, exchanges_list: List[GWFGWF]) -> None:
     def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
-            model_name_1 = exchange.dataset["model_name_1"].values[()]
-            model_1 = self[model_name_1]
+            if isinstance(exchange, GWFGWF): 
+                model_name_1 = exchange.dataset["model_name_1"].values[()]
+                model_1 = self[model_name_1]
             if isinstance (model_1, GroundwaterFlowModel):
                 exchange.set_options(
                     save_flows=model_1["oc"].is_budget_output,
@@ -1052,12 +1053,23 @@ def __repr__(self) -> str:
     def _generate_gwfgwt_exchanges(self):
         flow_models = self.get_models_of_type("gwf6")
         transport_models = self.get_models_of_type("gwt6")
-
         # exchange for flow and transport
         exchanges = []
-        if len(flow_models) == 1 and len(transport_models) > 0:
-            flow_model_name = list(flow_models.keys())[0]
-            for transport_model_name in transport_models.keys():
-                exchanges.append(GWFGWT(flow_model_name, transport_model_name))
+
+        for flow_model_name in flow_models:
+            tpt_models_of_flow_model = []
+            flow_model = self[flow_model_name]
+            domain = flow_model.domain
+            for tpt_model_name in transport_models:
+                tpt_model = self[tpt_model_name]
+                if tpt_model.domain.equals(domain):
+                    tpt_models_of_flow_model.append(tpt_model_name)
+
+            if len(tpt_models_of_flow_model) > 0:
+                for transport_model_name in tpt_models_of_flow_model:
+                    exchanges.append(GWFGWT(flow_model_name, transport_model_name))
+                        
+
+
 
         return exchanges
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index 987c02b30..ec0ef88ea 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -101,4 +101,22 @@ def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulatio
     assert "species_d" in list(split_simulation["flow_1"]["chd"].dataset.keys())
     assert "species_d" in list(split_simulation["flow_1"]["rch"].dataset.keys())
     assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())
-    pass
+
+
+
+@pytest.mark.usefixtures("flow_transport_simulation")
+def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
+
+    simulation = flow_transport_simulation
+
+    flow_model = simulation["flow"]
+    active = flow_model.domain
+
+    submodel_labels = zeros_like(active)
+    submodel_labels = submodel_labels.drop_vars("layer")
+    submodel_labels.values[:, :, 50:] = 1
+    submodel_labels = submodel_labels.sel(layer=0, drop=True)
+
+    new_simulation = simulation.split(submodel_labels)
+    new_simulation.write(tmp_path, binary=False)
+    pass
\ No newline at end of file

From b64b59b64aa4a0407832ee8d9d7c56a0211e5652 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 10:25:11 +0100
Subject: [PATCH 20/47] fixed tests

---
 imod/tests/test_mf6/test_mf6_simulation.py    | 19 ++++++++++---------
 .../test_exchange_creator_structured.py       |  6 +++---
 .../test_exchange_creator_unstructured.py     |  8 ++++----
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 24f9c96d8..51a3bd202 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -14,6 +14,7 @@
 
 import imod
 from imod.mf6.model import Modflow6Model
+from imod.mf6.model_gwf import GroundwaterFlowModel
 from imod.mf6.multimodel.modelsplitter import PartitionInfo
 from imod.mf6.simulation import get_models, get_packages
 from imod.mf6.statusinfo import NestedStatusInfo, StatusInfo
@@ -298,10 +299,10 @@ def test_split_multiple_models(
 
         simulation = setup_simulation
 
-        model_mock1 = MagicMock(spec_set=Modflow6Model)
+        model_mock1 = MagicMock(spec_set=GroundwaterFlowModel)
         model_mock1._model_id = "test_model_id1"
 
-        model_mock2 = MagicMock(spec_set=Modflow6Model)
+        model_mock2 = MagicMock(spec_set=GroundwaterFlowModel)
         model_mock2._model_id = "test_model_id2"
 
         simulation["test_model1"] = model_mock1
@@ -309,7 +310,7 @@ def test_split_multiple_models(
 
         simulation["solver"]["modelnames"] = ["test_model1", "test_model2"]
 
-        slice_model_mock.return_value = MagicMock(spec_set=Modflow6Model)
+        slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
 
         active = idomain.sel(layer=1)
         submodel_labels = xu.zeros_like(active).where(active.grid.face_y > 0.0, 1)
@@ -360,11 +361,11 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
 
         simulation = setup_simulation
 
-        model_mock1 = MagicMock(spec_set=Modflow6Model)
+        model_mock1 = MagicMock(spec_set=GroundwaterFlowModel)
         model_mock1._model_id = "test_model_id1"
         model_mock1.domain = idomain
 
-        model_mock2 = MagicMock(spec_set=Modflow6Model)
+        model_mock2 = MagicMock(spec_set=GroundwaterFlowModel)
         model_mock2._model_id = "test_model_id2"
         model_mock2.domain = idomain
 
@@ -373,7 +374,7 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
 
         simulation["solver"]["modelnames"] = ["test_model1", "test_model2"]
 
-        slice_model_mock.return_value = MagicMock(spec_set=Modflow6Model)
+        slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
 
         active = idomain.sel(layer=1)
         submodel_labels = xr.zeros_like(active).where(active.y > 50, 1)
@@ -390,9 +391,9 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
             submodel_labels, create_partition_info_mock()
         )
 
-        assert exchange_creator_mock.return_value.create_exchanges.call_count == 2
-        call1 = exchange_creator_mock.return_value.create_exchanges.call_args_list[0][0]
-        call2 = exchange_creator_mock.return_value.create_exchanges.call_args_list[1][0]
+        assert exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2
+        call1 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[0][0]
+        call2 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[1][0]
 
         assert call1[0] == "test_model1"
         xr.testing.assert_equal(call1[1], idomain.layer)
diff --git a/imod/tests/test_mf6/test_multimodel/test_exchange_creator_structured.py b/imod/tests/test_mf6/test_multimodel/test_exchange_creator_structured.py
index 119222fce..2c35d6b70 100644
--- a/imod/tests/test_mf6/test_multimodel/test_exchange_creator_structured.py
+++ b/imod/tests/test_mf6/test_multimodel/test_exchange_creator_structured.py
@@ -68,7 +68,7 @@ def test_create_exchanges_validate_number_of_exchanges(
         layer = idomain.layer
 
         # Act.
-        exchanges = exchange_creator.create_exchanges(model_name, layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges(model_name, layer)
 
         # Assert.
         num_exchanges_x_direction = y_number_partitions * (x_number_partitions - 1)
@@ -296,7 +296,7 @@ def test_create_exchanges_validate_local_cell_ids(
         layer = idomain.layer
 
         # Act.
-        exchanges = exchange_creator.create_exchanges(model_name, layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges(model_name, layer)
 
         # Assert.
         assert len(exchanges) == len(expected_exchanges)
@@ -337,7 +337,7 @@ def test_exchange_geometric_information(
         layer = idomain.layer
 
         # Act.
-        exchanges = exchange_creator.create_exchanges(model_name, layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges(model_name, layer)
 
         # Assert.
         assert len(exchanges) == len(expected_exchanges)
diff --git a/imod/tests/test_mf6/test_multimodel/test_exchange_creator_unstructured.py b/imod/tests/test_mf6/test_multimodel/test_exchange_creator_unstructured.py
index d63b71f04..f928e1751 100644
--- a/imod/tests/test_mf6/test_multimodel/test_exchange_creator_unstructured.py
+++ b/imod/tests/test_mf6/test_multimodel/test_exchange_creator_unstructured.py
@@ -48,7 +48,7 @@ def test_create_exchanges_unstructured_validate_number_of_exchanges(
         exchange_creator = ExchangeCreator_Unstructured(submodel_labels, partition_info)
 
         # Act.
-        exchanges = exchange_creator.create_exchanges("flow", idomain.layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges("flow", idomain.layer)
 
         # Assert.
         assert len(exchanges) == number_partitions - 1
@@ -75,7 +75,7 @@ def test_create_exchanges_unstructured_validate_exchange_locations(
         exchange_creator = ExchangeCreator_Unstructured(submodel_labels, partition_info)
 
         # Act.
-        exchanges = exchange_creator.create_exchanges("flow", idomain.layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges("flow", idomain.layer)
 
         # Assert.
         nlayer = 3
@@ -118,7 +118,7 @@ def test_create_exchanges_unstructured_validate_geometric_coefficients(
         exchange_creator = ExchangeCreator_Unstructured(submodel_labels, partition_info)
 
         # Act.
-        exchanges = exchange_creator.create_exchanges("flow", idomain.layer)
+        exchanges = exchange_creator.create_gwfgwf_exchanges("flow", idomain.layer)
 
         # Assert.
         assert np.allclose(exchanges[0].dataset["cl1"], expected_cl1)
@@ -159,7 +159,7 @@ def test_create_exchanges_unstructured_validate_auxiliary_coefficients(
         exchange_creator = ExchangeCreator_Unstructured(submodel_labels, partition_info)
 
         # Act.
-        _ = exchange_creator.create_exchanges("flow", idomain.layer)
+        _ = exchange_creator.create_gwfgwf_exchanges("flow", idomain.layer)
         """
         exchanges = exchange_creator.create_exchanges("flow", idomain.layer)
         """

From a4743e0a27017d4880571ecf5fb41a3cc19fd076 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 13:04:55 +0100
Subject: [PATCH 21/47] now can write the model

---
 imod/mf6/ims.py        | 18 ++++++++++++++++++
 imod/mf6/simulation.py | 18 +++++++++++++-----
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/imod/mf6/ims.py b/imod/mf6/ims.py
index d5ed53b37..f44ab4c68 100644
--- a/imod/mf6/ims.py
+++ b/imod/mf6/ims.py
@@ -427,6 +427,24 @@ def __init__(
         self.dataset["no_ptc"] = no_ptc
         self._validate_init_schemata(validate)
 
+    def remove_model_from_solution(self, modelname: str):
+        models_in_solution = self.dataset["modelnames"].values
+        filtered_models = []
+        for m in models_in_solution:
+            if m != modelname:
+                filtered_models.append(m)
+        if len(filtered_models) == 0:
+            self.dataset = self.dataset.drop_vars("modelnames")
+        else:
+            self.dataset.update({ "modelnames": ( "model" , filtered_models) })
+
+    def add_model_to_solution(self, modelname: str):
+        models_in_solution = []
+        if  "modelnames" in self.dataset.keys():
+            models_in_solution =  list(self.dataset["modelnames"].values)
+        models_in_solution.append(modelname)
+        # self.dataset["modelnames"] = xr.DataArray(dims = {"model"}, data = models_in_solution)
+        self.dataset.update({ "modelnames": ( "model" , models_in_solution) })
 
 def SolutionPresetSimple(
     modelnames, print_option="summary", csv_output=False, no_ptc=False
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 4baa17b38..3df5b8fec 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -37,6 +37,7 @@
 from imod.schemata import ValidationError
 from imod.typing import GridDataArray, GridDataset
 from imod.typing.grid import concat, is_unstructured, merge, merge_partitions, nan_like
+from imod.mf6.ims import Solution
 
 OUTPUT_FUNC_MAPPING = {
     "head": (open_hds, GroundwaterFlowModel),
@@ -916,11 +917,14 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             new_simulation[package_name] = package
 
         for model_name, model in original_models.items():
+            solution_name, solution_group = self.get_solution( model_name)
+            new_simulation[solution_name].remove_model_from_solution(model_name)
             for submodel_partition_info in partition_info:
                 new_model_name = f"{model_name}_{submodel_partition_info.id}"
                 new_simulation[new_model_name] = slice_model(
                     submodel_partition_info, model
                 )
+                new_simulation[solution_name].add_model_to_solution(new_model_name)
 
         exchanges = []
         flowmodels = self.get_models_of_type("gwf6")
@@ -928,11 +932,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             if isinstance(model, GroundwaterFlowModel):
                 exchanges += exchange_creator.create_gwfgwf_exchanges(
                     model_name, model.domain.layer
-                )
-
-        new_simulation["solver"]["modelnames"] = xr.DataArray(
-            list(get_models(new_simulation).keys())
-        )                
+                )        
 
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()
@@ -1031,6 +1031,14 @@ def _filter_inactive_cells_exchange_domain(self, ex: GWFGWF, i: int) -> None:
         active_exchange_domain = active_exchange_domain.dropna("index")
         ex.dataset = ex.dataset.sel(index=active_exchange_domain["index"])
 
+    def get_solution(self, model_name): 
+        for k, v in self.items():
+            if isinstance(v, Solution):
+                if model_name in v["modelnames"]:
+                    return k, v
+        return None
+
+
     def __repr__(self) -> str:
         typename = type(self).__name__
         INDENT = "    "

From 32455cc95669d66f6ce585088437441c274d3060 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 13:34:34 +0100
Subject: [PATCH 22/47] adding tests

---
 imod/mf6/ims.py                               | 10 +--
 imod/tests/test_mf6/test_mf6_ims.py           | 63 +++++++++++++++++++
 .../test_multimodel/test_mf6_modelsplitter.py |  1 +
 3 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/imod/mf6/ims.py b/imod/mf6/ims.py
index f44ab4c68..a8d29ec0c 100644
--- a/imod/mf6/ims.py
+++ b/imod/mf6/ims.py
@@ -429,10 +429,10 @@ def __init__(
 
     def remove_model_from_solution(self, modelname: str):
         models_in_solution = self.dataset["modelnames"].values
-        filtered_models = []
-        for m in models_in_solution:
-            if m != modelname:
-                filtered_models.append(m)
+        if modelname not in models_in_solution:
+            raise ValueError(f"attempted to remove model {modelname} from solution, but it was not found.")
+        filtered_models = [m for m in models_in_solution if m != modelname]
+
         if len(filtered_models) == 0:
             self.dataset = self.dataset.drop_vars("modelnames")
         else:
@@ -442,6 +442,8 @@ def add_model_to_solution(self, modelname: str):
         models_in_solution = []
         if  "modelnames" in self.dataset.keys():
             models_in_solution =  list(self.dataset["modelnames"].values)
+        if modelname in models_in_solution:
+            raise ValueError(f"attempted to add model {modelname} to solution, but it was already in it.")
         models_in_solution.append(modelname)
         # self.dataset["modelnames"] = xr.DataArray(dims = {"model"}, data = models_in_solution)
         self.dataset.update({ "modelnames": ( "model" , models_in_solution) })
diff --git a/imod/tests/test_mf6/test_mf6_ims.py b/imod/tests/test_mf6/test_mf6_ims.py
index 60d197419..e75827bce 100644
--- a/imod/tests/test_mf6/test_mf6_ims.py
+++ b/imod/tests/test_mf6/test_mf6_ims.py
@@ -65,3 +65,66 @@ def test_wrong_dtype():
             reordering_method=None,
             relaxation_factor=0.97,
         )
+
+
+def test_drop_and_add_model():
+    ims = imod.mf6.Solution(
+        modelnames=["GWF_1"],
+        print_option="summary",
+        csv_output=False,
+        no_ptc=True,
+        outer_dvclose=1.0e-4,
+        outer_maximum=500,
+        under_relaxation=None,
+        inner_dvclose=1.0e-4,
+        inner_rclose=0.001,
+        inner_maximum=100,
+        linear_acceleration="cg",
+        scaling_method=None,
+        reordering_method=None,
+        relaxation_factor=0.97,
+    )
+    ims.remove_model_from_solution("GWF_1")
+    assert  "modelnames" not in ims.dataset.keys() 
+    ims.add_model_to_solution("GWF_2")   
+    assert "GWF_2"  in ims.dataset["modelnames"].values
+
+def test_remove_non_present_model():
+    ims = imod.mf6.Solution(
+        modelnames=["GWF_1"],
+        print_option="summary",
+        csv_output=False,
+        no_ptc=True,
+        outer_dvclose=1.0e-4,
+        outer_maximum=500,
+        under_relaxation=None,
+        inner_dvclose=1.0e-4,
+        inner_rclose=0.001,
+        inner_maximum=100,
+        linear_acceleration="cg",
+        scaling_method=None,
+        reordering_method=None,
+        relaxation_factor=0.97,
+    )
+    with pytest.raises(ValueError):
+        ims.remove_model_from_solution("non_existing_model")
+
+def test_add_already_present_model():
+    ims = imod.mf6.Solution(
+        modelnames=["preexisting_model"],
+        print_option="summary",
+        csv_output=False,
+        no_ptc=True,
+        outer_dvclose=1.0e-4,
+        outer_maximum=500,
+        under_relaxation=None,
+        inner_dvclose=1.0e-4,
+        inner_rclose=0.001,
+        inner_maximum=100,
+        linear_acceleration="cg",
+        scaling_method=None,
+        reordering_method=None,
+        relaxation_factor=0.97,
+    )
+    with pytest.raises(ValueError):
+        ims.add_model_to_solution("preexisting_model")
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index ec0ef88ea..d31494ca7 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -119,4 +119,5 @@ def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
 
     new_simulation = simulation.split(submodel_labels)
     new_simulation.write(tmp_path, binary=False)
+    assert_simulation_can_run(new_simulation, "dis", tmp_path)
     pass
\ No newline at end of file

From bfce7a92ee4d57aa1c86aabb540b3dd91d1abb61 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 14:47:47 +0100
Subject: [PATCH 23/47] formatting

---
 imod/mf6/ims.py                               | 22 ++---
 imod/mf6/multimodel/exchange_creator.py       |  6 +-
 imod/mf6/simulation.py                        | 16 ++--
 imod/tests/test_mf6/test_mf6_ims.py           | 80 +++++++------------
 imod/tests/test_mf6/test_mf6_simulation.py    | 16 +++-
 .../test_multimodel/test_mf6_modelsplitter.py |  4 +-
 6 files changed, 65 insertions(+), 79 deletions(-)

diff --git a/imod/mf6/ims.py b/imod/mf6/ims.py
index a8d29ec0c..02f43c2c1 100644
--- a/imod/mf6/ims.py
+++ b/imod/mf6/ims.py
@@ -427,26 +427,30 @@ def __init__(
         self.dataset["no_ptc"] = no_ptc
         self._validate_init_schemata(validate)
 
-    def remove_model_from_solution(self, modelname: str):
+    def remove_model_from_solution(self, modelname: str) -> None:
         models_in_solution = self.dataset["modelnames"].values
         if modelname not in models_in_solution:
-            raise ValueError(f"attempted to remove model {modelname} from solution, but it was not found.")
+            raise ValueError(
+                f"attempted to remove model {modelname} from solution, but it was not found."
+            )
         filtered_models = [m for m in models_in_solution if m != modelname]
 
         if len(filtered_models) == 0:
             self.dataset = self.dataset.drop_vars("modelnames")
         else:
-            self.dataset.update({ "modelnames": ( "model" , filtered_models) })
+            self.dataset.update({"modelnames": ("model", filtered_models)})
 
-    def add_model_to_solution(self, modelname: str):
+    def add_model_to_solution(self, modelname: str) -> None:
         models_in_solution = []
-        if  "modelnames" in self.dataset.keys():
-            models_in_solution =  list(self.dataset["modelnames"].values)
+        if "modelnames" in self.dataset.keys():
+            models_in_solution = list(self.dataset["modelnames"].values)
         if modelname in models_in_solution:
-            raise ValueError(f"attempted to add model {modelname} to solution, but it was already in it.")
+            raise ValueError(
+                f"attempted to add model {modelname} to solution, but it was already in it."
+            )
         models_in_solution.append(modelname)
-        # self.dataset["modelnames"] = xr.DataArray(dims = {"model"}, data = models_in_solution)
-        self.dataset.update({ "modelnames": ( "model" , models_in_solution) })
+        self.dataset.update({"modelnames": ("model", models_in_solution)})
+
 
 def SolutionPresetSimple(
     modelnames, print_option="summary", csv_output=False, no_ptc=False
diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 2fd2c5d36..5f290bba5 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -6,11 +6,11 @@
 import xarray as xr
 
 from imod.mf6.gwfgwf import GWFGWF
+from imod.mf6.gwfgwt import GWFGWT
 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
 
-from imod.mf6.gwfgwt import GWFGWT
 
 def _adjust_gridblock_indexing(connected_cells: xr.Dataset) -> xr.Dataset:
     """
@@ -111,7 +111,9 @@ def __init__(
 
         self._geometric_information = self._compute_geometric_information()
 
-    def create_gwfgwf_exchanges(self, model_name: str, layers: GridDataArray) -> List[GWFGWF]:
+    def create_gwfgwf_exchanges(
+        self, model_name: str, layers: GridDataArray
+    ) -> List[GWFGWF]:
         """
         Create GroundWaterFlow-GroundWaterFlow exchanges based on the submodel_labels array provided in the class
         constructor. The layer parameter is used to extrude the cell connection through all the layers. An exchange
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 3df5b8fec..4caf83078 100644
--- 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.ims import Solution
 from imod.mf6.model import Modflow6Model
 from imod.mf6.model_gwf import GroundwaterFlowModel
 from imod.mf6.model_gwt import GroundwaterTransportModel
@@ -37,7 +38,6 @@
 from imod.schemata import ValidationError
 from imod.typing import GridDataArray, GridDataset
 from imod.typing.grid import concat, is_unstructured, merge, merge_partitions, nan_like
-from imod.mf6.ims import Solution
 
 OUTPUT_FUNC_MAPPING = {
     "head": (open_hds, GroundwaterFlowModel),
@@ -917,7 +917,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             new_simulation[package_name] = package
 
         for model_name, model in original_models.items():
-            solution_name, solution_group = self.get_solution( model_name)
+            solution_name, solution_group = self.get_solution(model_name)
             new_simulation[solution_name].remove_model_from_solution(model_name)
             for submodel_partition_info in partition_info:
                 new_model_name = f"{model_name}_{submodel_partition_info.id}"
@@ -932,7 +932,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             if isinstance(model, GroundwaterFlowModel):
                 exchanges += exchange_creator.create_gwfgwf_exchanges(
                     model_name, model.domain.layer
-                )        
+                )
 
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()
@@ -991,10 +991,10 @@ def _add_modelsplit_exchanges(self, exchanges_list: List[GWFGWF]) -> None:
     def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
-            if isinstance(exchange, GWFGWF): 
+            if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
                 model_1 = self[model_name_1]
-            if isinstance (model_1, GroundwaterFlowModel):
+            if isinstance(model_1, GroundwaterFlowModel):
                 exchange.set_options(
                     save_flows=model_1["oc"].is_budget_output,
                     dewatered=model_1["npf"].is_dewatered,
@@ -1031,14 +1031,13 @@ def _filter_inactive_cells_exchange_domain(self, ex: GWFGWF, i: int) -> None:
         active_exchange_domain = active_exchange_domain.dropna("index")
         ex.dataset = ex.dataset.sel(index=active_exchange_domain["index"])
 
-    def get_solution(self, model_name): 
+    def get_solution(self, model_name):
         for k, v in self.items():
             if isinstance(v, Solution):
                 if model_name in v["modelnames"]:
                     return k, v
         return None
 
-
     def __repr__(self) -> str:
         typename = type(self).__name__
         INDENT = "    "
@@ -1076,8 +1075,5 @@ def _generate_gwfgwt_exchanges(self):
             if len(tpt_models_of_flow_model) > 0:
                 for transport_model_name in tpt_models_of_flow_model:
                     exchanges.append(GWFGWT(flow_model_name, transport_model_name))
-                        
-
-
 
         return exchanges
diff --git a/imod/tests/test_mf6/test_mf6_ims.py b/imod/tests/test_mf6/test_mf6_ims.py
index e75827bce..d8d7523a2 100644
--- a/imod/tests/test_mf6/test_mf6_ims.py
+++ b/imod/tests/test_mf6/test_mf6_ims.py
@@ -6,6 +6,25 @@
 from imod.schemata import ValidationError
 
 
+def create_ims() -> imod.mf6.Solution:
+    return imod.mf6.Solution(
+        modelnames=["GWF_1"],
+        print_option="summary",
+        csv_output=False,
+        no_ptc=True,
+        outer_dvclose=1.0e-4,
+        outer_maximum=500,
+        under_relaxation=None,
+        inner_dvclose=1.0e-4,
+        inner_rclose=0.001,
+        inner_maximum=100,
+        linear_acceleration="cg",
+        scaling_method=None,
+        reordering_method=None,
+        relaxation_factor=0.97,
+    )
+
+
 def test_render():
     ims = imod.mf6.Solution(
         modelnames=["GWF_1"],
@@ -68,63 +87,20 @@ def test_wrong_dtype():
 
 
 def test_drop_and_add_model():
-    ims = imod.mf6.Solution(
-        modelnames=["GWF_1"],
-        print_option="summary",
-        csv_output=False,
-        no_ptc=True,
-        outer_dvclose=1.0e-4,
-        outer_maximum=500,
-        under_relaxation=None,
-        inner_dvclose=1.0e-4,
-        inner_rclose=0.001,
-        inner_maximum=100,
-        linear_acceleration="cg",
-        scaling_method=None,
-        reordering_method=None,
-        relaxation_factor=0.97,
-    )
+    ims = create_ims()
     ims.remove_model_from_solution("GWF_1")
-    assert  "modelnames" not in ims.dataset.keys() 
-    ims.add_model_to_solution("GWF_2")   
-    assert "GWF_2"  in ims.dataset["modelnames"].values
+    assert "modelnames" not in ims.dataset.keys()
+    ims.add_model_to_solution("GWF_2")
+    assert "GWF_2" in ims.dataset["modelnames"].values
+
 
 def test_remove_non_present_model():
-    ims = imod.mf6.Solution(
-        modelnames=["GWF_1"],
-        print_option="summary",
-        csv_output=False,
-        no_ptc=True,
-        outer_dvclose=1.0e-4,
-        outer_maximum=500,
-        under_relaxation=None,
-        inner_dvclose=1.0e-4,
-        inner_rclose=0.001,
-        inner_maximum=100,
-        linear_acceleration="cg",
-        scaling_method=None,
-        reordering_method=None,
-        relaxation_factor=0.97,
-    )
+    ims = create_ims()
     with pytest.raises(ValueError):
         ims.remove_model_from_solution("non_existing_model")
 
+
 def test_add_already_present_model():
-    ims = imod.mf6.Solution(
-        modelnames=["preexisting_model"],
-        print_option="summary",
-        csv_output=False,
-        no_ptc=True,
-        outer_dvclose=1.0e-4,
-        outer_maximum=500,
-        under_relaxation=None,
-        inner_dvclose=1.0e-4,
-        inner_rclose=0.001,
-        inner_maximum=100,
-        linear_acceleration="cg",
-        scaling_method=None,
-        reordering_method=None,
-        relaxation_factor=0.97,
-    )
+    ims = create_ims()
     with pytest.raises(ValueError):
-        ims.add_model_to_solution("preexisting_model")
+        ims.add_model_to_solution("GWF_1")
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 51a3bd202..79ae288e3 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -391,9 +391,19 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
             submodel_labels, create_partition_info_mock()
         )
 
-        assert exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2
-        call1 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[0][0]
-        call2 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[1][0]
+        assert (
+            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2
+        )
+        call1 = (
+            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[
+                0
+            ][0]
+        )
+        call2 = (
+            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[
+                1
+            ][0]
+        )
 
         assert call1[0] == "test_model1"
         xr.testing.assert_equal(call1[1], idomain.layer)
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index d31494ca7..a759fd872 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -103,10 +103,8 @@ def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulatio
     assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())
 
 
-
 @pytest.mark.usefixtures("flow_transport_simulation")
 def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
-
     simulation = flow_transport_simulation
 
     flow_model = simulation["flow"]
@@ -120,4 +118,4 @@ def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
     new_simulation = simulation.split(submodel_labels)
     new_simulation.write(tmp_path, binary=False)
     assert_simulation_can_run(new_simulation, "dis", tmp_path)
-    pass
\ No newline at end of file
+    pass

From 7b3ebef16854063f829f1bab4dc384a80854639a Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 15:22:15 +0100
Subject: [PATCH 24/47] cleanup

---
 imod/mf6/multimodel/exchange_creator.py       |  1 -
 imod/mf6/simulation.py                        |  1 -
 .../test_multimodel/test_mf6_modelsplitter.py | 18 --------
 .../test_mf6_modelsplitter_transport.py       | 46 +++++++++++++++++++
 4 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 5f290bba5..33598cbeb 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -6,7 +6,6 @@
 import xarray as xr
 
 from imod.mf6.gwfgwf import GWFGWF
-from imod.mf6.gwfgwt import GWFGWT
 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
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 4caf83078..d9ef9053b 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -927,7 +927,6 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 new_simulation[solution_name].add_model_to_solution(new_model_name)
 
         exchanges = []
-        flowmodels = self.get_models_of_type("gwf6")
         for model_name, model in original_models.items():
             if isinstance(model, GroundwaterFlowModel):
                 exchanges += exchange_creator.create_gwfgwf_exchanges(
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
index a759fd872..61543ed37 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py
@@ -101,21 +101,3 @@ def test_slice_model_with_auxiliary_variables(tmp_path, flow_transport_simulatio
     assert "species_d" in list(split_simulation["flow_1"]["chd"].dataset.keys())
     assert "species_d" in list(split_simulation["flow_1"]["rch"].dataset.keys())
     assert "concentration" in list(split_simulation["flow_1"]["well"].dataset.keys())
-
-
-@pytest.mark.usefixtures("flow_transport_simulation")
-def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
-    simulation = flow_transport_simulation
-
-    flow_model = simulation["flow"]
-    active = flow_model.domain
-
-    submodel_labels = zeros_like(active)
-    submodel_labels = submodel_labels.drop_vars("layer")
-    submodel_labels.values[:, :, 50:] = 1
-    submodel_labels = submodel_labels.sel(layer=0, drop=True)
-
-    new_simulation = simulation.split(submodel_labels)
-    new_simulation.write(tmp_path, binary=False)
-    assert_simulation_can_run(new_simulation, "dis", tmp_path)
-    pass
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
index 599dbe6d8..c91efcf5b 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
@@ -1,4 +1,7 @@
+import pytest
+
 from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
+from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 from imod.typing.grid import zeros_like
 
 
@@ -21,3 +24,46 @@ def test_slice_model_structured(flow_transport_simulation):
     for submodel in submodel_list:
         for package_name in list(transport_model.keys()):
             assert package_name in list(submodel.keys())
+
+
+@pytest.mark.usefixtures("flow_transport_simulation")
+def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
+    simulation = flow_transport_simulation
+
+    flow_model = simulation["flow"]
+    active = flow_model.domain
+
+    submodel_labels = zeros_like(active)
+    submodel_labels = submodel_labels.drop_vars("layer")
+    submodel_labels.values[:, :, 50:] = 1
+    submodel_labels = submodel_labels.sel(layer=0, drop=True)
+
+    new_simulation = simulation.split(submodel_labels)
+    new_simulation.write(tmp_path, binary=False)
+    assert len(new_simulation["gwtgwf_exchanges"]) == 8
+
+    assert new_simulation["gwtgwf_exchanges"][0]["model_name_1"].values[()] == "flow_0"
+    assert new_simulation["gwtgwf_exchanges"][0]["model_name_2"].values[()] == "tpt_a_0"
+
+    assert new_simulation["gwtgwf_exchanges"][1]["model_name_1"].values[()] == "flow_0"
+    assert new_simulation["gwtgwf_exchanges"][1]["model_name_2"].values[()] == "tpt_b_0"
+
+    assert new_simulation["gwtgwf_exchanges"][2]["model_name_1"].values[()] == "flow_0"
+    assert new_simulation["gwtgwf_exchanges"][2]["model_name_2"].values[()] == "tpt_c_0"
+
+    assert new_simulation["gwtgwf_exchanges"][3]["model_name_1"].values[()] == "flow_0"
+    assert new_simulation["gwtgwf_exchanges"][3]["model_name_2"].values[()] == "tpt_d_0"
+
+    assert new_simulation["gwtgwf_exchanges"][4]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][4]["model_name_2"].values[()] == "tpt_a_1"
+
+    assert new_simulation["gwtgwf_exchanges"][5]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][5]["model_name_2"].values[()] == "tpt_b_1"
+
+    assert new_simulation["gwtgwf_exchanges"][6]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][6]["model_name_2"].values[()] == "tpt_c_1"
+
+    assert new_simulation["gwtgwf_exchanges"][7]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][7]["model_name_2"].values[()] == "tpt_d_1"
+
+    assert_simulation_can_run(new_simulation, "dis", tmp_path)

From 3bb9a2c6d0c1b97c0d12150c843ab35a5a88795b Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 15:46:11 +0100
Subject: [PATCH 25/47] type hints

---
 imod/mf6/simulation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index d9ef9053b..6970ee2e8 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1056,7 +1056,7 @@ def __repr__(self) -> str:
             content = attrs + ["){}"]
         return "\n".join(content)
 
-    def _generate_gwfgwt_exchanges(self):
+    def _generate_gwfgwt_exchanges(self)->list[GWFGWT]:
         flow_models = self.get_models_of_type("gwf6")
         transport_models = self.get_models_of_type("gwt6")
         # exchange for flow and transport

From ff13c9ffebe780c8e91c9698124b676d65fa9b8b Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 16:39:25 +0100
Subject: [PATCH 26/47] cleanup

---
 imod/mf6/simulation.py                     | 10 +++++-----
 imod/tests/test_mf6/test_mf6_simulation.py |  8 ++++++--
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 6970ee2e8..09433d189 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -917,7 +917,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             new_simulation[package_name] = package
 
         for model_name, model in original_models.items():
-            solution_name, solution_group = self.get_solution(model_name)
+            solution_name = self.get_solution(model_name)
             new_simulation[solution_name].remove_model_from_solution(model_name)
             for submodel_partition_info in partition_info:
                 new_model_name = f"{model_name}_{submodel_partition_info.id}"
@@ -1030,11 +1030,11 @@ def _filter_inactive_cells_exchange_domain(self, ex: GWFGWF, i: int) -> None:
         active_exchange_domain = active_exchange_domain.dropna("index")
         ex.dataset = ex.dataset.sel(index=active_exchange_domain["index"])
 
-    def get_solution(self, model_name):
+    def get_solution(self, model_name: str) -> str:
         for k, v in self.items():
             if isinstance(v, Solution):
-                if model_name in v["modelnames"]:
-                    return k, v
+                if model_name in v.dataset["modelnames"]:
+                    return k
         return None
 
     def __repr__(self) -> str:
@@ -1056,7 +1056,7 @@ def __repr__(self) -> str:
             content = attrs + ["){}"]
         return "\n".join(content)
 
-    def _generate_gwfgwt_exchanges(self)->list[GWFGWT]:
+    def _generate_gwfgwt_exchanges(self) -> list[GWFGWT]:
         flow_models = self.get_models_of_type("gwf6")
         transport_models = self.get_models_of_type("gwt6")
         # exchange for flow and transport
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 79ae288e3..9f171c9d2 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -308,7 +308,9 @@ def test_split_multiple_models(
         simulation["test_model1"] = model_mock1
         simulation["test_model2"] = model_mock2
 
-        simulation["solver"]["modelnames"] = ["test_model1", "test_model2"]
+        simulation["solver"].dataset = xr.Dataset(
+            {"modelnames": ["test_model1", "test_model2"]}
+        )
 
         slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
 
@@ -372,7 +374,9 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
         simulation["test_model1"] = model_mock1
         simulation["test_model2"] = model_mock2
 
-        simulation["solver"]["modelnames"] = ["test_model1", "test_model2"]
+        simulation["solver"].dataset = xr.Dataset(
+            {"modelnames": ["test_model1", "test_model2"]}
+        )
 
         slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
 

From bbdfb23035c50c16f98c0b03447a68bd8de4a509 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Tue, 30 Jan 2024 16:49:01 +0100
Subject: [PATCH 27/47] cleanup

---
 imod/mf6/ims.py                     | 4 +++-
 imod/mf6/simulation.py              | 4 ++--
 imod/tests/test_mf6/test_mf6_ims.py | 3 ++-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/imod/mf6/ims.py b/imod/mf6/ims.py
index 02f43c2c1..374530234 100644
--- a/imod/mf6/ims.py
+++ b/imod/mf6/ims.py
@@ -428,7 +428,9 @@ def __init__(
         self._validate_init_schemata(validate)
 
     def remove_model_from_solution(self, modelname: str) -> None:
-        models_in_solution = self.dataset["modelnames"].values
+        models_in_solution = []
+        if "modelnames" in self.dataset.keys():
+            models_in_solution = self.dataset["modelnames"].values
         if modelname not in models_in_solution:
             raise ValueError(
                 f"attempted to remove model {modelname} from solution, but it was not found."
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 09433d189..9b809ad10 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -917,7 +917,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             new_simulation[package_name] = package
 
         for model_name, model in original_models.items():
-            solution_name = self.get_solution(model_name)
+            solution_name = self.get_solution_name(model_name)
             new_simulation[solution_name].remove_model_from_solution(model_name)
             for submodel_partition_info in partition_info:
                 new_model_name = f"{model_name}_{submodel_partition_info.id}"
@@ -1030,7 +1030,7 @@ def _filter_inactive_cells_exchange_domain(self, ex: GWFGWF, i: int) -> None:
         active_exchange_domain = active_exchange_domain.dropna("index")
         ex.dataset = ex.dataset.sel(index=active_exchange_domain["index"])
 
-    def get_solution(self, model_name: str) -> str:
+    def get_solution_name(self, model_name: str) -> str:
         for k, v in self.items():
             if isinstance(v, Solution):
                 if model_name in v.dataset["modelnames"]:
diff --git a/imod/tests/test_mf6/test_mf6_ims.py b/imod/tests/test_mf6/test_mf6_ims.py
index d8d7523a2..9bdb40b38 100644
--- a/imod/tests/test_mf6/test_mf6_ims.py
+++ b/imod/tests/test_mf6/test_mf6_ims.py
@@ -96,8 +96,9 @@ def test_drop_and_add_model():
 
 def test_remove_non_present_model():
     ims = create_ims()
+    ims.remove_model_from_solution("GWF_1")
     with pytest.raises(ValueError):
-        ims.remove_model_from_solution("non_existing_model")
+        ims.remove_model_from_solution("GWF_1")
 
 
 def test_add_already_present_model():

From 6267c0394b3c068e4f83d9377b377ee630096a2f Mon Sep 17 00:00:00 2001
From: luitjansl <64598682+luitjansl@users.noreply.github.com>
Date: Wed, 31 Jan 2024 13:46:22 +0100
Subject: [PATCH 28/47] Update imod/tests/test_mf6/test_mf6_simulation.py

review comment

Co-authored-by: Joeri van Engelen <joerivanengelen@hotmail.com>
---
 imod/tests/test_mf6/test_mf6_simulation.py | 17 ++++-------------
 1 file changed, 4 insertions(+), 13 deletions(-)

diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 9f171c9d2..2d5ebbcf0 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -395,19 +395,10 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
             submodel_labels, create_partition_info_mock()
         )
 
-        assert (
-            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2
-        )
-        call1 = (
-            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[
-                0
-            ][0]
-        )
-        call2 = (
-            exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[
-                1
-            ][0]
-        )
+        assert exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2  # noqa: E501
+        call1 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[0][0]  # noqa: E501
+        call2 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[1][0]  # noqa: E501
+
 
         assert call1[0] == "test_model1"
         xr.testing.assert_equal(call1[1], idomain.layer)

From 580ed42563971c6b92f5d7723547e8bf6fe00c58 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 31 Jan 2024 13:46:47 +0100
Subject: [PATCH 29/47] started work

---
 imod/mf6/gwtgwt.py                      | 85 +++++++++++++++++++++++++
 imod/mf6/multimodel/exchange_creator.py | 68 ++++++++++++++++++++
 imod/mf6/simulation.py                  |  4 ++
 3 files changed, 157 insertions(+)
 create mode 100644 imod/mf6/gwtgwt.py

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
new file mode 100644
index 000000000..c0756f669
--- /dev/null
+++ b/imod/mf6/gwtgwt.py
@@ -0,0 +1,85 @@
+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 = "gwfgwf"
+    _template = Package._initialize_template(_pkg_id)
+
+    def __init__(
+        self,
+        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,
+    ):
+        super().__init__(locals())
+        self.dataset["cell_id1"] = cell_id1
+        self.dataset["cell_id2"] = cell_id2
+        self.dataset["layer"] = layer
+        self.dataset["model_name_1"] = model_id1
+        self.dataset["model_name_2"] = model_id2
+        self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
+        self.dataset["cl1"] = cl1
+        self.dataset["cl2"] = cl2
+        self.dataset["hwva"] = hwva
+
+        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
index 33598cbeb..ecc688b1f 100644
--- 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
@@ -194,6 +195,73 @@ def create_gwfgwf_exchanges(
 
         return exchanges
 
+    def create_gwtgwt_exchanges(
+        self, 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"
+            ):
+                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"}
+                    )
+                )
+
+                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 = (
+                    connected_domain_pair.merge(mapping1)
+                    .merge(mapping2)
+                    .filter(
+                        [
+                            "cell_id1",
+                            "cell_id2",
+                            "cl1",
+                            "cl2",
+                            "hwva",
+                            "angldegx",
+                            "cdist",
+                        ]
+                    )
+                )
+
+                connected_cells = pd.merge(layers, connected_cells, how="cross")
+
+                connected_cells_dataset = self._to_xarray(connected_cells)
+
+                _adjust_gridblock_indexing(connected_cells_dataset)
+
+                exchanges.append(
+                    GWTGWT(
+                        f"{model_name}_{model_id1}",
+                        f"{model_name}_{model_id2}",
+                        **connected_cells_dataset,
+                    )
+                )
+
+        return exchanges
+
     def _create_global_cellidx_to_local_cellid_mapping(
         self, partition_info: List[PartitionInfo]
     ) -> Dict[int, pd.DataFrame]:
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 9b809ad10..03ef32da2 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -932,6 +932,10 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 exchanges += exchange_creator.create_gwfgwf_exchanges(
                     model_name, model.domain.layer
                 )
+            if isinstance(model, GroundwaterTransportModel):
+                exchanges += exchange_creator.create_gwtgwt_exchanges(
+                    model_name, model.domain.layer
+                )
 
         new_simulation._add_modelsplit_exchanges(exchanges)
         new_simulation._set_exchange_options()

From 7b20d0299d38fbf7f48e873efad1badae7e1d5ee Mon Sep 17 00:00:00 2001
From: luitjansl <64598682+luitjansl@users.noreply.github.com>
Date: Wed, 31 Jan 2024 13:50:47 +0100
Subject: [PATCH 30/47] Update imod/mf6/simulation.py

review comment

Co-authored-by: Joeri van Engelen <joerivanengelen@hotmail.com>
---
 imod/mf6/simulation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 9b809ad10..8bdd4f4c1 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1068,7 +1068,7 @@ def _generate_gwfgwt_exchanges(self) -> list[GWFGWT]:
             domain = flow_model.domain
             for tpt_model_name in transport_models:
                 tpt_model = self[tpt_model_name]
-                if tpt_model.domain.equals(domain):
+                if tpt_model.domain.equals(flow_model.domain):
                     tpt_models_of_flow_model.append(tpt_model_name)
 
             if len(tpt_models_of_flow_model) > 0:

From 74b1b6fb4aa3f73621521ee78d4b257653a04ff2 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 31 Jan 2024 13:52:37 +0100
Subject: [PATCH 31/47] review comment

---
 imod/mf6/simulation.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 8bdd4f4c1..1a65b9598 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1065,7 +1065,6 @@ def _generate_gwfgwt_exchanges(self) -> list[GWFGWT]:
         for flow_model_name in flow_models:
             tpt_models_of_flow_model = []
             flow_model = self[flow_model_name]
-            domain = flow_model.domain
             for tpt_model_name in transport_models:
                 tpt_model = self[tpt_model_name]
                 if tpt_model.domain.equals(flow_model.domain):

From 18f2998f9a227700ed9246a7016f2b66e1f3114e Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 31 Jan 2024 14:05:03 +0100
Subject: [PATCH 32/47] review comment

---
 imod/mf6/ims.py                            | 14 ++++++++------
 imod/tests/test_mf6/test_mf6_simulation.py |  3 ++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/imod/mf6/ims.py b/imod/mf6/ims.py
index 374530234..b3f619f7e 100644
--- a/imod/mf6/ims.py
+++ b/imod/mf6/ims.py
@@ -428,9 +428,7 @@ def __init__(
         self._validate_init_schemata(validate)
 
     def remove_model_from_solution(self, modelname: str) -> None:
-        models_in_solution = []
-        if "modelnames" in self.dataset.keys():
-            models_in_solution = self.dataset["modelnames"].values
+        models_in_solution = self.get_models_in_solution()
         if modelname not in models_in_solution:
             raise ValueError(
                 f"attempted to remove model {modelname} from solution, but it was not found."
@@ -443,9 +441,7 @@ def remove_model_from_solution(self, modelname: str) -> None:
             self.dataset.update({"modelnames": ("model", filtered_models)})
 
     def add_model_to_solution(self, modelname: str) -> None:
-        models_in_solution = []
-        if "modelnames" in self.dataset.keys():
-            models_in_solution = list(self.dataset["modelnames"].values)
+        models_in_solution = self.get_models_in_solution()
         if modelname in models_in_solution:
             raise ValueError(
                 f"attempted to add model {modelname} to solution, but it was already in it."
@@ -453,6 +449,12 @@ def add_model_to_solution(self, modelname: str) -> None:
         models_in_solution.append(modelname)
         self.dataset.update({"modelnames": ("model", models_in_solution)})
 
+    def get_models_in_solution(self) -> list[str]:
+        models_in_solution = []
+        if "modelnames" in self.dataset.keys():
+            models_in_solution = list(self.dataset["modelnames"].values)
+        return models_in_solution
+
 
 def SolutionPresetSimple(
     modelnames, print_option="summary", csv_output=False, no_ptc=False
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 2d5ebbcf0..b53ec1210 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -395,10 +395,11 @@ def test_split_multiple_models_creates_expected_number_of_exchanges(
             submodel_labels, create_partition_info_mock()
         )
 
+        # fmt: off
         assert exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2  # noqa: E501
         call1 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[0][0]  # noqa: E501
         call2 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[1][0]  # noqa: E501
-
+        # fmt: on
 
         assert call1[0] == "test_model1"
         xr.testing.assert_equal(call1[1], idomain.layer)

From c92df59e114806e89dfc21731b3cd63e6c3efdea Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 31 Jan 2024 14:27:55 +0100
Subject: [PATCH 33/47] skips gwfgwt exchanges when setting options

---
 imod/mf6/simulation.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 1a65b9598..b2e7bd772 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -993,6 +993,8 @@ def _set_exchange_options(self):
             if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
                 model_1 = self[model_name_1]
+            if isinstance(exchange, GWFGWT):
+                continue
             if isinstance(model_1, GroundwaterFlowModel):
                 exchange.set_options(
                     save_flows=model_1["oc"].is_budget_output,

From 2410958ab37a16886380586f9060f1339b817c4e Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Wed, 31 Jan 2024 18:22:57 +0100
Subject: [PATCH 34/47] first model runs

---
 imod/mf6/exchangebase.py                |  2 +-
 imod/mf6/gwtgwt.py                      | 10 ++++--
 imod/mf6/multimodel/exchange_creator.py |  8 +++--
 imod/mf6/pkgbase.py                     |  2 +-
 imod/mf6/simulation.py                  | 45 +++++++++++++++----------
 imod/templates/mf6/exg-gwtgwt.j2        | 37 ++++++++++++++++++++
 6 files changed, 78 insertions(+), 26 deletions(-)
 create mode 100644 imod/templates/mf6/exg-gwtgwt.j2

diff --git a/imod/mf6/exchangebase.py b/imod/mf6/exchangebase.py
index 7a434df5b..ef34f69c6 100644
--- 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
index c0756f669..df079aa03 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -18,11 +18,13 @@ class GWTGWT(ExchangeBase):
     simulation class."""
 
     _auxiliary_data = {"auxiliary_data": "variable"}
-    _pkg_id = "gwfgwf"
+    _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,
@@ -38,8 +40,10 @@ def __init__(
         self.dataset["cell_id1"] = cell_id1
         self.dataset["cell_id2"] = cell_id2
         self.dataset["layer"] = layer
-        self.dataset["model_name_1"] = model_id1
-        self.dataset["model_name_2"] = model_id2
+        self.dataset["model_name_1"] = transport_model_id1
+        self.dataset["model_name_2"] = transport_model_id2        
+        self.dataset["flow_model_name_1"] = flow_model_id1
+        self.dataset["flow_model_name_2"] = flow_model_id2
         self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
         self.dataset["cl1"] = cl1
         self.dataset["cl2"] = cl2
diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index ecc688b1f..796c5419f 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -196,7 +196,7 @@ def create_gwfgwf_exchanges(
         return exchanges
 
     def create_gwtgwt_exchanges(
-        self, model_name: str, layers: GridDataArray
+        self, transport_model_name: str, flow_model_name: str, layers: GridDataArray
     ) -> List[GWTGWT]:
         layers = layers.to_dataframe().filter(["layer"])
 
@@ -254,8 +254,10 @@ def create_gwtgwt_exchanges(
 
                 exchanges.append(
                     GWTGWT(
-                        f"{model_name}_{model_id1}",
-                        f"{model_name}_{model_id2}",
+                        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
index 7c9c69a36..2cc1e5d49 100644
--- a/imod/mf6/pkgbase.py
+++ b/imod/mf6/pkgbase.py
@@ -10,7 +10,7 @@
 from imod.mf6.interfaces.ipackagebase import IPackageBase
 
 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
index 1299930fb..6e9379a17 100644
--- 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
@@ -808,7 +809,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):
@@ -927,14 +927,19 @@ 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
-                )
-            if isinstance(model, GroundwaterTransportModel):
+        flow_models = self.get_models_of_type("gwf6")
+        if len(flow_models) != 1:
+            raise ValueError("splitting of simulations with more (or less) than 1 flow model currently not supported")
+        flow_model_name = list(flow_models.keys())[0]
+        flow_model = self[flow_model_name]
+        exchanges += exchange_creator.create_gwfgwf_exchanges(
+            flow_model_name, flow_model.domain.layer
+        )
+        transport_models =  self.get_models_of_type("gwt6")
+        if len(transport_models)  >  0:
+            for tpt_model_name in transport_models:
                 exchanges += exchange_creator.create_gwtgwt_exchanges(
-                    model_name, model.domain.layer
+                    tpt_model_name, flow_model_name, model.domain.layer
                 )
 
         new_simulation._add_modelsplit_exchanges(exchanges)
@@ -994,19 +999,23 @@ def _add_modelsplit_exchanges(self, exchanges_list: List[GWFGWF]) -> None:
     def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
+
+            if isinstance(exchange, GWFGWT):
+                continue
             if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
-                model_1 = self[model_name_1]
-            if isinstance(exchange, GWFGWT):
+                model_1 = self[model_name_1]            
+                if isinstance(model_1, GroundwaterFlowModel):
+                    exchange.set_options(
+                        save_flows=model_1["oc"].is_budget_output,
+                        dewatered=model_1["npf"].is_dewatered,
+                        variablecv=model_1["npf"].is_variable_vertical_conductance,
+                        xt3d=model_1["npf"].get_xt3d_option(),
+                        newton=model_1.is_use_newton(),
+                    )
+            if isinstance(exchange, GWTGWT):
                 continue
-            if isinstance(model_1, GroundwaterFlowModel):
-                exchange.set_options(
-                    save_flows=model_1["oc"].is_budget_output,
-                    dewatered=model_1["npf"].is_dewatered,
-                    variablecv=model_1["npf"].is_variable_vertical_conductance,
-                    xt3d=model_1["npf"].get_xt3d_option(),
-                    newton=model_1.is_use_newton(),
-                )
+
 
     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
new file mode 100644
index 000000000..3c406c3c9
--- /dev/null
+++ 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

From 8ae078ce46a4d795e239fab9c27abecb0b570ce0 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 11:45:23 +0100
Subject: [PATCH 35/47] added test for output check

---
 .../test_mf6_modelsplitter_transport.py       | 54 +++++++++++++++++--
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
index c91efcf5b..2e048dd3e 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
@@ -3,7 +3,7 @@
 from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
 from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 from imod.typing.grid import zeros_like
-
+import numpy as np
 
 def test_slice_model_structured(flow_transport_simulation):
     # Arrange.
@@ -63,7 +63,53 @@ def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
     assert new_simulation["gwtgwf_exchanges"][6]["model_name_1"].values[()] == "flow_1"
     assert new_simulation["gwtgwf_exchanges"][6]["model_name_2"].values[()] == "tpt_c_1"
 
-    assert new_simulation["gwtgwf_exchanges"][7]["model_name_1"].values[()] == "flow_1"
-    assert new_simulation["gwtgwf_exchanges"][7]["model_name_2"].values[()] == "tpt_d_1"
 
-    assert_simulation_can_run(new_simulation, "dis", tmp_path)
+@pytest.mark.usefixtures("flow_transport_simulation")
+def test_split_flow_and_transport_model_evaluate_output(tmp_path, flow_transport_simulation):
+    
+    simulation = flow_transport_simulation
+
+    flow_model = simulation["flow"]
+    active = flow_model.domain
+    simulation.pop("tpt_a")
+    simulation.pop("tpt_c")
+    simulation.pop("tpt_d")
+    simulation["transport_solver"].remove_model_from_solution("tpt_a")
+    simulation["transport_solver"].remove_model_from_solution("tpt_c")
+    simulation["transport_solver"].remove_model_from_solution("tpt_d")    
+    submodel_labels = zeros_like(active)
+    submodel_labels = submodel_labels.drop_vars("layer")
+    submodel_labels.values[:, :, 50:] = 1
+    submodel_labels = submodel_labels.sel(layer=0, drop=True)
+
+    simulation.write(tmp_path/"original", binary = False )
+    simulation.run ()
+
+    new_simulation = simulation.split(submodel_labels)
+    new_simulation.write(tmp_path, binary=False)
+
+    original_conc = simulation.open_concentration(species_ls=["b"])
+    original_head = simulation.open_head()
+
+    assert len(new_simulation["gwtgwf_exchanges"]) == 2
+
+    assert new_simulation["gwtgwf_exchanges"][0]["model_name_1"].values[()] == "flow_0"
+    assert new_simulation["gwtgwf_exchanges"][0]["model_name_2"].values[()] == "tpt_b_0"
+
+    assert new_simulation["gwtgwf_exchanges"][1]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][1]["model_name_2"].values[()] == "tpt_b_1"
+
+
+    new_simulation.run()
+    conc = new_simulation.open_concentration(species_ls=["a"])
+    head = new_simulation.open_head()
+
+    # Compare the head result of the original simulation with the result of the
+    # partitioned simulation.
+    np.testing.assert_allclose(
+        head.sel(time=2000)["head"].values, original_head.sel(time=200).values, rtol=1e-4, atol=1e-4
+    )
+    print( conc.sel(time = 2000)["concentration"].values - original_conc.sel(time=200).values)
+    np.testing.assert_allclose(
+        conc.sel(time = 2000)["concentration"].values, original_conc.sel(time=200).values, rtol= 6.5, atol=0.011
+    )

From 10bcc8c9be6e585493793647c72329c19d9b3f74 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 13:21:02 +0100
Subject: [PATCH 36/47] formatting

---
 imod/mf6/gwtgwt.py                            |  6 ++--
 imod/mf6/multimodel/exchange_creator.py       |  2 +-
 imod/mf6/simulation.py                        | 12 +++----
 .../flow_transport_simulation_fixture.py      | 10 +++---
 .../test_mf6_modelsplitter_transport.py       | 33 ++++++++++++-------
 5 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
index df079aa03..5cade694c 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -23,8 +23,8 @@ class GWTGWT(ExchangeBase):
 
     def __init__(
         self,
-        transport_model_id1: str, 
-        transport_model_id2: str,              
+        transport_model_id1: str,
+        transport_model_id2: str,
         flow_model_id1: str,
         flow_model_id2: str,
         cell_id1: xr.DataArray,
@@ -41,7 +41,7 @@ def __init__(
         self.dataset["cell_id2"] = cell_id2
         self.dataset["layer"] = layer
         self.dataset["model_name_1"] = transport_model_id1
-        self.dataset["model_name_2"] = transport_model_id2        
+        self.dataset["model_name_2"] = transport_model_id2
         self.dataset["flow_model_name_1"] = flow_model_id1
         self.dataset["flow_model_name_2"] = flow_model_id2
         self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 796c5419f..3273cb01c 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -255,7 +255,7 @@ def create_gwtgwt_exchanges(
                 exchanges.append(
                     GWTGWT(
                         f"{transport_model_name}_{model_id1}",
-                        f"{transport_model_name}_{model_id2}",                        
+                        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/simulation.py b/imod/mf6/simulation.py
index 6e9379a17..40b2dbf4b 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -929,14 +929,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
         exchanges = []
         flow_models = self.get_models_of_type("gwf6")
         if len(flow_models) != 1:
-            raise ValueError("splitting of simulations with more (or less) than 1 flow model currently not supported")
+            raise ValueError(
+                "splitting of simulations with more (or less) than 1 flow model currently not supported"
+            )
         flow_model_name = list(flow_models.keys())[0]
         flow_model = self[flow_model_name]
         exchanges += exchange_creator.create_gwfgwf_exchanges(
             flow_model_name, flow_model.domain.layer
         )
-        transport_models =  self.get_models_of_type("gwt6")
-        if len(transport_models)  >  0:
+        transport_models = self.get_models_of_type("gwt6")
+        if len(transport_models) > 0:
             for tpt_model_name in transport_models:
                 exchanges += exchange_creator.create_gwtgwt_exchanges(
                     tpt_model_name, flow_model_name, model.domain.layer
@@ -999,12 +1001,11 @@ def _add_modelsplit_exchanges(self, exchanges_list: List[GWFGWF]) -> None:
     def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
-
             if isinstance(exchange, GWFGWT):
                 continue
             if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
-                model_1 = self[model_name_1]            
+                model_1 = self[model_name_1]
                 if isinstance(model_1, GroundwaterFlowModel):
                     exchange.set_options(
                         save_flows=model_1["oc"].is_budget_output,
@@ -1016,7 +1017,6 @@ def _set_exchange_options(self):
             if isinstance(exchange, GWTGWT):
                 continue
 
-
     def _filter_inactive_cells_from_exchanges(self) -> None:
         for ex in self["split_exchanges"]:
             for i in [1, 2]:
diff --git a/imod/tests/fixtures/flow_transport_simulation_fixture.py b/imod/tests/fixtures/flow_transport_simulation_fixture.py
index 346ba321f..807453fba 100644
--- 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")
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
index 2e048dd3e..bfb6649a1 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
@@ -1,9 +1,10 @@
+import numpy as np
 import pytest
 
 from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model
 from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 from imod.typing.grid import zeros_like
-import numpy as np
+
 
 def test_slice_model_structured(flow_transport_simulation):
     # Arrange.
@@ -63,10 +64,14 @@ def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation):
     assert new_simulation["gwtgwf_exchanges"][6]["model_name_1"].values[()] == "flow_1"
     assert new_simulation["gwtgwf_exchanges"][6]["model_name_2"].values[()] == "tpt_c_1"
 
+    assert new_simulation["gwtgwf_exchanges"][7]["model_name_1"].values[()] == "flow_1"
+    assert new_simulation["gwtgwf_exchanges"][7]["model_name_2"].values[()] == "tpt_d_1"    
+    assert_simulation_can_run( new_simulation, "dis", tmp_path)
 
 @pytest.mark.usefixtures("flow_transport_simulation")
-def test_split_flow_and_transport_model_evaluate_output(tmp_path, flow_transport_simulation):
-    
+def test_split_flow_and_transport_model_evaluate_output(
+    tmp_path, flow_transport_simulation
+):
     simulation = flow_transport_simulation
 
     flow_model = simulation["flow"]
@@ -76,14 +81,16 @@ def test_split_flow_and_transport_model_evaluate_output(tmp_path, flow_transport
     simulation.pop("tpt_d")
     simulation["transport_solver"].remove_model_from_solution("tpt_a")
     simulation["transport_solver"].remove_model_from_solution("tpt_c")
-    simulation["transport_solver"].remove_model_from_solution("tpt_d")    
+    simulation["transport_solver"].remove_model_from_solution("tpt_d")
     submodel_labels = zeros_like(active)
     submodel_labels = submodel_labels.drop_vars("layer")
-    submodel_labels.values[:, :, 50:] = 1
+
+    submodel_labels.values[:, :, 30:] = 1
+
     submodel_labels = submodel_labels.sel(layer=0, drop=True)
 
-    simulation.write(tmp_path/"original", binary = False )
-    simulation.run ()
+    simulation.write(tmp_path / "original", binary=False)
+    simulation.run()
 
     new_simulation = simulation.split(submodel_labels)
     new_simulation.write(tmp_path, binary=False)
@@ -99,7 +106,6 @@ def test_split_flow_and_transport_model_evaluate_output(tmp_path, flow_transport
     assert new_simulation["gwtgwf_exchanges"][1]["model_name_1"].values[()] == "flow_1"
     assert new_simulation["gwtgwf_exchanges"][1]["model_name_2"].values[()] == "tpt_b_1"
 
-
     new_simulation.run()
     conc = new_simulation.open_concentration(species_ls=["a"])
     head = new_simulation.open_head()
@@ -107,9 +113,14 @@ def test_split_flow_and_transport_model_evaluate_output(tmp_path, flow_transport
     # Compare the head result of the original simulation with the result of the
     # partitioned simulation.
     np.testing.assert_allclose(
-        head.sel(time=2000)["head"].values, original_head.sel(time=200).values, rtol=1e-4, atol=1e-4
+        head.sel(time=2000)["head"].values,
+        original_head.sel(time=200).values,
+        rtol=1e-4,
+        atol=1e-4,
     )
-    print( conc.sel(time = 2000)["concentration"].values - original_conc.sel(time=200).values)
     np.testing.assert_allclose(
-        conc.sel(time = 2000)["concentration"].values, original_conc.sel(time=200).values, rtol= 6.5, atol=0.011
+        conc.sel(time=2000)["concentration"].values,
+        original_conc.sel(time=200).values,
+        rtol=1e-4,
+        atol=0.011,
     )

From 41dc3bf52a224bbff11902992b6e3e95ab83baac Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 15:03:16 +0100
Subject: [PATCH 37/47] fix build error

---
 imod/mf6/multimodel/exchange_creator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 54787c1b0..a3a3c3c78 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -197,7 +197,7 @@ def create_gwfgwf_exchanges(
 
     def create_gwtgwt_exchanges(
         self, transport_model_name: str, flow_model_name: str, layers: GridDataArray
-    ) -> List[GWTGWT]:
+    ) -> list[GWTGWT]:
         layers = layers.to_dataframe().filter(["layer"])
 
         connected_cells_with_geometric_info = pd.merge(

From 181d7d845146a0b14aff367089348220d0e4bb36 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 15:42:05 +0100
Subject: [PATCH 38/47] function for geometric dataframe

---
 imod/mf6/multimodel/exchange_creator.py | 120 +++++++++---------------
 1 file changed, 45 insertions(+), 75 deletions(-)

diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index a3a3c3c78..566ac3962 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -145,46 +145,12 @@ 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"}
-                    )
-                )
-
-                connected_cells = (
-                    connected_domain_pair.merge(mapping1)
-                    .merge(mapping2)
-                    .filter(
-                        [
-                            "cell_id1",
-                            "cell_id2",
-                            "cl1",
-                            "cl2",
-                            "hwva",
-                            "angldegx",
-                            "cdist",
-                        ]
-                    )
-                )
-
-                connected_cells = pd.merge(layers, connected_cells, how="cross")
-
-                connected_cells_dataset = self._to_xarray(connected_cells)
-
-                _adjust_gridblock_indexing(connected_cells_dataset)
-
                 exchanges.append(
                     GWFGWF(
                         f"{model_name}_{model_id1}",
@@ -195,6 +161,45 @@ def create_gwfgwf_exchanges(
 
         return exchanges
 
+    def _collect_geometric_constants_connected_cells(
+        self, model_id1: int, model_id2: int, connected_domain_pair: pd.DataFrame, layers: GridDataArray
+    ):
+        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 = (
+            connected_domain_pair.merge(mapping1)
+            .merge(mapping2)
+            .filter(
+                [
+                    "cell_id1",
+                    "cell_id2",
+                    "cl1",
+                    "cl2",
+                    "hwva",
+                    "angldegx",
+                    "cdist",
+                ]
+            )
+        )
+
+        connected_cells = pd.merge(layers, connected_cells, how="cross")
+
+        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]:
@@ -212,46 +217,11 @@ def create_gwtgwt_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"}
-                    )
-                )
-
-                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_dataset = (
+                    self._collect_geometric_constants_connected_cells(
+                        model_id1, model_id2, connected_domain_pair, layers
                     )
                 )
-
-                connected_cells = (
-                    connected_domain_pair.merge(mapping1)
-                    .merge(mapping2)
-                    .filter(
-                        [
-                            "cell_id1",
-                            "cell_id2",
-                            "cl1",
-                            "cl2",
-                            "hwva",
-                            "angldegx",
-                            "cdist",
-                        ]
-                    )
-                )
-
-                connected_cells = pd.merge(layers, connected_cells, how="cross")
-
-                connected_cells_dataset = self._to_xarray(connected_cells)
-
-                _adjust_gridblock_indexing(connected_cells_dataset)
-
                 exchanges.append(
                     GWTGWT(
                         f"{transport_model_name}_{model_id1}",

From 74ba2929815945ad5001db03c2e04c077fd7d991 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 15:48:11 +0100
Subject: [PATCH 39/47] type hints

---
 imod/mf6/multimodel/exchange_creator.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py
index 566ac3962..2efc69bad 100644
--- a/imod/mf6/multimodel/exchange_creator.py
+++ b/imod/mf6/multimodel/exchange_creator.py
@@ -162,8 +162,12 @@ def create_gwfgwf_exchanges(
         return exchanges
 
     def _collect_geometric_constants_connected_cells(
-        self, model_id1: int, model_id2: int, connected_domain_pair: pd.DataFrame, layers: GridDataArray
-    ):
+        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"])

From df50f22af65b65786985879fde1b145dddc604a5 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Thu, 1 Feb 2024 16:02:50 +0100
Subject: [PATCH 40/47] cleanup

---
 imod/mf6/simulation.py                        | 17 ++++++------
 .../test_mf6_modelsplitter_transport.py       | 27 +++++++------------
 2 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 72d4c4e58..bf59208c8 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1015,19 +1015,20 @@ def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
             if isinstance(exchange, GWFGWT):
+                # that package does not have options
                 continue
             if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
                 model_1 = self[model_name_1]
-                if isinstance(model_1, GroundwaterFlowModel):
-                    exchange.set_options(
-                        save_flows=model_1["oc"].is_budget_output,
-                        dewatered=model_1["npf"].is_dewatered,
-                        variablecv=model_1["npf"].is_variable_vertical_conductance,
-                        xt3d=model_1["npf"].get_xt3d_option(),
-                        newton=model_1.is_use_newton(),
-                    )
+                exchange.set_options(
+                    save_flows=model_1["oc"].is_budget_output,
+                    dewatered=model_1["npf"].is_dewatered,
+                    variablecv=model_1["npf"].is_variable_vertical_conductance,
+                    xt3d=model_1["npf"].get_xt3d_option(),
+                    newton=model_1.is_use_newton(),
+                )
             if isinstance(exchange, GWTGWT):
+                # TODO: issue #747
                 continue
 
     def _filter_inactive_cells_from_exchanges(self) -> None:
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
index 5f971f14d..3a2b777ef 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter_transport.py
@@ -77,42 +77,35 @@ def test_split_flow_and_transport_model_evaluate_output(
 
     flow_model = simulation["flow"]
     active = flow_model.domain
+
+    # TODO: put the other transport models back when #797 is solved
     simulation.pop("tpt_a")
     simulation.pop("tpt_c")
     simulation.pop("tpt_d")
     simulation["transport_solver"].remove_model_from_solution("tpt_a")
     simulation["transport_solver"].remove_model_from_solution("tpt_c")
     simulation["transport_solver"].remove_model_from_solution("tpt_d")
+
+    # create label array
     submodel_labels = zeros_like(active)
     submodel_labels = submodel_labels.drop_vars("layer")
-
     submodel_labels.values[:, :, 30:] = 1
-
     submodel_labels = submodel_labels.sel(layer=0, drop=True)
 
+    # for reference run the original model and load the results
     simulation.write(tmp_path / "original", binary=False)
     simulation.run()
-
-    new_simulation = simulation.split(submodel_labels)
-    new_simulation.write(tmp_path, binary=False)
-
     original_conc = simulation.open_concentration(species_ls=["b"])
     original_head = simulation.open_head()
 
-    assert len(new_simulation["gwtgwf_exchanges"]) == 2
-
-    assert new_simulation["gwtgwf_exchanges"][0]["model_name_1"].values[()] == "flow_0"
-    assert new_simulation["gwtgwf_exchanges"][0]["model_name_2"].values[()] == "tpt_b_0"
-
-    assert new_simulation["gwtgwf_exchanges"][1]["model_name_1"].values[()] == "flow_1"
-    assert new_simulation["gwtgwf_exchanges"][1]["model_name_2"].values[()] == "tpt_b_1"
-
+    # split the model , run the split model and load the results
+    new_simulation = simulation.split(submodel_labels)
+    new_simulation.write(tmp_path, binary=False)
     new_simulation.run()
-    conc = new_simulation.open_concentration(species_ls=["a"])
+    conc = new_simulation.open_concentration(species_ls=["b"])
     head = new_simulation.open_head()
 
-    # Compare the head result of the original simulation with the result of the
-    # partitioned simulation.
+    # Compare
     np.testing.assert_allclose(
         head.sel(time=2000)["head"].values,
         original_head.sel(time=200).values,

From 9410c3020c2b1222969500d45e6b72581455012a Mon Sep 17 00:00:00 2001
From: luitjansl <64598682+luitjansl@users.noreply.github.com>
Date: Fri, 2 Feb 2024 11:23:10 +0100
Subject: [PATCH 41/47] Update imod/mf6/gwtgwt.py

review comment

Co-authored-by: Joeri van Engelen <joerivanengelen@hotmail.com>
---
 imod/mf6/gwtgwt.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
index 5cade694c..bc80698a0 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -15,7 +15,8 @@ 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."""
+    simulation class.
+    """
 
     _auxiliary_data = {"auxiliary_data": "variable"}
     _pkg_id = "gwtgwt"

From 3481c6d25da1e7ca691f68c7d4ce8c9f1eb03cfd Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 2 Feb 2024 11:31:52 +0100
Subject: [PATCH 42/47] review comments

---
 imod/mf6/gwtgwt.py     | 26 ++++++++++++++------------
 imod/mf6/simulation.py |  3 ---
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
index bc80698a0..cf801fc2c 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -37,18 +37,20 @@ def __init__(
         angldegx: Optional[xr.DataArray] = None,
         cdist: Optional[xr.DataArray] = None,
     ):
-        super().__init__(locals())
-        self.dataset["cell_id1"] = cell_id1
-        self.dataset["cell_id2"] = cell_id2
-        self.dataset["layer"] = layer
-        self.dataset["model_name_1"] = transport_model_id1
-        self.dataset["model_name_2"] = transport_model_id2
-        self.dataset["flow_model_name_1"] = flow_model_id1
-        self.dataset["flow_model_name_2"] = flow_model_id2
-        self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
-        self.dataset["cl1"] = cl1
-        self.dataset["cl2"] = cl2
-        self.dataset["hwva"] = hwva
+        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:
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index bf59208c8..c100230e4 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1014,9 +1014,6 @@ def _add_modelsplit_exchanges(self, exchanges_list: list[GWFGWF]) -> None:
     def _set_exchange_options(self):
         # collect some options that we will auto-set
         for exchange in self["split_exchanges"]:
-            if isinstance(exchange, GWFGWT):
-                # that package does not have options
-                continue
             if isinstance(exchange, GWFGWF):
                 model_name_1 = exchange.dataset["model_name_1"].values[()]
                 model_1 = self[model_name_1]

From 5139fffdb8e449323c57019753684d20a4b40eb8 Mon Sep 17 00:00:00 2001
From: luitjansl <64598682+luitjansl@users.noreply.github.com>
Date: Fri, 2 Feb 2024 11:36:49 +0100
Subject: [PATCH 43/47] Update imod/mf6/simulation.py

review comment

Co-authored-by: Joeri van Engelen <joerivanengelen@hotmail.com>
---
 imod/mf6/simulation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index c100230e4..077b894c7 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1024,7 +1024,7 @@ def _set_exchange_options(self):
                     xt3d=model_1["npf"].get_xt3d_option(),
                     newton=model_1.is_use_newton(),
                 )
-            if isinstance(exchange, GWTGWT):
+            elif isinstance(exchange, GWTGWT):
                 # TODO: issue #747
                 continue
 

From 4f17342e9fe56dc3b3697fcae33f77e6c5673606 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 2 Feb 2024 11:42:06 +0100
Subject: [PATCH 44/47] review comments

---
 imod/mf6/gwtgwt.py     | 6 +++---
 imod/mf6/simulation.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/imod/mf6/gwtgwt.py b/imod/mf6/gwtgwt.py
index cf801fc2c..11d931da1 100644
--- a/imod/mf6/gwtgwt.py
+++ b/imod/mf6/gwtgwt.py
@@ -43,14 +43,14 @@ def __init__(
             "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,            
+            "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)        
+        super().__init__(dict_dataset)
 
         auxiliary_variables = [var for var in [angldegx, cdist] if var is not None]
         if auxiliary_variables:
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 077b894c7..38d89e028 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -952,7 +952,7 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
             flow_model_name, flow_model.domain.layer
         )
         transport_models = self.get_models_of_type("gwt6")
-        if len(transport_models) > 0:
+        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

From e3cfe8d38c7964035f9e7cdb552113fcda2432fc Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Fri, 2 Feb 2024 16:52:03 +0100
Subject: [PATCH 45/47] fix failing tests

---
 imod/mf6/simulation.py                     |  24 +--
 imod/tests/test_mf6/test_mf6_simulation.py | 165 +++++----------------
 2 files changed, 52 insertions(+), 137 deletions(-)

diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 38d89e028..76ae5e8b9 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -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,17 +951,11 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation:
                 new_simulation[solution_name].add_model_to_solution(new_model_name)
 
         exchanges = []
-        flow_models = self.get_models_of_type("gwf6")
-        if len(flow_models) != 1:
-            raise ValueError(
-                "splitting of simulations with more (or less) than 1 flow model currently not supported"
+
+        for flow_model_name, flow_model in flow_models.items():
+            exchanges += exchange_creator.create_gwfgwf_exchanges(
+                flow_model_name, flow_model.domain.layer
             )
-        flow_model_name = list(flow_models.keys())[0]
-        flow_model = self[flow_model_name]
-        exchanges += exchange_creator.create_gwfgwf_exchanges(
-            flow_model_name, flow_model.domain.layer
-        )
-        transport_models = self.get_models_of_type("gwt6")
         if any(transport_models):
             for tpt_model_name in transport_models:
                 exchanges += exchange_creator.create_gwtgwt_exchanges(
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index b53ec1210..3038ce537 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -2,6 +2,7 @@
 import re
 import sys
 import textwrap
+from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
 from unittest import mock
@@ -16,9 +17,9 @@
 from imod.mf6.model import Modflow6Model
 from imod.mf6.model_gwf import GroundwaterFlowModel
 from imod.mf6.multimodel.modelsplitter import PartitionInfo
-from imod.mf6.simulation import get_models, get_packages
 from imod.mf6.statusinfo import NestedStatusInfo, StatusInfo
 from imod.schemata import ValidationError
+from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run
 from imod.typing.grid import zeros_like
 
 
@@ -274,138 +275,48 @@ def test_split_simulation_only_has_packages(
         submodel_labels = xu.zeros_like(active).where(active.grid.face_y > 0.0, 1)
 
         # Act.
-        new_simulation = simulation.split(submodel_labels)
+        with pytest.raises(ValueError):
+            _ = simulation.split(submodel_labels)
 
-        # Assert.
-        assert len(get_models(new_simulation)) == 0
-        assert len(get_packages(new_simulation)) == 3
-        assert new_simulation["solver"] is simulation["solver"]
-        assert (
-            new_simulation["time_discretization"] is simulation["time_discretization"]
-        )
-        assert new_simulation["disv"] is simulation["disv"]
-
-    @mock.patch("imod.mf6.simulation.slice_model", autospec=True)
-    @mock.patch("imod.mf6.simulation.ExchangeCreator_Unstructured")
-    def test_split_multiple_models(
-        self,
-        exchange_creator_unstructured_mock,
-        slice_model_mock,
-        circle_dis,
-        setup_simulation,
-    ):
+    def test_split_multiple_models(self, tmp_path, circle_model):
         # Arrange.
-        idomain, _, _ = circle_dis
-
-        simulation = setup_simulation
-
-        model_mock1 = MagicMock(spec_set=GroundwaterFlowModel)
-        model_mock1._model_id = "test_model_id1"
-
-        model_mock2 = MagicMock(spec_set=GroundwaterFlowModel)
-        model_mock2._model_id = "test_model_id2"
+        oc2 = deepcopy(circle_model["GWF_1"]["oc"])
+        npf2 = deepcopy(circle_model["GWF_1"]["npf"])
+        disv2 = deepcopy(circle_model["GWF_1"]["disv"])
+        sto2 = deepcopy(circle_model["GWF_1"]["sto"])
+        chd2 = deepcopy(circle_model["GWF_1"]["chd"])
+        rch2 = deepcopy(circle_model["GWF_1"]["rch"])
+        ic2 = deepcopy(circle_model["GWF_1"]["ic"])
+        gwf_2 = GroundwaterFlowModel()
+        gwf_2["oc"] = oc2
+        gwf_2["npf"] = npf2
+        gwf_2["disv"] = disv2
+        gwf_2["sto"] = sto2
+        gwf_2["chd"] = chd2
+        gwf_2["rch"] = rch2
+        gwf_2["ic"] = ic2
+        circle_model["GWF_2"] = gwf_2
+        circle_model["solver"].add_model_to_solution("GWF_2")
+
+        active = circle_model["GWF_1"].domain.sel(layer=1)
+        submodel_labels = xu.zeros_like(active)
+        submodel_labels.values[90:] = 1
+
+        new_simulation = circle_model.split(submodel_labels)
 
-        simulation["test_model1"] = model_mock1
-        simulation["test_model2"] = model_mock2
-
-        simulation["solver"].dataset = xr.Dataset(
-            {"modelnames": ["test_model1", "test_model2"]}
+        assert (
+            new_simulation["split_exchanges"][0]["model_name_1"].values[()] == "GWF_1_0"
         )
-
-        slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
-
-        active = idomain.sel(layer=1)
-        submodel_labels = xu.zeros_like(active).where(active.grid.face_y > 0.0, 1)
-
-        # Act.
-        new_simulation = simulation.split(submodel_labels)
-
-        # Assert.
-        new_models = get_models(new_simulation)
-        assert slice_model_mock.call_count == 4
-        assert len(new_models) == 4
-
-        # fmt: off
-        assert len([model_name for model_name in new_models.keys() if "test_model1" in model_name]) == 2
-        assert len([model_name for model_name in new_models.keys() if "test_model2" in model_name]) == 2
-
-        active_domain1 = submodel_labels.where(submodel_labels == 0, 0).where(submodel_labels != 0, 1)
-        active_domain2 = submodel_labels.where(submodel_labels == 1, 0).where(submodel_labels != 1, 1)
-        # fmt: on
-
-        expected_slice_model_calls = [
-            (PartitionInfo(id=0, active_domain=active_domain1), model_mock1),
-            (PartitionInfo(id=0, active_domain=active_domain1), model_mock2),
-            (PartitionInfo(id=1, active_domain=active_domain2), model_mock1),
-            (PartitionInfo(id=1, active_domain=active_domain2), model_mock2),
-        ]
-
-        for expected_call in expected_slice_model_calls:
-            assert any(
-                compare_submodel_partition_info(expected_call[0], call_args[0][0])
-                and (expected_call[1] is call_args[0][1])
-                for call_args in slice_model_mock.call_args_list
-            )
-
-    @mock.patch("imod.mf6.simulation.slice_model", autospec=True)
-    @mock.patch("imod.mf6.simulation.ExchangeCreator_Structured", autospec=True)
-    @mock.patch("imod.mf6.simulation.create_partition_info")
-    def test_split_multiple_models_creates_expected_number_of_exchanges(
-        self,
-        create_partition_info_mock,
-        exchange_creator_mock,
-        slice_model_mock,
-        basic_dis,
-        setup_simulation,
-    ):
-        # Arrange.
-        idomain, top, bottom = basic_dis
-
-        simulation = setup_simulation
-
-        model_mock1 = MagicMock(spec_set=GroundwaterFlowModel)
-        model_mock1._model_id = "test_model_id1"
-        model_mock1.domain = idomain
-
-        model_mock2 = MagicMock(spec_set=GroundwaterFlowModel)
-        model_mock2._model_id = "test_model_id2"
-        model_mock2.domain = idomain
-
-        simulation["test_model1"] = model_mock1
-        simulation["test_model2"] = model_mock2
-
-        simulation["solver"].dataset = xr.Dataset(
-            {"modelnames": ["test_model1", "test_model2"]}
+        assert (
+            new_simulation["split_exchanges"][0]["model_name_2"].values[()] == "GWF_1_1"
         )
-
-        slice_model_mock.return_value = MagicMock(spec_set=GroundwaterFlowModel)
-
-        active = idomain.sel(layer=1)
-        submodel_labels = xr.zeros_like(active).where(active.y > 50, 1)
-
-        create_partition_info_mock.return_value = [
-            PartitionInfo(id=0, active_domain=xr.DataArray(0)),
-            PartitionInfo(id=1, active_domain=xr.DataArray(1)),
-        ]
-        # Act.
-        _ = simulation.split(submodel_labels)
-
-        # Assert.
-        exchange_creator_mock.assert_called_with(
-            submodel_labels, create_partition_info_mock()
+        assert (
+            new_simulation["split_exchanges"][1]["model_name_1"].values[()] == "GWF_2_0"
         )
-
-        # fmt: off
-        assert exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_count == 2  # noqa: E501
-        call1 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[0][0]  # noqa: E501
-        call2 = exchange_creator_mock.return_value.create_gwfgwf_exchanges.call_args_list[1][0]  # noqa: E501
-        # fmt: on
-
-        assert call1[0] == "test_model1"
-        xr.testing.assert_equal(call1[1], idomain.layer)
-
-        assert call2[0] == "test_model2"
-        xr.testing.assert_equal(call2[1], idomain.layer)
+        assert (
+            new_simulation["split_exchanges"][1]["model_name_2"].values[()] == "GWF_2_1"
+        )
+        assert_simulation_can_run(new_simulation, "disv", tmp_path)
 
     @pytest.mark.usefixtures("transient_twri_model")
     def test_exchanges_in_simulation_file(self, transient_twri_model, tmp_path):

From 71d55e4139c9d6087357caf263452df0142211a7 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Mon, 5 Feb 2024 11:15:42 +0100
Subject: [PATCH 46/47] cleanup

---
 imod/tests/test_mf6/test_mf6_simulation.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 3038ce537..be04e1354 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -302,8 +302,10 @@ def test_split_multiple_models(self, tmp_path, circle_model):
         submodel_labels = xu.zeros_like(active)
         submodel_labels.values[90:] = 1
 
+        # Act
         new_simulation = circle_model.split(submodel_labels)
-
+ 
+        # Assert
         assert (
             new_simulation["split_exchanges"][0]["model_name_1"].values[()] == "GWF_1_0"
         )

From 2df1dc517e21cfa57c3b1779d9912f44c9a305b4 Mon Sep 17 00:00:00 2001
From: luitjan <luitjan.slooten@deltares.nl>
Date: Mon, 5 Feb 2024 11:21:43 +0100
Subject: [PATCH 47/47] formatting

---
 imod/tests/test_mf6/test_mf6_simulation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index be04e1354..be9439f63 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -304,7 +304,7 @@ def test_split_multiple_models(self, tmp_path, circle_model):
 
         # Act
         new_simulation = circle_model.split(submodel_labels)
- 
+
         # Assert
         assert (
             new_simulation["split_exchanges"][0]["model_name_1"].values[()] == "GWF_1_0"