diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 2fc2228bf..6d51baed9 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -17,6 +17,7 @@ Added :meth:`imod.mf6.Modflow6Simulation.clip_box` to consider a package empty if its first times step is all nodata. This can save a lot of clipping or masking transient models with many timesteps. +- Added :meth:`imod.msw.MetaSwapModel.split` to split MetaSWAP models. Fixed ~~~~~ @@ -53,7 +54,7 @@ Changed - ``proportion_depth`` and ``proportion_rate`` in :class:`imod.mf6.Evapotranspiration` are now optional variables. If provided, now require ``"segment"`` dimension when ``proportion_depth`` and - ``proportion_rate``. + ``proportion_rate``. [1.0.0] - 2025-11-11 @@ -132,7 +133,7 @@ Fixed - Fixed bug where :meth:`imod.mf6.Modflow6Simulation.split`, :meth:`imod.mf6.Modflow6Simulation.regrid_like`, and :meth:`imod.mf6.Modflow6Simulation.clip_box` would not copy - :class:`imod.mf6.ValidationSettings`. + :class:`imod.mf6.ValidationSettings`. - ``landuse``, ``soil_physical_unit``, ``active`` for :class:`imod.msw.GridData` are now properly regridded with the ``mode`` statistic when using :meth:`imod.msw.GridData.regrid_like`. @@ -189,7 +190,7 @@ Fixed ``states_for_boundary`` argument would place these bc at the incorrect places with unstructured grids. - Fixed bug where :meth:`imod.mf6.SourceSinkMixing.from_flow_model` would return - an error upon adding a package which cannot have a ``concentration``, such as + an error upon adding a package which cannot have a ``concentration``, such as :class:`imod.mf6.HorizontalFlowBarrierResistance`. - Broken names for ``outer_csvfile`` and ``inner_csvfile`` in the :class:`imod.mf6.Solution` MODFLOW 6 template file. @@ -343,8 +344,8 @@ Changed :meth:`imod.mf6.Recharge.from_imod5_data` got extra arguments for ``period_data``, ``time_min`` and ``time_max``. - :func:`imod.visualize.read_imod_legend` now also returns the labels as an extra - argument. Update your code by changing - ``colors, levels = read_imod_legend(...)`` to + argument. Update your code by changing + ``colors, levels = read_imod_legend(...)`` to ``colors, levels, labels = read_imod_legend(...)``. diff --git a/imod/common/utilities/partitioninfo.py b/imod/common/utilities/partitioninfo.py new file mode 100644 index 000000000..0bd89c419 --- /dev/null +++ b/imod/common/utilities/partitioninfo.py @@ -0,0 +1,53 @@ +from typing import List, NamedTuple + +import numpy as np + +from imod.typing import GridDataArray +from imod.typing.grid import ones_like + + +class PartitionInfo(NamedTuple): + active_domain: GridDataArray + id: int + + +def create_partition_info(submodel_labels: GridDataArray) -> List[PartitionInfo]: + """ + A PartitionInfo is used to partition a model or package. The partition + info's of a domain are created using a submodel_labels array. The + submodel_labels provided as input should have the same shape as a single + layer of the model grid (all layers are split the same way), and contains an + integer value in each cell. Each cell in the model grid will end up in the + submodel with the index specified by the corresponding label of that cell. + The labels should be numbers between 0 and the number of partitions. + """ + _validate_submodel_label_array(submodel_labels) + + unique_labels = np.unique(submodel_labels.values) + + partition_infos = [] + for label_id in unique_labels: + active_domain = submodel_labels.where(submodel_labels.values == label_id) + active_domain = ones_like(active_domain).where(active_domain.notnull(), 0) + active_domain = active_domain.astype(submodel_labels.dtype) + + submodel_partition_info = PartitionInfo( + id=label_id, active_domain=active_domain + ) + partition_infos.append(submodel_partition_info) + + return partition_infos + + +def _validate_submodel_label_array(submodel_labels: GridDataArray) -> None: + unique_labels = np.unique(submodel_labels.values) + + if not ( + len(unique_labels) == unique_labels.max() + 1 + and unique_labels.min() == 0 + and np.issubdtype(submodel_labels.dtype, np.integer) + ): + raise ValueError( + "The submodel_label array should be integer and contain all the numbers between 0 and the number of " + "partitions minus 1." + ) diff --git a/imod/mf6/multimodel/modelsplitter.py b/imod/mf6/multimodel/modelsplitter.py index 45fac27ea..fd6e080c4 100644 --- a/imod/mf6/multimodel/modelsplitter.py +++ b/imod/mf6/multimodel/modelsplitter.py @@ -1,11 +1,10 @@ -from typing import List, NamedTuple, cast - -import numpy as np +from typing import List, cast from imod.common.interfaces.iagnosticpackage import IAgnosticPackage from imod.common.interfaces.imodel import IModel from imod.common.interfaces.ipackage import IPackage from imod.common.utilities.clip import clip_by_grid +from imod.common.utilities.partitioninfo import PartitionInfo from imod.mf6.boundary_condition import BoundaryCondition from imod.mf6.ims import Solution from imod.mf6.model_gwf import GroundwaterFlowModel @@ -13,57 +12,7 @@ from imod.mf6.ssm import SourceSinkMixing from imod.mf6.validation_settings import ValidationSettings, trim_time_dimension from imod.typing import GridDataArray -from imod.typing.grid import ( - get_non_spatial_dimension_names, - ones_like, -) - - -class PartitionInfo(NamedTuple): - active_domain: GridDataArray - id: int - - -def create_partition_info(submodel_labels: GridDataArray) -> List[PartitionInfo]: - """ - A PartitionInfo is used to partition a model or package. The partition - info's of a domain are created using a submodel_labels array. The - submodel_labels provided as input should have the same shape as a single - layer of the model grid (all layers are split the same way), and contains an - integer value in each cell. Each cell in the model grid will end up in the - submodel with the index specified by the corresponding label of that cell. - The labels should be numbers between 0 and the number of partitions. - """ - _validate_submodel_label_array(submodel_labels) - - unique_labels = np.unique(submodel_labels.values) - - partition_infos = [] - for label_id in unique_labels: - active_domain = submodel_labels.where(submodel_labels.values == label_id) - active_domain = ones_like(active_domain).where(active_domain.notnull(), 0) - active_domain = active_domain.astype(submodel_labels.dtype) - - submodel_partition_info = PartitionInfo( - id=label_id, active_domain=active_domain - ) - partition_infos.append(submodel_partition_info) - - return partition_infos - - -def _validate_submodel_label_array(submodel_labels: GridDataArray) -> None: - unique_labels = np.unique(submodel_labels.values) - - if not ( - len(unique_labels) == unique_labels.max() + 1 - and unique_labels.min() == 0 - and np.issubdtype(submodel_labels.dtype, np.integer) - ): - raise ValueError( - "The submodel_label array should be integer and contain all the numbers between 0 and the number of " - "partitions minus 1." - ) +from imod.typing.grid import get_non_spatial_dimension_names class ModelSplitter: diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index 976219fbd..0e2013a8a 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -25,6 +25,7 @@ from imod.common.statusinfo import NestedStatusInfo from imod.common.utilities.dataclass_type import DataclassType from imod.common.utilities.mask import mask_all_models +from imod.common.utilities.partitioninfo import create_partition_info from imod.common.utilities.regrid import _regrid_like from imod.common.utilities.version import ( get_version, @@ -43,10 +44,7 @@ from imod.mf6.multimodel.exchange_creator_unstructured import ( ExchangeCreator_Unstructured, ) -from imod.mf6.multimodel.modelsplitter import ( - ModelSplitter, - create_partition_info, -) +from imod.mf6.multimodel.modelsplitter import ModelSplitter from imod.mf6.out import open_cbc, open_conc, open_hds from imod.mf6.package import Package from imod.mf6.validation_settings import ValidationSettings diff --git a/imod/msw/coupler_mapping.py b/imod/msw/coupler_mapping.py index 635e4210e..53084e2ee 100644 --- a/imod/msw/coupler_mapping.py +++ b/imod/msw/coupler_mapping.py @@ -4,6 +4,7 @@ import pandas as pd import xarray as xr +from imod.common.interfaces.ipackagebase import IPackageBase from imod.mf6.dis import StructuredDiscretization from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.msw.fixed_format import VariableMetaData @@ -11,7 +12,7 @@ from imod.typing import IntArray -class CouplerMapping(MetaSwapPackage): +class CouplerMapping(MetaSwapPackage, IPackageBase): """ This contains the data to connect MODFLOW 6 cells to MetaSWAP svats. diff --git a/imod/msw/grid_data.py b/imod/msw/grid_data.py index fe6786f3b..89d1aacb7 100644 --- a/imod/msw/grid_data.py +++ b/imod/msw/grid_data.py @@ -83,7 +83,29 @@ def __init__( self._pkgcheck() - def generate_index_array(self) -> tuple[np.ndarray, xr.DataArray]: + def generate_index_array(self) -> np.ndarray: + """ + Generate index array to be used on other packages. + + Returns + ------- + np.ndarray + Index array and svat grid. + The index array is a 1D array with the index of the active cells. + """ + area = self.dataset["area"] + active = self.dataset["active"] + + isactive = area.where(active).notnull() + # Load into memory to avoid dask issue + # https://github.com/dask/dask/issues/11753 + isactive.load() + + index = isactive.values.ravel() + + return index + + def generate_index_svat_array(self) -> tuple[np.ndarray, xr.DataArray]: """ Generate index array and svat grid to be used on other packages. @@ -94,17 +116,18 @@ def generate_index_array(self) -> tuple[np.ndarray, xr.DataArray]: The index array is a 1D array with the index of the active cells. The svat grid is a 2D array with the SVAT numbers for each cell. """ + index = self.generate_index_array() + area = self.dataset["area"] active = self.dataset["active"] - isactive = area.where(active).notnull() + svat = xr.full_like(area, fill_value=0, dtype=np.int64).rename("svat") # Load into memory to avoid dask issue # https://github.com/dask/dask/issues/11753 isactive.load() svat.load() - index = isactive.values.ravel() svat.data[isactive.data] = np.arange(1, index.sum() + 1) return index, svat diff --git a/imod/msw/meteo_mapping.py b/imod/msw/meteo_mapping.py index ed4471b0e..28efa0a5e 100644 --- a/imod/msw/meteo_mapping.py +++ b/imod/msw/meteo_mapping.py @@ -9,6 +9,7 @@ import xarray as xr import imod +from imod.common.interfaces.ipackagebase import IPackageBase from imod.common.utilities.clip import clip_spatial_box, clip_time_slice from imod.common.utilities.dataclass_type import DataclassType from imod.msw.fixed_format import VariableMetaData @@ -65,7 +66,7 @@ def open_first_meteo_grid_from_imod5_data(imod5_data: Imod5DataDict, column_nr: return open_first_meteo_grid(metegrid_path, column_nr=column_nr) -class MeteoMapping(MetaSwapPackage): +class MeteoMapping(MetaSwapPackage, IPackageBase): """ This class provides common methods for creating mappings between meteorological data and MetaSWAP grids. It should not be instantiated diff --git a/imod/msw/model.py b/imod/msw/model.py index 31fb48953..608d31467 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -10,6 +10,8 @@ import numpy as np import xarray as xr +from imod.common.utilities.clip import clip_by_grid +from imod.common.utilities.partitioninfo import create_partition_info from imod.common.utilities.value_filters import enforce_scalar from imod.common.utilities.version import prepend_content_with_version_info from imod.mf6.dis import StructuredDiscretization @@ -43,23 +45,26 @@ from imod.msw.utilities.mask import MaskValues, mask_and_broadcast_cap_data from imod.msw.utilities.parse import read_para_sim from imod.msw.vegetation import AnnualCropFactors -from imod.typing import Imod5DataDict +from imod.typing import GridDataArray, Imod5DataDict from imod.util.dims import drop_layer_dim_cap_data from imod.util.regrid import RegridderWeightsCache +from imod.util.time import to_datetime_internal REQUIRED_PACKAGES = ( GridData, CouplerMapping, Infiltration, LanduseOptions, - MeteoGrid, EvapotranspirationMapping, PrecipitationMapping, IdfMapping, TimeOutputControl, AnnualCropFactors, ) - +METEO_PACKAGES = ( + MeteoGrid, + MeteoGridCopy, +) INITIAL_CONDITIONS_PACKAGES = ( InitialConditionsEquilibrium, InitialConditionsPercolation, @@ -113,6 +118,7 @@ class MetaSwapModel(Model): _pkg_id = "model" _file_name = "para_sim.inp" + _model_name = "MSW" _template = jinja2.Template( "{%for setting, value in settings.items()%}{{setting}} = {{value}}\n{%endfor%}" @@ -122,6 +128,7 @@ def __init__( self, unsaturated_database: Path | str, settings: Optional[dict[str, Any]] = None, + starttime: Optional[str] = None, ): super().__init__() @@ -130,6 +137,7 @@ def __init__( else: self.simulation_settings = settings + self.starttime = starttime self.simulation_settings["unsa_svat_path"] = unsaturated_database def _render_unsaturated_database_path(self, unsaturated_database: Union[str, Path]): @@ -152,6 +160,10 @@ def _check_required_packages(self) -> None: f"Missing the following required packages: {missing_packages}" ) + meteo_set = pkg_types_included & set(METEO_PACKAGES) + if len(meteo_set) < 1: + raise ValueError(f"Missing meteo package, assign one of {METEO_PACKAGES}") + initial_condition_set = pkg_types_included & set(INITIAL_CONDITIONS_PACKAGES) if len(initial_condition_set) < 1: raise ValueError( @@ -205,6 +217,19 @@ def _has_file_copier(self) -> bool: return FileCopier in pkg_types_included def _get_starttime(self): + if self.starttime is not None: + starttime = to_datetime_internal(self.starttime, use_cftime=False) + else: + starttime = self._get_minimum_package_time() + + year, time_since_start_year = to_metaswap_timeformat([starttime]) + + year = int(enforce_scalar(year.values)) + time_since_start_year = float(enforce_scalar(time_since_start_year.values)) + + return year, time_since_start_year + + def _get_minimum_package_time(self): """ Loop over all packages to get the minimum time. @@ -218,14 +243,13 @@ def _get_starttime(self): if "time" in ds.coords: starttimes.append(ds["time"].min().values) - starttime = min(starttimes) - - year, time_since_start_year = to_metaswap_timeformat([starttime]) - - year = int(enforce_scalar(year.values)) - time_since_start_year = float(enforce_scalar(time_since_start_year.values)) + if len(starttimes) == 0: + raise ValueError( + "No package with a coordinate 'time' found. Please add a MeteoGrid or TimeOutputControl package." + ) - return year, time_since_start_year + starttime = min(starttimes) + return starttime def get_pkgkey( self, pkg_type: type[MetaSwapPackage], optional_package: bool = False @@ -268,7 +292,7 @@ def _get_pkg_key( ) return self.get_pkgkey(pkg_type, optional_package) - def _model_checks(self, validate: bool): + def _model_checks(self, validate: Optional[bool] = True): if validate and not self._has_file_copier(): self._check_required_packages() self._check_vegetation_indices_in_annual_crop_factors() @@ -311,8 +335,8 @@ def write( self, directory: Union[str, Path], mf6_dis: StructuredDiscretization, - mf6_wel: Mf6Wel, - validate: bool = True, + mf6_wel: Optional[Mf6Wel] = None, + validate: Optional[bool] = True, ): """ Write packages and simulation settings (para_sim.inp). @@ -336,7 +360,7 @@ def write( # Get index and svat grid_key = self.get_pkgkey(GridData) grid_pkg = cast(GridData, self[grid_key]) - index, svat = grid_pkg.generate_index_array() + index, svat = grid_pkg.generate_index_svat_array() # write package contents for pkgname in self: @@ -477,6 +501,103 @@ def clip_box( ) return clipped + def split( + self, + submodel_labels: GridDataArray, + ) -> dict[str, "MetaSwapModel"]: + """ + Split a MetaSWAP model in different partitions using a submodel_labels + array. Note: for specifying meteorological grid data, splitting is only + supported when the MeteoGridCopy instance is being used. + + Parameters + ---------- + submodel_labels: xr.DataArray or xu.UgridDataArray + A grid that defines how the simulation will be split. The array + should have the same topology as the domain being split, i.e. + similar shape as a layer in the domain. The values in the array + indicate to which partition a cell belongs. The values should be + zero or greater. + + Returns + ------- + dict[str, MetaSwapModel] + A mapping from generated submodel names to the corresponding clipped + MetaSWAP submodel. + Examples + -------- + >>> submodel_labels = mf6_sim.create_partition_labels(n_partitions=4) + >>> msw_splitted = msw.split(submodel_labels) + """ + + has_meteogrid = MeteoGrid in [type(pkg) for pkg in self.values()] + if has_meteogrid: + raise ValueError( + "Splitting packages of type MeteoGrid is not supported, use MeteoGridCopy instead." + ) + + has_meteogrid_copy = MeteoGridCopy in [type(pkg) for pkg in self.values()] + + settings = deepcopy(self.simulation_settings) + starttime = self.starttime + unsa_svat_path = settings.pop("unsa_svat_path") + + partition_info_list = create_partition_info(submodel_labels) + + # Initialize mapping from partition IDs to models + partition_id_to_models: dict[int, dict[str, MetaSwapModel]] = {} + for submodel_partition_info in partition_info_list: + partition_id_to_models[submodel_partition_info.id] = {} + partitioned_submodels = {} + submodel_to_partition = {} + + # Create empty MetaSWAP model for each partition + for submodel_partition_info in partition_info_list: + submodel_name = f"{self._model_name}_{submodel_partition_info.id}" + + submodel = MetaSwapModel(unsa_svat_path, settings, starttime) + partitioned_submodels[submodel_name] = submodel + submodel_to_partition[submodel_name] = submodel_partition_info + partition_id_to_models[submodel_partition_info.id][submodel_name] = submodel + + # First, handle the grid data and determine the overlap. + grid_key = self.get_pkgkey(GridData) + grid_pkg = cast(GridData, self[grid_key]) + has_overlap = {} + + for submodel_name, submodel in partitioned_submodels.items(): + partition_info = submodel_to_partition[submodel_name] + sliced_grid_pkg = clip_by_grid(grid_pkg, partition_info.active_domain) + sliced_index = sliced_grid_pkg.generate_index_array() + + # Add package to model if it has data in the overlap. + if bool(sliced_index.any()): + has_overlap[submodel_name] = True + submodel[grid_key] = sliced_grid_pkg + else: + has_overlap[submodel_name] = False + + # Second, add other packages to each partitioned submodel. + for submodel_name, submodel in partitioned_submodels.items(): + partition_info = submodel_to_partition[submodel_name] + + if not has_overlap[submodel_name]: + continue + + # Add packages to models + for pkg_name, pkg in self.items(): + if isinstance(pkg, GridData): + continue + + if has_meteogrid_copy and isinstance(pkg, MeteoMapping): + submodel[pkg_name] = deepcopy(pkg) + else: + # Slice and add the package to the partitioned model + sliced_package = clip_by_grid(pkg, partition_info.active_domain) + submodel[pkg_name] = sliced_package + + return partitioned_submodels + @classmethod def from_imod5_data( cls, diff --git a/imod/tests/conftest.py b/imod/tests/conftest.py index cdc3eb56c..0670d85d0 100644 --- a/imod/tests/conftest.py +++ b/imod/tests/conftest.py @@ -101,4 +101,10 @@ from .fixtures.msw_fixture import fixed_format_parser, simple_2d_grid_with_subunits from .fixtures.msw_imod5_cap_fixture import imod5_cap_data from .fixtures.msw_meteo_fixture import meteo_grids -from .fixtures.msw_model_fixture import coupled_mf6_model, coupled_mf6wel, msw_model +from .fixtures.msw_model_fixture import ( + coupled_mf6_model, + coupled_mf6wel, + msw_model, + msw_model_get_starttime, + msw_model_no_sprinkling, +) diff --git a/imod/tests/fixtures/msw_model_fixture.py b/imod/tests/fixtures/msw_model_fixture.py index 33b9e8723..06f7db8b7 100644 --- a/imod/tests/fixtures/msw_model_fixture.py +++ b/imod/tests/fixtures/msw_model_fixture.py @@ -201,13 +201,6 @@ def make_msw_model(): msw_model["meteo_grid"] = msw.MeteoGrid(precipitation, evapotranspiration) msw_model["mapping_prec"] = msw.PrecipitationMapping(precipitation) msw_model["mapping_evt"] = msw.EvapotranspirationMapping(precipitation * 1.5) - - # %% Sprinkling - msw_model["sprinkling"] = msw.Sprinkling( - max_abstraction_groundwater=xr.full_like(area, 100.0), - max_abstraction_surfacewater=xr.full_like(area, 100.0), - ) - # %% Ponding msw_model["ponding"] = msw.Ponding( ponding_depth=xr.full_like(area, 0.0), @@ -279,6 +272,17 @@ def make_msw_model(): return msw_model +def msw_add_sprinkling(msw_model): + # %% Sprinkling + area = msw_model["grid"].dataset["area"] + + msw_model["sprinkling"] = msw.Sprinkling( + max_abstraction_groundwater=xr.full_like(area, 100.0), + max_abstraction_surfacewater=xr.full_like(area, 100.0), + ) + return msw_model + + @pytest.fixture(scope="function") def coupled_mf6wel(): x = [1.0, 2.0, 3.0] @@ -316,4 +320,15 @@ def coupled_mf6_model(): @pytest.fixture(scope="function") def msw_model(): + msw_model = make_msw_model() + return msw_add_sprinkling(msw_model) + + +@pytest.fixture(scope="function") +def msw_model_no_sprinkling(): return make_msw_model() + + +@pytest.fixture(scope="function") +def msw_model_get_starttime(): + return "1/1/1971" diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py index 945fe1871..2fc5f42c9 100644 --- a/imod/tests/test_mf6/test_mf6_simulation.py +++ b/imod/tests/test_mf6/test_mf6_simulation.py @@ -21,11 +21,11 @@ import imod from imod.common.statusinfo import NestedStatusInfo, StatusInfo +from imod.common.utilities.partitioninfo import PartitionInfo from imod.common.utilities.version import get_version from imod.logging import LoggerType, LogLevel from imod.mf6 import LayeredWell, Well from imod.mf6.model import Modflow6Model -from imod.mf6.multimodel.modelsplitter import PartitionInfo from imod.mf6.oc import OutputControl from imod.mf6.regrid.regrid_schemes import ( DiscretizationRegridMethod, 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 c78f855e2..11d9b0b3e 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 @@ -4,8 +4,8 @@ import pytest from pytest_cases import parametrize_with_cases +from imod.common.utilities.partitioninfo import create_partition_info from imod.mf6.multimodel.exchange_creator_structured import ExchangeCreator_Structured -from imod.mf6.multimodel.modelsplitter import create_partition_info from imod.tests.fixtures.flow_basic_fixture import BasicDisSettings from imod.typing.grid import zeros_like from imod.util.spatial import spatial_reference 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 23bd0102a..aee3e5049 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 @@ -3,10 +3,10 @@ import xugrid as xu import imod +from imod.common.utilities.partitioninfo import create_partition_info from imod.mf6.multimodel.exchange_creator_unstructured import ( ExchangeCreator_Unstructured, ) -from imod.mf6.multimodel.modelsplitter import create_partition_info from imod.tests.fixtures.flow_basic_fixture import BasicDisSettings from imod.typing.grid import zeros_like 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 942550755..fcccecccf 100644 --- a/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py +++ b/imod/tests/test_mf6/test_multimodel/test_mf6_modelsplitter.py @@ -1,6 +1,7 @@ import numpy as np -from imod.mf6.multimodel.modelsplitter import ModelSplitter, create_partition_info +from imod.common.utilities.partitioninfo import create_partition_info +from imod.mf6.multimodel.modelsplitter import ModelSplitter from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run from imod.typing.grid import zeros_like 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 0fecb995e..c834c5c81 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 @@ -4,9 +4,10 @@ import numpy as np import pytest +from imod.common.utilities.partitioninfo import create_partition_info from imod.mf6 import AdvectionCentral, AdvectionTVD, AdvectionUpstream from imod.mf6.dsp import Dispersion -from imod.mf6.multimodel.modelsplitter import ModelSplitter, create_partition_info +from imod.mf6.multimodel.modelsplitter import ModelSplitter from imod.mf6.package import Package from imod.mf6.simulation import Modflow6Simulation from imod.tests.fixtures.mf6_modelrun_fixture import assert_simulation_can_run diff --git a/imod/tests/test_msw/test_grid_data.py b/imod/tests/test_msw/test_grid_data.py index cd2392efa..d50490b70 100644 --- a/imod/tests/test_msw/test_grid_data.py +++ b/imod/tests/test_msw/test_grid_data.py @@ -64,7 +64,7 @@ def test_write( xr.full_like(like, True, dtype=bool), ) - index, svat = grid_data.generate_index_array() + index, svat = grid_data.generate_index_svat_array() with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) @@ -293,12 +293,12 @@ def case_grid_data_two_subunits__dask( @parametrize_with_cases("grid_data_dict", cases=".", has_tag="two_subunit") -def test_generate_index_array( +def test_generate_index_svat_array( grid_data_dict: dict[str, xr.DataArray], coords_two_subunit: dict ): grid_data = GridData(**grid_data_dict) - index, svat = grid_data.generate_index_array() + index, svat = grid_data.generate_index_svat_array() index_expected = [ False, @@ -346,7 +346,7 @@ def test_generate_index_array( def test_simple_model(fixed_format_parser, grid_data_dict: dict[str, xr.DataArray]): grid_data = GridData(**grid_data_dict) - index, svat = grid_data.generate_index_array() + index, svat = grid_data.generate_index_svat_array() with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) @@ -388,7 +388,7 @@ def test_simple_model_1_subunit( ): grid_data = GridData(**grid_data_dict) - index, svat = grid_data.generate_index_array() + index, svat = grid_data.generate_index_svat_array() with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index abd74786e..80f136466 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -8,10 +8,12 @@ from imod import mf6, msw from imod.msw.copy_files import FileCopier +from imod.msw.meteo_grid import MeteoGridCopy from imod.msw.meteo_mapping import MeteoMapping from imod.msw.model import DEFAULT_SETTINGS from imod.msw.utilities.parse import read_para_sim from imod.typing import GridDataArray, Imod5DataDict +from imod.typing.grid import zeros_like from imod.util.regrid import RegridderWeightsCache @@ -143,6 +145,49 @@ def test_clip_box(msw_model): assert clipped_settings == model_settings +def test_msw_model_split_write( + msw_model_no_sprinkling, msw_model_get_starttime, coupled_mf6_model, tmp_path +): + # Note that this test is without sprinkling wells. + + # Set the MetaSWAP output directory. + output_dir = tmp_path / "metaswap" + output_dir.mkdir(exist_ok=True, parents=True) + + # Set the starting time. This should be explicitly set since the MetaSWAP submodels + # will not have further time information during writing. + if msw_model_no_sprinkling.starttime is None: + msw_model_no_sprinkling.starttime = msw_model_get_starttime + + # Create meteo output directory, write the meteo data, and copy metegrid. + meteo_output_dir = tmp_path / "meteo" + meteo_output_dir.mkdir(exist_ok=True, parents=True) + msw_model_no_sprinkling["meteo_grid"].write(meteo_output_dir) + del msw_model_no_sprinkling["meteo_grid"] + mete_grid = meteo_output_dir / Path("mete_grid.inp") + msw_model_no_sprinkling["meteo_grid"] = MeteoGridCopy(mete_grid) + + # Set the partition mask. + active = coupled_mf6_model["GWF_1"].domain.sel(layer=1) + submodel_labels = zeros_like(active) + submodel_labels = submodel_labels.where((submodel_labels["y"] > 2.5), 1) + submodel_labels = submodel_labels.where((submodel_labels["y"] > 1.5), 2) + + # Split the MODFLOW and MetaSWAP models. + mf6_splitted = coupled_mf6_model.split(submodel_labels) + msw_splitted = msw_model_no_sprinkling.split(submodel_labels) + + # Write the MetaSWAP submodels. + for msw_submodelname, msw_submodel in msw_splitted.items(): + mf6_submodelname = f"GWF_1{msw_submodelname.split('MSW')[1]}" + mf6_dis = mf6_splitted[mf6_submodelname]["dis"] + submodel_output_dir = output_dir / msw_submodelname + msw_submodel.write(directory=submodel_output_dir, mf6_dis=mf6_dis) + + # Assert by checking the number of MetaSWAP inp file counts. + assert len(list(submodel_output_dir.rglob(r"*.inp"))) == 15 + + def test_render_unsat_database_path(msw_model, tmp_path): rel_path = msw_model._render_unsaturated_database_path("./unsat_database")