diff --git a/imod/mf6/exchangebase.py b/imod/mf6/exchangebase.py index efaa47bb4..d3b0dd03b 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 new file mode 100644 index 000000000..11d931da1 --- /dev/null +++ b/imod/mf6/gwtgwt.py @@ -0,0 +1,92 @@ +from typing import Optional + +import cftime +import numpy as np +import xarray as xr + +from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables +from imod.mf6.exchangebase import ExchangeBase +from imod.mf6.package import Package +from imod.typing import GridDataArray + + +class GWTGWT(ExchangeBase): + """ + This package is for writing an exchange file, used for splitting up a model + into different submodels (that can be solved in parallel). It (usually) + is not instantiated by users, but created by the "split" method of the + simulation class. + """ + + _auxiliary_data = {"auxiliary_data": "variable"} + _pkg_id = "gwtgwt" + _template = Package._initialize_template(_pkg_id) + + def __init__( + self, + transport_model_id1: str, + transport_model_id2: str, + flow_model_id1: str, + flow_model_id2: str, + cell_id1: xr.DataArray, + cell_id2: xr.DataArray, + layer: xr.DataArray, + cl1: xr.DataArray, + cl2: xr.DataArray, + hwva: xr.DataArray, + angldegx: Optional[xr.DataArray] = None, + cdist: Optional[xr.DataArray] = None, + ): + dict_dataset = { + "cell_id1": cell_id1, + "cell_id2": cell_id2, + "layer": layer, + "model_name_1": transport_model_id1, + "model_name_2": transport_model_id2, + "flow_model_name_1": flow_model_id1, + "flow_model_name_2": flow_model_id2, + "ihc": xr.DataArray(np.ones_like(cl1, dtype=int)), + "cl1": cl1, + "cl2": cl2, + "hwva": hwva, + } + super().__init__(dict_dataset) + + auxiliary_variables = [var for var in [angldegx, cdist] if var is not None] + if auxiliary_variables: + self.dataset["auxiliary_data"] = xr.merge(auxiliary_variables).to_array( + name="auxiliary_data" + ) + expand_transient_auxiliary_variables(self) + + def set_options( + self, + print_input: Optional[bool] = None, + print_flows: Optional[bool] = None, + save_flows: Optional[bool] = None, + adv_scheme: Optional[str] = None, + dsp_xt3d_off: Optional[bool] = None, + dsp_xt3d_rhs: Optional[bool] = None, + ): + self.dataset["print_input"] = print_input + self.dataset["print_flows"] = print_flows + self.dataset["save_flows"] = save_flows + self.dataset["adv_scheme"] = adv_scheme + self.dataset["dsp_xt3d_off"] = dsp_xt3d_off + self.dataset["dsp_xt3d_rhs"] = dsp_xt3d_rhs + + def clip_box( + self, + time_min: Optional[cftime.datetime | np.datetime64 | str] = None, + time_max: Optional[cftime.datetime | np.datetime64 | str] = None, + layer_min: Optional[int] = None, + layer_max: Optional[int] = None, + x_min: Optional[float] = None, + x_max: Optional[float] = None, + y_min: Optional[float] = None, + y_max: Optional[float] = None, + top: Optional[GridDataArray] = None, + bottom: Optional[GridDataArray] = None, + state_for_boundary: Optional[GridDataArray] = None, + ) -> Package: + raise NotImplementedError("this package cannot be clipped") diff --git a/imod/mf6/multimodel/exchange_creator.py b/imod/mf6/multimodel/exchange_creator.py index bbde93867..2efc69bad 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 @@ -144,50 +145,93 @@ def create_gwfgwf_exchanges( for model_id2, connected_domain_pair in grouped_connected_models.groupby( "cell_label2" ): - model_id1 = int(model_id1) - model_id2 = int(model_id2) - mapping1 = ( - self._global_to_local_mapping[model_id1] - .drop(columns=["local_idx"]) - .rename( - columns={"global_idx": "cell_idx1", "local_cell_id": "cell_id1"} + connected_cells_dataset = ( + self._collect_geometric_constants_connected_cells( + model_id1, model_id2, connected_domain_pair, layers ) ) - mapping2 = ( - self._global_to_local_mapping[model_id2] - .drop(columns=["local_idx"]) - .rename( - columns={"global_idx": "cell_idx2", "local_cell_id": "cell_id2"} + exchanges.append( + GWFGWF( + f"{model_name}_{model_id1}", + f"{model_name}_{model_id2}", + **connected_cells_dataset, ) ) - connected_cells = ( - connected_domain_pair.merge(mapping1) - .merge(mapping2) - .filter( - [ - "cell_id1", - "cell_id2", - "cl1", - "cl2", - "hwva", - "angldegx", - "cdist", - ] - ) - ) + return exchanges + + def _collect_geometric_constants_connected_cells( + self, + model_id1: int, + model_id2: int, + connected_domain_pair: pd.DataFrame, + layers: GridDataArray, + ) -> xr.Dataset: + mapping1 = ( + self._global_to_local_mapping[model_id1] + .drop(columns=["local_idx"]) + .rename(columns={"global_idx": "cell_idx1", "local_cell_id": "cell_id1"}) + ) + + mapping2 = ( + self._global_to_local_mapping[model_id2] + .drop(columns=["local_idx"]) + .rename(columns={"global_idx": "cell_idx2", "local_cell_id": "cell_id2"}) + ) - connected_cells = pd.merge(layers, connected_cells, how="cross") + connected_cells = ( + connected_domain_pair.merge(mapping1) + .merge(mapping2) + .filter( + [ + "cell_id1", + "cell_id2", + "cl1", + "cl2", + "hwva", + "angldegx", + "cdist", + ] + ) + ) - connected_cells_dataset = self._to_xarray(connected_cells) + connected_cells = pd.merge(layers, connected_cells, how="cross") - _adjust_gridblock_indexing(connected_cells_dataset) + connected_cells_dataset = self._to_xarray(connected_cells) + + _adjust_gridblock_indexing(connected_cells_dataset) + + return connected_cells_dataset + + def create_gwtgwt_exchanges( + self, transport_model_name: str, flow_model_name: str, layers: GridDataArray + ) -> list[GWTGWT]: + layers = layers.to_dataframe().filter(["layer"]) + connected_cells_with_geometric_info = pd.merge( + self._connected_cells, self._geometric_information + ) + + exchanges = [] + for ( + model_id1, + grouped_connected_models, + ) in connected_cells_with_geometric_info.groupby("cell_label1"): + for model_id2, connected_domain_pair in grouped_connected_models.groupby( + "cell_label2" + ): + connected_cells_dataset = ( + self._collect_geometric_constants_connected_cells( + model_id1, model_id2, connected_domain_pair, layers + ) + ) exchanges.append( - GWFGWF( - f"{model_name}_{model_id1}", - f"{model_name}_{model_id2}", + GWTGWT( + f"{transport_model_name}_{model_id1}", + f"{transport_model_name}_{model_id2}", + f"{flow_model_name}_{model_id1}", + f"{flow_model_name}_{model_id2}", **connected_cells_dataset, ) ) diff --git a/imod/mf6/pkgbase.py b/imod/mf6/pkgbase.py index fe59e0781..d144dfb8c 100644 --- a/imod/mf6/pkgbase.py +++ b/imod/mf6/pkgbase.py @@ -12,7 +12,7 @@ from imod.typing.grid import GridDataArray, GridDataset, merge_with_dictionary TRANSPORT_PACKAGES = ("adv", "dsp", "ssm", "mst", "ist", "src") -EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt") +EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt", "gwtgwt") class PackageBase(IPackageBase, abc.ABC): diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index b418f9231..76ae5e8b9 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 @@ -822,7 +823,6 @@ def get_exchange_relationships(self): if is_split(self): for exchange in self["split_exchanges"]: result.append(exchange.get_specification()) - return result def get_models_of_type(self, modeltype): @@ -912,6 +912,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation: "Unable to split simulation. Splitting can only be done on simulations that haven't been split." ) + flow_models = self.get_models_of_type("gwf6") + transport_models = self.get_models_of_type("gwt6") + if any(transport_models) and len(flow_models) != 1: + raise ValueError( + "splitting of simulations with more (or less) than 1 flow model currently not supported, if a transport model is present" + ) + + if not any(flow_models) and not any(transport_models): + raise ValueError("a simulation without any models cannot be split.") + original_models = get_models(self) original_packages = get_packages(self) @@ -941,12 +951,16 @@ def split(self, submodel_labels: xr.DataArray) -> Modflow6Simulation: new_simulation[solution_name].add_model_to_solution(new_model_name) exchanges = [] - for model_name, model in original_models.items(): - if isinstance(model, GroundwaterFlowModel): - exchanges += exchange_creator.create_gwfgwf_exchanges( - model_name, model.domain.layer - ) + for flow_model_name, flow_model in flow_models.items(): + exchanges += exchange_creator.create_gwfgwf_exchanges( + flow_model_name, flow_model.domain.layer + ) + if any(transport_models): + for tpt_model_name in transport_models: + exchanges += exchange_creator.create_gwtgwt_exchanges( + tpt_model_name, flow_model_name, model.domain.layer + ) new_simulation._add_modelsplit_exchanges(exchanges) new_simulation._set_exchange_options() @@ -1014,6 +1028,9 @@ def _set_exchange_options(self): xt3d=model_1["npf"].get_xt3d_option(), newton=model_1.is_use_newton(), ) + elif isinstance(exchange, GWTGWT): + # TODO: issue #747 + continue def _filter_inactive_cells_from_exchanges(self) -> None: for ex in self["split_exchanges"]: diff --git a/imod/templates/mf6/exg-gwtgwt.j2 b/imod/templates/mf6/exg-gwtgwt.j2 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 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_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py index b53ec1210..be9439f63 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,50 @@ 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 + 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 - 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" - - simulation["test_model1"] = model_mock1 - simulation["test_model2"] = model_mock2 + # Act + new_simulation = circle_model.split(submodel_labels) - simulation["solver"].dataset = xr.Dataset( - {"modelnames": ["test_model1", "test_model2"]} + # Assert + 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): 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..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 @@ -1,3 +1,4 @@ +import numpy as np import pytest from imod.mf6.multimodel.modelsplitter import create_partition_info, slice_model @@ -65,5 +66,55 @@ def test_split_flow_and_transport_model(tmp_path, flow_transport_simulation): 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 + + # 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() + original_conc = simulation.open_concentration(species_ls=["b"]) + original_head = simulation.open_head() + + # 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=["b"]) + head = new_simulation.open_head() + + # Compare + np.testing.assert_allclose( + head.sel(time=2000)["head"].values, + original_head.sel(time=200).values, + rtol=1e-4, + atol=1e-4, + ) + np.testing.assert_allclose( + conc.sel(time=2000)["concentration"].values, + original_conc.sel(time=200).values, + rtol=1e-4, + atol=0.011, + )