From b894f7aea875be1be5a1369a23d264c0b33e1215 Mon Sep 17 00:00:00 2001 From: rleander Date: Mon, 16 Mar 2026 18:07:48 +0100 Subject: [PATCH 1/9] the dump model method of ModflowModel now moved to a common function in utilities, accepting the IModel interface as argument. The MetaSWAP model now also implements this interface, so can be prepared to use the same functionality --- imod/common/utilities/dump_model.py | 82 +++++++++++++++++++++++++++ imod/mf6/model.py | 68 ---------------------- imod/mf6/simulation.py | 5 +- imod/msw/model.py | 3 +- imod/tests/test_mf6/test_mf6_model.py | 3 +- 5 files changed, 89 insertions(+), 72 deletions(-) create mode 100644 imod/common/utilities/dump_model.py diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py new file mode 100644 index 000000000..acd619f4c --- /dev/null +++ b/imod/common/utilities/dump_model.py @@ -0,0 +1,82 @@ +import collections +from pathlib import Path +from typing import Any, Optional +import tomli_w +from imod.logging.logging_decorators import standard_log_decorator +from imod.common.serializer import EngineType, create_package_serializer +from imod.mf6.validation_settings import ValidationSettings +from imod.schemata import ValidationError +from imod.common.interfaces.imodel import IModel + +def testfun(): + print("test") + +@standard_log_decorator() +def dump_modelpkgs( + model: IModel, + directory, + modelname, + validate: bool = True, + mdal_compliant: bool = False, + crs: Optional[Any] = None, + engine: EngineType = "netcdf4", +) -> Path: + """ + Dump model to files. Writes a model definition as .TOML file, which + points to data for each package. Each package is stored as a separate + NetCDF. Structured grids are saved as regular NetCDFs, unstructured + grids are saved as UGRID NetCDF. Structured grids are always made GDAL + compliant, unstructured grids can be made MDAL compliant optionally. + + Parameters + ---------- + directory: str or Path + directory to dump simulation into. + modelname: str + modelname, will be used to create a subdirectory. + validate: bool, optional + Whether to validate simulation data. Defaults to True. + mdal_compliant: bool, optional + Convert data with + :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL + compliant unstructured grids. Defaults to False. + crs: Any, optional + Anything accepted by rasterio.crs.CRS.from_user_input + Requires ``rioxarray`` installed. + engine : str, optional + File engine used to write packages. Options are ``'netcdf4'``, + ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other + softwares, for example QGIS. Zarr is optimized for big data, cloud + storage and parallel access. The ``'zarr.zip'`` option is an + experimental option which creates a zipped zarr store in a single + file, which is easier to copy and automatically compresses data as + well. Default is ``'netcdf4'``. + + """ + modeldirectory = Path(directory) / modelname + modeldirectory.mkdir(exist_ok=True, parents=True) + + # validation currently only supports MF6, but we want to keep the option to turn it on for other + validation_context = ValidationSettings(validate=validate) + if validation_context.validate: + statusinfo = model.validate(modelname, validation_context) + if statusinfo.has_errors(): + raise ValidationError(statusinfo.to_string()) + + toml_content: dict = collections.defaultdict(dict) + + for pkgname, pkg in model.items(): + pkg_path = pkg.to_file( + modeldirectory, + pkgname, + mdal_compliant=mdal_compliant, + crs=crs, + engine=engine, + ) + toml_content[type(pkg).__name__][pkgname] = pkg_path.name + + toml_path = modeldirectory / f"{modelname}.toml" + with open(toml_path, "wb") as f: + tomli_w.dump(toml_content, f) + + return toml_path \ No newline at end of file diff --git a/imod/mf6/model.py b/imod/mf6/model.py index b1a9b6f37..efb2b7bf1 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -606,74 +606,6 @@ def _write( return NestedStatusInfo(modelname) - @standard_log_decorator() - def dump( - self, - directory, - modelname, - validate: bool = True, - mdal_compliant: bool = False, - crs: Optional[Any] = None, - engine: EngineType = "netcdf4", - ) -> Path: - """ - Dump simulation to files. Writes a model definition as .TOML file, which - points to data for each package. Each package is stored as a separate - NetCDF. Structured grids are saved as regular NetCDFs, unstructured - grids are saved as UGRID NetCDF. Structured grids are always made GDAL - compliant, unstructured grids can be made MDAL compliant optionally. - - Parameters - ---------- - directory: str or Path - directory to dump simulation into. - modelname: str - modelname, will be used to create a subdirectory. - validate: bool, optional - Whether to validate simulation data. Defaults to True. - mdal_compliant: bool, optional - Convert data with - :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL - compliant unstructured grids. Defaults to False. - crs: Any, optional - Anything accepted by rasterio.crs.CRS.from_user_input - Requires ``rioxarray`` installed. - engine : str, optional - File engine used to write packages. Options are ``'netcdf4'``, - ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other - softwares, for example QGIS. Zarr is optimized for big data, cloud - storage and parallel access. The ``'zarr.zip'`` option is an - experimental option which creates a zipped zarr store in a single - file, which is easier to copy and automatically compresses data as - well. Default is ``'netcdf4'``. - - """ - modeldirectory = pathlib.Path(directory) / modelname - modeldirectory.mkdir(exist_ok=True, parents=True) - validation_context = ValidationSettings(validate=validate) - if validation_context.validate: - statusinfo = self.validate(modelname, validation_context) - if statusinfo.has_errors(): - raise ValidationError(statusinfo.to_string()) - - toml_content: dict = collections.defaultdict(dict) - - for pkgname, pkg in self.items(): - pkg_path = pkg.to_file( - modeldirectory, - pkgname, - mdal_compliant=mdal_compliant, - crs=crs, - engine=engine, - ) - toml_content[type(pkg).__name__][pkgname] = pkg_path.name - - toml_path = modeldirectory / f"{modelname}.toml" - with open(toml_path, "wb") as f: - tomli_w.dump(toml_content, f) - - return toml_path - @classmethod def from_file(cls, toml_path): pkg_classes = { diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index 976219fbd..c930f9db7 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.dump_model import dump_modelpkgs from imod.common.utilities.regrid import _regrid_like from imod.common.utilities.version import ( get_version, @@ -1039,8 +1040,8 @@ def dump( for key, value in self.items(): cls_name = type(value).__name__ if isinstance(value, Modflow6Model): - model_toml_path = value.dump( - directory, key, validate, mdal_compliant, crs, engine=engine + model_toml_path = dump_modelpkgs( + value, directory, key, validate, mdal_compliant, crs, engine=engine ) toml_content[cls_name][key] = model_toml_path.relative_to( directory diff --git a/imod/msw/model.py b/imod/msw/model.py index 31fb48953..4a15c91bb 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -46,6 +46,7 @@ from imod.typing import Imod5DataDict from imod.util.dims import drop_layer_dim_cap_data from imod.util.regrid import RegridderWeightsCache +from imod.common.interfaces.imodel import IModel REQUIRED_PACKAGES = ( GridData, @@ -90,7 +91,7 @@ } -class Model(collections.UserDict): +class Model(collections.UserDict,IModel): def __setitem__(self, key, value): # TODO: Add packagecheck super().__setitem__(key, value) diff --git a/imod/tests/test_mf6/test_mf6_model.py b/imod/tests/test_mf6/test_mf6_model.py index 7495cbe51..068730e4f 100644 --- a/imod/tests/test_mf6/test_mf6_model.py +++ b/imod/tests/test_mf6/test_mf6_model.py @@ -21,6 +21,7 @@ from imod.mf6.write_context import WriteContext from imod.schemata import ValidationError from imod.typing.grid import concat, nan_like +from imod.common.utilities.dump_model import dump_modelpkgs # Duplicate from test_mf6_dis.py @@ -91,7 +92,7 @@ def test_key_assign(): def roundtrip(model, tmp_path): - model.dump(tmp_path, "test") + dump_modelpkgs(tmp_path, "test") back = type(model).from_file(tmp_path / "test/test.toml") assert isinstance(back, type(model)) From d6eb9a11dcd4534ed20bbe0fb70c7446d5a456b6 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Mon, 4 May 2026 11:10:59 +0200 Subject: [PATCH 2/9] try out imodel interface --- imod/common/utilities/dump_model.py | 48 ++++++++++++++++++++++----- imod/mf6/model.py | 4 +-- imod/mf6/simulation.py | 2 +- imod/msw/model.py | 5 +-- imod/tests/test_mf6/test_mf6_model.py | 2 +- imod/tests/test_msw/test_model.py | 14 ++++++++ 6 files changed, 60 insertions(+), 15 deletions(-) diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py index acd619f4c..791de1ff1 100644 --- a/imod/common/utilities/dump_model.py +++ b/imod/common/utilities/dump_model.py @@ -1,16 +1,24 @@ import collections -from pathlib import Path +import inspect +import pathlib +from pathlib import Path from typing import Any, Optional + +import tomli import tomli_w + +import imod.mf6 +import imod.msw +from imod.common.interfaces.imodel import IModel +from imod.common.interfaces.ipackage import IPackage +from imod.common.serializer import EngineType from imod.logging.logging_decorators import standard_log_decorator -from imod.common.serializer import EngineType, create_package_serializer +from imod.mf6.model import Modflow6Model from imod.mf6.validation_settings import ValidationSettings +from imod.msw.model import MetaSwapModel from imod.schemata import ValidationError -from imod.common.interfaces.imodel import IModel -def testfun(): - print("test") - + @standard_log_decorator() def dump_modelpkgs( model: IModel, @@ -55,7 +63,7 @@ def dump_modelpkgs( """ modeldirectory = Path(directory) / modelname modeldirectory.mkdir(exist_ok=True, parents=True) - + # validation currently only supports MF6, but we want to keep the option to turn it on for other validation_context = ValidationSettings(validate=validate) if validation_context.validate: @@ -79,4 +87,28 @@ def dump_modelpkgs( with open(toml_path, "wb") as f: tomli_w.dump(toml_content, f) - return toml_path \ No newline at end of file + return toml_path + + +def from_file(instance, toml_path): + if isinstance(instance, Modflow6Model): + modref = imod.mf6 + if isinstance(instance, MetaSwapModel): + modref = imod.msw + pkg_classes = { + name: pkg_cls + for name, pkg_cls in inspect.getmembers(modref, inspect.isclass) + if issubclass(pkg_cls, IPackage) + } + + toml_path = pathlib.Path(toml_path) + with open(toml_path, "rb") as f: + toml_content = tomli.load(f) + + parentdir = toml_path.parent + for key, entry in toml_content.items(): + for pkgname, path in entry.items(): + pkg_cls = pkg_classes[key] + instance[pkgname] = pkg_cls.from_file(parentdir / path) + + return instance diff --git a/imod/mf6/model.py b/imod/mf6/model.py index efb2b7bf1..8da50d49e 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -6,20 +6,18 @@ import pathlib import warnings from pathlib import Path -from typing import Any, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import cftime import jinja2 import numpy as np import tomli -import tomli_w import xarray as xr import xugrid as xu from jinja2 import Template import imod from imod.common.interfaces.imodel import IModel -from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase from imod.common.utilities.clip import clip_box_dataset from imod.common.utilities.mask import mask_all_packages diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index c930f9db7..3d54a6a53 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -24,8 +24,8 @@ from imod.common.serializer import EngineType 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.dump_model import dump_modelpkgs +from imod.common.utilities.mask import mask_all_models from imod.common.utilities.regrid import _regrid_like from imod.common.utilities.version import ( get_version, diff --git a/imod/msw/model.py b/imod/msw/model.py index 4a15c91bb..56b0bf140 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -12,6 +12,7 @@ from imod.common.utilities.value_filters import enforce_scalar from imod.common.utilities.version import prepend_content_with_version_info +from imod.logging import standard_log_decorator from imod.mf6.dis import StructuredDiscretization from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.msw.copy_files import FileCopier @@ -91,7 +92,7 @@ } -class Model(collections.UserDict,IModel): +class Model(collections.UserDict): def __setitem__(self, key, value): # TODO: Add packagecheck super().__setitem__(key, value) @@ -101,7 +102,7 @@ def update(self, *args, **kwargs): self[k] = v -class MetaSwapModel(Model): +class MetaSwapModel(Model, IModel): """ Contains data and writes consistent model input files diff --git a/imod/tests/test_mf6/test_mf6_model.py b/imod/tests/test_mf6/test_mf6_model.py index 068730e4f..c73587743 100644 --- a/imod/tests/test_mf6/test_mf6_model.py +++ b/imod/tests/test_mf6/test_mf6_model.py @@ -11,6 +11,7 @@ from xugrid.core.wrap import UgridDataArray import imod +from imod.common.utilities.dump_model import dump_modelpkgs from imod.mf6 import ConstantConcentration, ConstantHead from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.model import Modflow6Model @@ -21,7 +22,6 @@ from imod.mf6.write_context import WriteContext from imod.schemata import ValidationError from imod.typing.grid import concat, nan_like -from imod.common.utilities.dump_model import dump_modelpkgs # Duplicate from test_mf6_dis.py diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index abd74786e..66f24b79b 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -7,14 +7,28 @@ from pytest_cases import parametrize_with_cases from imod import mf6, msw +from imod.common.utilities.dump_model import dump_modelpkgs from imod.msw.copy_files import FileCopier from imod.msw.meteo_mapping import MeteoMapping from imod.msw.model import DEFAULT_SETTINGS +from imod.msw.model import MetaSwapModel from imod.msw.utilities.parse import read_para_sim from imod.typing import GridDataArray, Imod5DataDict from imod.util.regrid import RegridderWeightsCache +def roundtrip(msw_model, tmpdir_factory, name, engine): + # TODO: look at the values? + tmp_path = tmpdir_factory.mktemp(name) + dump_modelpkgs(msw_model, tmp_path, name, engine=engine) + back = MetaSwapModel.from_file(tmp_path / f"{simulation.name}.toml") + assert isinstance(back, MetaSwapModel) + +def test_msw_pkgdump(msw_model, tmp_path): + dump_modelpkgs(msw_model, tmp_path, "testmodel", engine="netcdf4") +# back = MetaSwapModel.from_file(tmp_path / "testmodel.toml") +# assert isinstance(back, MetaSwapModel) + def test_msw_model_write(msw_model, coupled_mf6_model, coupled_mf6wel, tmp_path): mf6_dis = coupled_mf6_model["GWF_1"]["dis"] From 54d1301cf00ce65a2c10614d37d4c490c454c239 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Mon, 11 May 2026 17:15:34 +0200 Subject: [PATCH 3/9] roundtrip test dumping and loading msw packages added --- imod/common/utilities/dump_model.py | 3 + imod/msw/model.py | 184 +++++++++++++++++++++++++- imod/msw/pkgbase.py | 119 ++++++++++++++++- imod/tests/test_mf6/test_mf6_model.py | 2 +- imod/tests/test_msw/test_model.py | 15 +-- 5 files changed, 309 insertions(+), 14 deletions(-) diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py index 791de1ff1..76ffe6d09 100644 --- a/imod/common/utilities/dump_model.py +++ b/imod/common/utilities/dump_model.py @@ -83,6 +83,9 @@ def dump_modelpkgs( ) toml_content[type(pkg).__name__][pkgname] = pkg_path.name + if hasattr(model, "simulation_settings"): + toml_content["simulation_settings"] = model.simulation_settings + toml_path = modeldirectory / f"{modelname}.toml" with open(toml_path, "wb") as f: tomli_w.dump(toml_content, f) diff --git a/imod/msw/model.py b/imod/msw/model.py index 56b0bf140..30d00c1cb 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -1,18 +1,22 @@ import collections +import inspect import warnings from copy import copy, deepcopy from datetime import datetime from pathlib import Path -from typing import Any, Optional, Union, cast +from typing import Any, Optional, Tuple, Union, cast import cftime import jinja2 import numpy as np +import tomli import xarray as xr +import imod.msw +from imod.common.interfaces.imodel import IModel +from imod.common.utilities.mask import mask_all_packages from imod.common.utilities.value_filters import enforce_scalar from imod.common.utilities.version import prepend_content_with_version_info -from imod.logging import standard_log_decorator from imod.mf6.dis import StructuredDiscretization from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.msw.copy_files import FileCopier @@ -44,10 +48,9 @@ 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.common.interfaces.imodel import IModel REQUIRED_PACKAGES = ( GridData, @@ -344,6 +347,34 @@ def write( for pkgname in self: self[pkgname].write(directory, index, svat, mf6_dis, mf6_wel) + @classmethod + def from_file(cls, directory, modelname): + pkg_classes = { + name: pkg_cls + for name, pkg_cls in inspect.getmembers(imod.msw, inspect.isclass) + if issubclass(pkg_cls, MetaSwapPackage) + } + modeldirectory = Path(directory) / modelname + toml_path = modeldirectory / f"{modelname}.toml" + with open(toml_path, "rb") as f: + toml_content = tomli.load(f) + + parentdir = toml_path.parent + if issubclass(cls, MetaSwapModel): + simulation_settings = toml_content.get("simulation_settings", {}) + unsa_svat_path = simulation_settings.get("unsa_svat_path", "") + instance = cls(unsa_svat_path, simulation_settings) + else: + instance = cls() + + for key, entry in toml_content.items(): + if key != "simulation_settings": + for pkgname, path in entry.items(): + pkg_cls = pkg_classes[key] + instance[pkgname] = pkg_cls.from_file(parentdir / path) + + return instance + def regrid_like( self, mf6_regridded_dis: StructuredDiscretization, @@ -541,3 +572,148 @@ def from_imod5_data( model["time_oc"] = TimeOutputControl(times_da) return model + + def _is_splitting_supported(self) -> Tuple[bool, str]: + """ + Returns True if all the packages in the model supports splitting. If one + of the packages in the model does not support splitting, it returns the + name of the first one. + + Returns + ------- + Tuple[bool, str] + A tuple where the first element is a boolean indicating if splitting + is supported, and the second element is the name of the first package + that does not support splitting, or an empty string if all packages + support splitting. + """ + for package_name, package in self.items(): + if not package._is_splitting_supported(): + return False, package_name + return True, "" + + def _is_regridding_supported(self) -> Tuple[bool, str]: + """ + Returns True if all the packages in the model supports regridding. If one + of the packages in the model does not support regridding, it returns the + name of the first one. + + Returns + ------- + Tuple[bool, str] + A tuple where the first element is a boolean indicating if regridding + is supported, and the second element is the name of the first package + that does not support regridding, or an empty string if all packages + support regridding. + """ + for package_name, package in self.items(): + if not package._is_regridding_supported(): + return False, package_name + return True, "" + + def _is_clipping_supported(self) -> Tuple[bool, str]: + """ + Returns True if all the packages in the model supports clipping. If one + of the packages in the model does not support clipping, it returns the + name of the first one. + + Returns + ------- + Tuple[bool, str] + A tuple where the first element is a boolean indicating if clipping + is supported, and the second element is the name of the first package + that does not support clipping, or an empty string if all packages + support clipping. + """ + for package_name, package in self.items(): + if not package._is_clipping_supported(): + return False, package_name + return True, "" + + def purge_empty_packages( + self, model_name: Optional[str] = "", ignore_time: bool = False + ) -> None: + """ + This method removes empty packages from the model in place. + + Parameters + ---------- + model_name: str, optional + Name of the model, used for logging. + ignore_time: bool, optional + If False, packages are considered empty if they have no data at all + timesteps. If True, packages are considered empty if they have no + data at the first time step. The latter can increase performance + considerably. + """ + empty_packages = [ + package_name + for package_name, package in self.items() + if package.is_empty(ignore_time=ignore_time) + ] + for package_name in empty_packages: + self.pop(package_name) + + def mask_all_packages( + self, + mask: GridDataArray, + ignore_time_purge_empty: bool = False, + ): + """ + This function applies a mask to all packages in a model. The mask must + be presented as an idomain-like integer array that has 0 (inactive) or + <0 (vertical passthrough) values in filtered cells and >0 in active + cells. + Masking will overwrite idomain with the mask where the mask is <=0. + Where the mask is >0, the original value of idomain will be kept. Masking + will update the packages accordingly, blanking their input where needed, + and is therefore not a reversible operation. + + Parameters + ---------- + mask: xr.DataArray, xu.UgridDataArray of ints + idomain-like integer array. >0 sets cells to active, 0 sets cells to inactive, + <0 sets cells to vertical passthrough + ignore_time_purge_empty: bool, default False + Whether to ignore time dimension when purging empty packages. Can + improve performance when masking models with many time steps. + """ + + mask_all_packages(self, mask, ignore_time_purge_empty) + + @property + def model_id(self) -> str: + if self._model_id is None: + return "MetaSwapModel" + return self._model_id + + @property + def options(self) -> dict: + return self._options + + @property + def domain(self) -> GridDataArray: + """ + Get the active domain array as an integer array compatible with idomain. + + Returns + ------- + GridDataArray + Integer array where: + - 1 indicates active cells + - 0 indicates inactive cells + """ + grid_key = self.get_pkgkey(GridData) + active = self[grid_key]["active"] + return active.astype(int) + + def validate( + self, + model_name: str = "", + validation_context: Optional[Any] = None, + **kwargs, + ) -> None: + return + + +# make a read function for packages from netcdf diff --git a/imod/msw/pkgbase.py b/imod/msw/pkgbase.py index 442644fc2..6d7fe3eee 100644 --- a/imod/msw/pkgbase.py +++ b/imod/msw/pkgbase.py @@ -1,13 +1,17 @@ import abc +import numbers from copy import deepcopy from pathlib import Path -from typing import Any, Optional, TextIO, TypeAlias, Union +from typing import Any, Optional, Self, TextIO, TypeAlias, Union import cftime import numpy as np import pandas as pd import xarray as xr +import zarr # nb: in modflow6, this is a try...except +from xarray.core.utils import is_scalar +from imod.common.serializer import EngineType, create_package_serializer from imod.common.utilities.clip import clip_spatial_box, clip_time_slice from imod.common.utilities.dataclass_type import DataclassType, EmptyRegridMethod from imod.common.utilities.regrid import ( @@ -23,6 +27,17 @@ DataDictType: TypeAlias = dict[str, IntArray | int | str] +def _is_scalar_nan(da: GridDataArray): + """ + Test if is_scalar_nan, carefully avoid loading grids in memory + """ + scalar_data: bool = is_scalar(da) + if scalar_data: + stripped_value = da.to_numpy()[()] + return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value) # type: ignore[call-overload] + return False + + class MetaSwapPackage(abc.ABC): """ MetaSwapPackage is used to share methods for Metaswap packages. @@ -385,3 +400,105 @@ def from_imod5_data(self, *args, **kwargs): def _is_clipping_supported(self) -> bool: return True + + @classmethod + def from_file(cls, path: str | Path, **kwargs) -> Self: + """ + Loads an imod msw package from a file (netcdf or zarr). + Note that the checks upon package initialization are not done again! + + Parameters + ---------- + path : str, pathlib.Path + Path to the file. + **kwargs : keyword arguments + Arbitrary keyword arguments forwarded to ``xarray.open_dataset()``, or + ``xarray.open_zarr()``. + Refer to the examples. + + Returns + ------- + package : imod.msw.MetaSwapPackage + Returns a package with data loaded from file. + + Examples + -------- + + To load a package from a file, e.g. a River package: + + >>> river = imod.mf6.River.from_file("river.nc") + + For large datasets, you likely want to process it in chunks. You can + forward keyword arguments to ``xarray.open_dataset()`` or + ``xarray.open_zarr()``: + + >>> river = imod.mf6.River.from_file("river.nc", chunks={"time": 1}) + + Refer to the xarray documentation for the possible keyword arguments. + """ + path = Path(path) + if path.suffix in (".zip", ".zarr"): + if path.suffix == ".zip": + with zarr.storage.ZipStore(path, mode="r") as store: + dataset = xr.open_zarr(store, **kwargs) + else: + dataset = xr.open_zarr(str(path), **kwargs) + else: + dataset = xr.open_dataset(path, chunks="auto", **kwargs) + + # Replace NaNs by None + for key, value in dataset.items(): + if _is_scalar_nan(value): + dataset[key] = None + + return cls._from_dataset(dataset) + + def to_file( + self, + modeldirectory: Path, + pkgname: str, + mdal_compliant: bool = False, + crs: Optional[Any] = None, + engine: EngineType = "netcdf4", + **kwargs, + ) -> Path: + """ + Write package to file. + + Parameters + ---------- + modeldirectory : pathlib.Path + Directory where to write the package file. + pkgname : str + Name of the package, used to create the file name. + mdal_compliant : bool, optional + Whether to write the package in MDAL-compliant format. Only used if + ``engine="netcdf4"``. Default is False. + crs : optional + Coordinate reference system to use when writing the package. Only + used if ``engine="netcdf4"`` Default is None. + engine : str, optional + File engine used to write packages. Options are ``'netcdf4'``, + ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other + softwares, for example QGIS. Zarr is optimized for big data, cloud + storage and parallel access. The ``'zarr.zip'`` option is an + experimental option which creates a zipped zarr store in a single + file, which is easier to copy and automatically compresses data as + well. Default is ``'netcdf4'``. + **kwargs : keyword arguments + Additional keyword arguments forwarded to the package serializer's + `to_file` method. + + Returns + ------- + pathlib.Path + Path to the written package file. + """ + + # All serialization logic is in the package serializers do not override + # this method (hence final decorator). + filewriter = create_package_serializer( + engine, mdal_compliant=mdal_compliant, crs=crs + ) + path = filewriter.to_file(self, modeldirectory, pkgname, **kwargs) + return path diff --git a/imod/tests/test_mf6/test_mf6_model.py b/imod/tests/test_mf6/test_mf6_model.py index c73587743..dc6dd6c4d 100644 --- a/imod/tests/test_mf6/test_mf6_model.py +++ b/imod/tests/test_mf6/test_mf6_model.py @@ -92,7 +92,7 @@ def test_key_assign(): def roundtrip(model, tmp_path): - dump_modelpkgs(tmp_path, "test") + dump_modelpkgs(model, tmp_path, "test") back = type(model).from_file(tmp_path / "test/test.toml") assert isinstance(back, type(model)) diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index 66f24b79b..0150be454 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -10,8 +10,7 @@ from imod.common.utilities.dump_model import dump_modelpkgs from imod.msw.copy_files import FileCopier from imod.msw.meteo_mapping import MeteoMapping -from imod.msw.model import DEFAULT_SETTINGS -from imod.msw.model import MetaSwapModel +from imod.msw.model import DEFAULT_SETTINGS, MetaSwapModel from imod.msw.utilities.parse import read_para_sim from imod.typing import GridDataArray, Imod5DataDict from imod.util.regrid import RegridderWeightsCache @@ -20,14 +19,14 @@ def roundtrip(msw_model, tmpdir_factory, name, engine): # TODO: look at the values? tmp_path = tmpdir_factory.mktemp(name) - dump_modelpkgs(msw_model, tmp_path, name, engine=engine) - back = MetaSwapModel.from_file(tmp_path / f"{simulation.name}.toml") + dump_modelpkgs(msw_model, tmp_path, name, engine=engine, validate=False) + back = MetaSwapModel.from_file(tmp_path, name) assert isinstance(back, MetaSwapModel) -def test_msw_pkgdump(msw_model, tmp_path): - dump_modelpkgs(msw_model, tmp_path, "testmodel", engine="netcdf4") -# back = MetaSwapModel.from_file(tmp_path / "testmodel.toml") -# assert isinstance(back, MetaSwapModel) + +def test_msw_pkgdump(msw_model, tmpdir_factory): + roundtrip(msw_model, tmpdir_factory, name="testmodel", engine="netcdf4") + def test_msw_model_write(msw_model, coupled_mf6_model, coupled_mf6wel, tmp_path): mf6_dis = coupled_mf6_model["GWF_1"]["dis"] From 0bf629048dd7353566d0ba4a536954c122076627 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Tue, 19 May 2026 10:15:49 +0200 Subject: [PATCH 4/9] fixing lint --- imod/msw/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imod/msw/model.py b/imod/msw/model.py index 3f8555820..9983cfc2b 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -13,10 +13,10 @@ import xarray as xr import imod.msw -from imod.common.interfaces.imodel import IModel -from imod.common.utilities.mask import mask_all_packages from imod.common.constants import MaskValues +from imod.common.interfaces.imodel import IModel from imod.common.utilities.clip import clip_by_grid +from imod.common.utilities.mask import mask_all_packages from imod.common.utilities.partitioninfo import create_partition_info from imod.common.utilities.regrid import regrid_imod5_cap_data from imod.common.utilities.value_filters import enforce_scalar From 817b309c3b41c36a2bd10d89e36327a8607a7a5c Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Tue, 19 May 2026 10:34:12 +0200 Subject: [PATCH 5/9] restored mf6 model.py, accidentally changed --- imod/mf6/model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/imod/mf6/model.py b/imod/mf6/model.py index b3486e0bd..e895a29f1 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -6,18 +6,20 @@ import pathlib import warnings from pathlib import Path -from typing import List, Optional, Tuple, Union +from typing import Any, List, Optional, Tuple, Union import cftime import jinja2 import numpy as np import tomli +import tomli_w import xarray as xr import xugrid as xu from jinja2 import Template import imod from imod.common.interfaces.imodel import IModel +from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase from imod.common.utilities.clip import clip_box_dataset from imod.common.utilities.mask import mask_all_packages From aaa56baff4bfdbd8d9a1686b3f357dd7ecaa9332 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Wed, 20 May 2026 09:41:09 +0200 Subject: [PATCH 6/9] Extending the checks on the restored metaswap model --- imod/tests/test_msw/test_model.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index 50b56e9af..77ee59463 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -26,6 +26,14 @@ def roundtrip(msw_model, tmpdir_factory, name, engine): dump_modelpkgs(msw_model, tmp_path, name, engine=engine, validate=False) back = MetaSwapModel.from_file(tmp_path, name) assert isinstance(back, MetaSwapModel) + pkgkeycheck = set(msw_model.keys()) == set(back.keys()) + assert pkgkeycheck + pkgcheck = all( + msw_model[pkgname].dataset.equals(back[pkgname].dataset) + for pkgname in msw_model.keys() + ) + assert pkgcheck + assert msw_model.simulation_settings == back.simulation_settings def test_msw_pkgdump(msw_model, tmpdir_factory): From 5fb4dacdb011d147589bf14dbf4676148784d516 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Mon, 1 Jun 2026 16:11:29 +0200 Subject: [PATCH 7/9] update api ref --- docs/api/changelog.rst | 3 ++ docs/api/msw.rst | 3 +- imod/common/utilities/dump_model.py | 34 +---------------------- imod/mf6/model.py | 11 ++++++++ imod/mf6/simulation.py | 4 +-- imod/msw/model.py | 40 +++++++++++++++++++++++++++ imod/tests/test_mf6/test_mf6_model.py | 4 +-- imod/tests/test_msw/test_model.py | 3 +- 8 files changed, 62 insertions(+), 40 deletions(-) diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index a1bfc7b5e..3778a8ff7 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -102,6 +102,9 @@ Added - Functionality to dump and load MODFLOW 6 simulations to/from zarr and zipstore formats. See :meth:`imod.mf6.Modflow6Simulation.dump` and :meth:`imod.mf6.Modflow6Simulation.from_file` for more information. +- Functionality to dump and load MetaSwap models to/from netcdf + format. See :meth:`imod.msw.MetaSwapModel.dump` and + :meth:`imod.msw.MetaSwapModel.from_file` for more information. Changed ~~~~~~~ diff --git a/docs/api/msw.rst b/docs/api/msw.rst index c42fe5ab9..ca1a857a2 100644 --- a/docs/api/msw.rst +++ b/docs/api/msw.rst @@ -11,6 +11,7 @@ Model objects & methods MetaSwapModel MetaSwapModel.write + MetaSwapModel.dump MetaSwapModel.from_imod5_data MetaSwapModel.regrid_like MetaSwapModel.clip_box @@ -172,4 +173,4 @@ Mappings CouplerMapping.clip_box CouplerMapping.from_imod5_data CouplerMapping.get_regrid_methods - CouplerMapping.write \ No newline at end of file + CouplerMapping.write diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py index 76ffe6d09..a43d559aa 100644 --- a/imod/common/utilities/dump_model.py +++ b/imod/common/utilities/dump_model.py @@ -1,26 +1,18 @@ import collections -import inspect -import pathlib from pathlib import Path from typing import Any, Optional -import tomli import tomli_w -import imod.mf6 -import imod.msw from imod.common.interfaces.imodel import IModel -from imod.common.interfaces.ipackage import IPackage from imod.common.serializer import EngineType from imod.logging.logging_decorators import standard_log_decorator -from imod.mf6.model import Modflow6Model from imod.mf6.validation_settings import ValidationSettings -from imod.msw.model import MetaSwapModel from imod.schemata import ValidationError @standard_log_decorator() -def dump_modelpkgs( +def _dump_model( model: IModel, directory, modelname, @@ -91,27 +83,3 @@ def dump_modelpkgs( tomli_w.dump(toml_content, f) return toml_path - - -def from_file(instance, toml_path): - if isinstance(instance, Modflow6Model): - modref = imod.mf6 - if isinstance(instance, MetaSwapModel): - modref = imod.msw - pkg_classes = { - name: pkg_cls - for name, pkg_cls in inspect.getmembers(modref, inspect.isclass) - if issubclass(pkg_cls, IPackage) - } - - toml_path = pathlib.Path(toml_path) - with open(toml_path, "rb") as f: - toml_content = tomli.load(f) - - parentdir = toml_path.parent - for key, entry in toml_content.items(): - for pkgname, path in entry.items(): - pkg_cls = pkg_classes[key] - instance[pkgname] = pkg_cls.from_file(parentdir / path) - - return instance diff --git a/imod/mf6/model.py b/imod/mf6/model.py index e895a29f1..825c3ef3d 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -22,6 +22,7 @@ from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase from imod.common.utilities.clip import clip_box_dataset +from imod.common.utilities.dump_model import _dump_model from imod.common.utilities.mask import mask_all_packages from imod.common.utilities.regrid import _regrid_like from imod.common.utilities.schemata import ( @@ -658,6 +659,16 @@ def dump( toml_content: dict[str, Any] = collections.defaultdict(dict) + _dump_model( + self, + modeldirectory, + modelname, + validate=validate, + mdal_compliant=mdal_compliant, + crs=crs, + engine=engine, + ) + for pkgname, pkg in self.items(): pkg_path = pkg.to_file( modeldirectory, diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index 02f70c93d..cd95a583a 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -24,7 +24,7 @@ from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo from imod.common.utilities.dataclass_type import DataclassType -from imod.common.utilities.dump_model import dump_modelpkgs +from imod.common.utilities.dump_model import _dump_model 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 @@ -1038,7 +1038,7 @@ def dump( for key, value in self.items(): cls_name = type(value).__name__ if isinstance(value, Modflow6Model): - model_toml_path = dump_modelpkgs( + model_toml_path = _dump_model( value, directory, key, validate, mdal_compliant, crs, engine=engine ) toml_content[cls_name][key] = model_toml_path.relative_to( diff --git a/imod/msw/model.py b/imod/msw/model.py index 9983cfc2b..f286577f8 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -15,7 +15,9 @@ import imod.msw from imod.common.constants import MaskValues from imod.common.interfaces.imodel import IModel +from imod.common.serializer import EngineType from imod.common.utilities.clip import clip_by_grid +from imod.common.utilities.dump_model import _dump_model from imod.common.utilities.mask import mask_all_packages from imod.common.utilities.partitioninfo import create_partition_info from imod.common.utilities.regrid import regrid_imod5_cap_data @@ -403,6 +405,44 @@ def from_file(cls, directory, modelname): return instance + def dump( + self, + directory: Union[str, Path], + modelname: Optional[str] = None, + validate: Optional[bool] = True, + mdal_compliant: bool = False, + crs: Optional[str] = None, + engine: EngineType = "netcdf4", + ): + """ + Dump model packages to netCDF files and create a toml file with paths to these files. + The MetaSWAP model can be reloaded from the dumped files using the :func:`from_file` method. + + Parameters + ---------- + directory: Path or str + Directory to dump model in. A subdirectory with the name of the model will be created in this directory, and the files will be dumped there. + modelname: str, optional + Name of the model. This will be used as the name of the subdirectory where the files are dumped, and in the name of the toml file. If not provided, it defaults to the value of ``self._model_name``. + validate: bool, optional + Whether to perform validation before dumping the model. If True, the model will be validated using the validation functionality in iMOD. If validation errors are found, a ValidationError is raised and the model is not dumped. Default is True. + mdal_compliant: bool, optional + Whether to write the files in a format compliant with the MDAL specification. This can be used to make the files compatible with software that supports MDAL, such as QGIS. Default is False. + crs: str, optional + Coordinate reference system to use in the dumped files. This should be a string in a format recognized by the pyproj library, for example "EPSG:28992". If not provided, no CRS information is included in the files. + engine: EngineType, optional + File engine used to write packages. + """ + _dump_model( + self, + directory=directory, + modelname=modelname, + validate=validate, + mdal_compliant=mdal_compliant, + crs=crs, + engine=engine, + ) + def regrid_like( self, mf6_regridded_dis: StructuredDiscretization, diff --git a/imod/tests/test_mf6/test_mf6_model.py b/imod/tests/test_mf6/test_mf6_model.py index dc6dd6c4d..345dfd6b9 100644 --- a/imod/tests/test_mf6/test_mf6_model.py +++ b/imod/tests/test_mf6/test_mf6_model.py @@ -11,7 +11,7 @@ from xugrid.core.wrap import UgridDataArray import imod -from imod.common.utilities.dump_model import dump_modelpkgs +from imod.common.utilities.dump_model import _dump_model from imod.mf6 import ConstantConcentration, ConstantHead from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.model import Modflow6Model @@ -92,7 +92,7 @@ def test_key_assign(): def roundtrip(model, tmp_path): - dump_modelpkgs(model, tmp_path, "test") + _dump_model(model, tmp_path, "test") back = type(model).from_file(tmp_path / "test/test.toml") assert isinstance(back, type(model)) diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index 77ee59463..6481f9ed8 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -8,7 +8,6 @@ from pytest_cases import parametrize_with_cases from imod import mf6, msw -from imod.common.utilities.dump_model import dump_modelpkgs from imod.msw.copy_files import FileCopier from imod.msw.meteo_grid import MeteoGridCopy from imod.msw.meteo_mapping import MeteoMapping @@ -23,7 +22,7 @@ def roundtrip(msw_model, tmpdir_factory, name, engine): # TODO: look at the values? tmp_path = tmpdir_factory.mktemp(name) - dump_modelpkgs(msw_model, tmp_path, name, engine=engine, validate=False) + msw_model.dump(tmp_path, name, engine=engine, validate=False) back = MetaSwapModel.from_file(tmp_path, name) assert isinstance(back, MetaSwapModel) pkgkeycheck = set(msw_model.keys()) == set(back.keys()) From b5d33a1e3e5b7424365bcd614105d1b958aa561e Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Thu, 4 Jun 2026 14:13:22 +0200 Subject: [PATCH 8/9] added tests for other formats and example line to docstring --- imod/msw/model.py | 19 +++++++++++++++++++ imod/tests/test_msw/test_model.py | 10 +++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/imod/msw/model.py b/imod/msw/model.py index f286577f8..2d6172adb 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -432,7 +432,26 @@ def dump( Coordinate reference system to use in the dumped files. This should be a string in a format recognized by the pyproj library, for example "EPSG:28992". If not provided, no CRS information is included in the files. engine: EngineType, optional File engine used to write packages. + engine : str, optional + File engine used to write packages. Options are ``'netcdf4'``, + ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other + softwares, for example QGIS. Zarr is optimized for big data, cloud + storage and parallel access. The ``'zarr.zip'`` option is an + experimental option which creates a zipped zarr store in a single + file, which is easier to copy and automatically compresses data as + well. Default is ``'netcdf4'``. + + Returns + ------- + Path + Path to the created toml file which contains the paths to the dumped package files. The package files are dumped in the same directory as the toml file. + + Example + ------- + >>> tmp_path = tmpdir_factory.mktemp(name) + >>> msw_model.dump(tmp_path, name, engine=engine, validate=False) """ + _dump_model( self, directory=directory, diff --git a/imod/tests/test_msw/test_model.py b/imod/tests/test_msw/test_model.py index 6481f9ed8..771fed5ba 100644 --- a/imod/tests/test_msw/test_model.py +++ b/imod/tests/test_msw/test_model.py @@ -35,10 +35,18 @@ def roundtrip(msw_model, tmpdir_factory, name, engine): assert msw_model.simulation_settings == back.simulation_settings -def test_msw_pkgdump(msw_model, tmpdir_factory): +def test_msw_pkgdump_nc(msw_model, tmpdir_factory): roundtrip(msw_model, tmpdir_factory, name="testmodel", engine="netcdf4") +def test_msw_pkgdump_zarr(msw_model, tmpdir_factory): + roundtrip(msw_model, tmpdir_factory, name="testmodel", engine="zarr") + + +def test_msw_pkgdump_zarrzip(msw_model, tmpdir_factory): + roundtrip(msw_model, tmpdir_factory, name="testmodel", engine="zarr.zip") + + def test_msw_model_write(msw_model, coupled_mf6_model, coupled_mf6wel, tmp_path): mf6_dis = coupled_mf6_model["GWF_1"]["dis"] From 480148551b9111f695416335b6ab09b83d953079 Mon Sep 17 00:00:00 2001 From: Robert Leander Date: Thu, 4 Jun 2026 18:22:22 +0200 Subject: [PATCH 9/9] processing review comments --- imod/common/utilities/dump_model.py | 3 +- imod/common/utilities/value_filters.py | 14 ++- imod/mf6/model.py | 98 ++++++++---------- imod/mf6/pkgbase.py | 17 +--- imod/mf6/simulation.py | 5 +- imod/msw/model.py | 135 +++---------------------- imod/msw/pkgbase.py | 16 +-- imod/tests/test_mf6/test_mf6_model.py | 3 +- 8 files changed, 80 insertions(+), 211 deletions(-) diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py index a43d559aa..2c4481b08 100644 --- a/imod/common/utilities/dump_model.py +++ b/imod/common/utilities/dump_model.py @@ -12,7 +12,7 @@ @standard_log_decorator() -def _dump_model( +def dump_model( model: IModel, directory, modelname, @@ -75,6 +75,7 @@ def _dump_model( ) toml_content[type(pkg).__name__][pkgname] = pkg_path.name + # simulation settings are only relevant/present for MetaSwap models (msw) if hasattr(model, "simulation_settings"): toml_content["simulation_settings"] = model.simulation_settings diff --git a/imod/common/utilities/value_filters.py b/imod/common/utilities/value_filters.py index 6f94cacfb..2c8285920 100644 --- a/imod/common/utilities/value_filters.py +++ b/imod/common/utilities/value_filters.py @@ -1,10 +1,22 @@ +import numbers from typing import Any import numpy as np import xarray as xr from xarray.core.utils import is_scalar -from imod.typing import GridDataset +from imod.typing import GridDataArray, GridDataset + + +def is_scalar_nan(da: GridDataArray): + """ + Test if is_scalar_nan, carefully avoid loading grids in memory + """ + scalar_data: bool = is_scalar(da) + if scalar_data: + stripped_value = da.to_numpy()[()] + return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value) # type: ignore[call-overload] + return False def is_valid(value: Any) -> bool: diff --git a/imod/mf6/model.py b/imod/mf6/model.py index 825c3ef3d..ee6a5ad9b 100644 --- a/imod/mf6/model.py +++ b/imod/mf6/model.py @@ -12,7 +12,6 @@ import jinja2 import numpy as np import tomli -import tomli_w import xarray as xr import xugrid as xu from jinja2 import Template @@ -22,7 +21,7 @@ from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase from imod.common.utilities.clip import clip_box_dataset -from imod.common.utilities.dump_model import _dump_model +from imod.common.utilities.dump_model import dump_model from imod.common.utilities.mask import mask_all_packages from imod.common.utilities.regrid import _regrid_like from imod.common.utilities.schemata import ( @@ -618,50 +617,51 @@ def dump( engine: EngineType = "netcdf4", ) -> Path: """ - Dump simulation to files. Writes a model definition as .TOML file, which - points to data for each package. Each package is stored as a separate - NetCDF. Structured grids are saved as regular NetCDFs, unstructured - grids are saved as UGRID NetCDF. Structured grids are always made GDAL - compliant, unstructured grids can be made MDAL compliant optionally. - - Parameters - ---------- - directory: str or Path - directory to dump simulation into. - modelname: str - modelname, will be used to create a subdirectory. - validate: bool, optional - Whether to validate simulation data. Defaults to True. - mdal_compliant: bool, optional - Convert data with - :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL - compliant unstructured grids. Defaults to False. - crs: Any, optional - Anything accepted by rasterio.crs.CRS.from_user_input - Requires ``rioxarray`` installed. - engine : str, optional - File engine used to write packages. Options are ``'netcdf4'``, - ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other - softwares, for example QGIS. Zarr is optimized for big data, cloud - storage and parallel access. The ``'zarr.zip'`` option is an - experimental option which creates a zipped zarr store in a single - file, which is easier to copy and automatically compresses data as - well. Default is ``'netcdf4'``. + Dump simulation to files. Writes a model definition as .TOML file, which + points to data for each package. Each package is stored as a separate + NetCDF. Structured grids are saved as regular NetCDFs, unstructured + grids are saved as UGRID NetCDF. Structured grids are always made GDAL + compliant, unstructured grids can be made MDAL compliant optionally. + + Parameters + ---------- + directory: str or Path + directory to dump simulation into. + modelname: str + modelname, will be used to create a subdirectory. + validate: bool, optional + Whether to validate simulation data. Defaults to True. + mdal_compliant: bool, optional + Convert data with + :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL + compliant unstructured grids. Defaults to False. + crs: Any, optional + Anything accepted by rasterio.crs.CRS.from_user_input + Requires ``rioxarray`` installed. + engine : str, optional + File engine used to write packages. Options are ``'netcdf4'``, + ``'zarr'``, and ``'zarr.zip'``. NetCDF4 is readable by many other + softwares, for example QGIS. Zarr is optimized for big data, cloud + storage and parallel access. The ``'zarr.zip'`` option is an + experimental option which creates a zipped zarr store in a single + file, which is easier to copy and automatically compresses data as + well. Default is ``'netcdf4'``. + Returns + ------- + Path + Path to the created toml file which contains the paths to the dumped package files. The package files are dumped in the same directory as the toml file. + + Example + ------- + >>> tmp_path = tmpdir_factory.mktemp(name) + >>> toml_path = mf6_model.dump(tmp_path, name, engine=engine, validate=False) + >>> back = ModflowModel.from_file(tmp_path, name) """ - modeldirectory = pathlib.Path(directory) / modelname - modeldirectory.mkdir(exist_ok=True, parents=True) - validation_context = ValidationSettings(validate=validate) - if validation_context.validate: - statusinfo = self.validate(modelname, validation_context) - if statusinfo.has_errors(): - raise ValidationError(statusinfo.to_string()) - toml_content: dict[str, Any] = collections.defaultdict(dict) - - _dump_model( + toml_path = dump_model( self, - modeldirectory, + directory, modelname, validate=validate, mdal_compliant=mdal_compliant, @@ -669,20 +669,6 @@ def dump( engine=engine, ) - for pkgname, pkg in self.items(): - pkg_path = pkg.to_file( - modeldirectory, - pkgname, - mdal_compliant=mdal_compliant, - crs=crs, - engine=engine, - ) - toml_content[type(pkg).__name__][pkgname] = pkg_path.name - - toml_path = modeldirectory / f"{modelname}.toml" - with open(toml_path, "wb") as f: - tomli_w.dump(toml_content, f) - return toml_path @classmethod diff --git a/imod/mf6/pkgbase.py b/imod/mf6/pkgbase.py index 6a5f212ca..cef68dc6e 100644 --- a/imod/mf6/pkgbase.py +++ b/imod/mf6/pkgbase.py @@ -1,16 +1,14 @@ import abc -import numbers from pathlib import Path from typing import TYPE_CHECKING, Any, Mapping, Optional, Self, final -import numpy as np import xarray as xr import xugrid as xu -from xarray.core.utils import is_scalar import imod from imod.common.interfaces.ipackagebase import IPackageBase from imod.common.serializer import EngineType, create_package_serializer +from imod.common.utilities.value_filters import is_scalar_nan from imod.typing.grid import ( GridDataArray, GridDataset, @@ -32,17 +30,6 @@ UTIL_PACKAGES = ("ats", "hpc") -def _is_scalar_nan(da: GridDataArray): - """ - Test if is_scalar_nan, carefully avoid loading grids in memory - """ - scalar_data: bool = is_scalar(da) - if scalar_data: - stripped_value = da.to_numpy()[()] - return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value) # type: ignore[call-overload] - return False - - class PackageBase(IPackageBase, abc.ABC): """ This class is used for storing a collection of Xarray DataArrays or UgridDataArrays @@ -148,7 +135,7 @@ def from_file(cls, path: str | Path, **kwargs) -> Self: # Replace NaNs by None for key, value in dataset.items(): - if _is_scalar_nan(value): + if is_scalar_nan(value): dataset[key] = None # to_netcdf converts strings into NetCDF "variable‑length UTF‑8 strings" diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py index cd95a583a..b2d254b7e 100644 --- a/imod/mf6/simulation.py +++ b/imod/mf6/simulation.py @@ -24,7 +24,6 @@ from imod.common.serializer import EngineType from imod.common.statusinfo import NestedStatusInfo from imod.common.utilities.dataclass_type import DataclassType -from imod.common.utilities.dump_model import _dump_model 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 @@ -1038,8 +1037,8 @@ def dump( for key, value in self.items(): cls_name = type(value).__name__ if isinstance(value, Modflow6Model): - model_toml_path = _dump_model( - value, directory, key, validate, mdal_compliant, crs, engine=engine + model_toml_path = value.dump( + directory, key, validate, mdal_compliant, crs, engine=engine ) toml_content[cls_name][key] = model_toml_path.relative_to( directory diff --git a/imod/msw/model.py b/imod/msw/model.py index 2d6172adb..bba41533c 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -17,8 +17,7 @@ from imod.common.interfaces.imodel import IModel from imod.common.serializer import EngineType from imod.common.utilities.clip import clip_by_grid -from imod.common.utilities.dump_model import _dump_model -from imod.common.utilities.mask import mask_all_packages +from imod.common.utilities.dump_model import dump_model from imod.common.utilities.partitioninfo import create_partition_info from imod.common.utilities.regrid import regrid_imod5_cap_data from imod.common.utilities.value_filters import enforce_scalar @@ -390,12 +389,9 @@ def from_file(cls, directory, modelname): toml_content = tomli.load(f) parentdir = toml_path.parent - if issubclass(cls, MetaSwapModel): - simulation_settings = toml_content.get("simulation_settings", {}) - unsa_svat_path = simulation_settings.get("unsa_svat_path", "") - instance = cls(unsa_svat_path, simulation_settings) - else: - instance = cls() + simulation_settings = toml_content.get("simulation_settings", {}) + unsa_svat_path = simulation_settings.get("unsa_svat_path", "") + instance = cls(unsa_svat_path, simulation_settings) for key, entry in toml_content.items(): if key != "simulation_settings": @@ -449,10 +445,11 @@ def dump( Example ------- >>> tmp_path = tmpdir_factory.mktemp(name) - >>> msw_model.dump(tmp_path, name, engine=engine, validate=False) + >>> toml_path = msw_model.dump(tmp_path, name, engine=engine, validate=False) + >>> back = MetaSwapModel.from_file(tmp_path, name) """ - _dump_model( + toml_path = dump_model( self, directory=directory, modelname=modelname, @@ -461,6 +458,7 @@ def dump( crs=crs, engine=engine, ) + return toml_path def regrid_like( self, @@ -774,138 +772,37 @@ def from_imod5_data( return model def _is_splitting_supported(self) -> Tuple[bool, str]: - """ - Returns True if all the packages in the model supports splitting. If one - of the packages in the model does not support splitting, it returns the - name of the first one. - - Returns - ------- - Tuple[bool, str] - A tuple where the first element is a boolean indicating if splitting - is supported, and the second element is the name of the first package - that does not support splitting, or an empty string if all packages - support splitting. - """ - for package_name, package in self.items(): - if not package._is_splitting_supported(): - return False, package_name - return True, "" + raise NotImplementedError def _is_regridding_supported(self) -> Tuple[bool, str]: - """ - Returns True if all the packages in the model supports regridding. If one - of the packages in the model does not support regridding, it returns the - name of the first one. - - Returns - ------- - Tuple[bool, str] - A tuple where the first element is a boolean indicating if regridding - is supported, and the second element is the name of the first package - that does not support regridding, or an empty string if all packages - support regridding. - """ - for package_name, package in self.items(): - if not package._is_regridding_supported(): - return False, package_name - return True, "" + raise NotImplementedError def _is_clipping_supported(self) -> Tuple[bool, str]: - """ - Returns True if all the packages in the model supports clipping. If one - of the packages in the model does not support clipping, it returns the - name of the first one. - - Returns - ------- - Tuple[bool, str] - A tuple where the first element is a boolean indicating if clipping - is supported, and the second element is the name of the first package - that does not support clipping, or an empty string if all packages - support clipping. - """ - for package_name, package in self.items(): - if not package._is_clipping_supported(): - return False, package_name - return True, "" + raise NotImplementedError def purge_empty_packages( self, model_name: Optional[str] = "", ignore_time: bool = False ) -> None: - """ - This method removes empty packages from the model in place. - - Parameters - ---------- - model_name: str, optional - Name of the model, used for logging. - ignore_time: bool, optional - If False, packages are considered empty if they have no data at all - timesteps. If True, packages are considered empty if they have no - data at the first time step. The latter can increase performance - considerably. - """ - empty_packages = [ - package_name - for package_name, package in self.items() - if package.is_empty(ignore_time=ignore_time) - ] - for package_name in empty_packages: - self.pop(package_name) + raise NotImplementedError def mask_all_packages( self, mask: GridDataArray, ignore_time_purge_empty: bool = False, ): - """ - This function applies a mask to all packages in a model. The mask must - be presented as an idomain-like integer array that has 0 (inactive) or - <0 (vertical passthrough) values in filtered cells and >0 in active - cells. - Masking will overwrite idomain with the mask where the mask is <=0. - Where the mask is >0, the original value of idomain will be kept. Masking - will update the packages accordingly, blanking their input where needed, - and is therefore not a reversible operation. - - Parameters - ---------- - mask: xr.DataArray, xu.UgridDataArray of ints - idomain-like integer array. >0 sets cells to active, 0 sets cells to inactive, - <0 sets cells to vertical passthrough - ignore_time_purge_empty: bool, default False - Whether to ignore time dimension when purging empty packages. Can - improve performance when masking models with many time steps. - """ - - mask_all_packages(self, mask, ignore_time_purge_empty) + raise NotImplementedError @property def model_id(self) -> str: - if self._model_id is None: - return "MetaSwapModel" - return self._model_id + raise NotImplementedError @property def options(self) -> dict: - return self._options + raise NotImplementedError @property def domain(self) -> GridDataArray: - """ - Get the active domain array as an integer array compatible with idomain. - - Returns - ------- - GridDataArray - Integer array where: - - 1 indicates active cells - - 0 indicates inactive cells - """ - grid_key = self.get_pkgkey(GridData) - active = self[grid_key]["active"] - return active.astype(int) + raise NotImplementedError def validate( self, diff --git a/imod/msw/pkgbase.py b/imod/msw/pkgbase.py index 55b362673..73686b7b1 100644 --- a/imod/msw/pkgbase.py +++ b/imod/msw/pkgbase.py @@ -1,5 +1,4 @@ import abc -import numbers from copy import deepcopy from pathlib import Path from typing import Any, Optional, Self, TextIO, TypeAlias, Union @@ -9,7 +8,6 @@ import pandas as pd import xarray as xr import zarr # nb: in modflow6, this is a try...except -from xarray.core.utils import is_scalar from imod.common.serializer import EngineType, create_package_serializer from imod.common.utilities.clip import clip_spatial_box, clip_time_slice @@ -17,6 +15,7 @@ from imod.common.utilities.regrid import ( _regrid_like, ) +from imod.common.utilities.value_filters import is_scalar_nan from imod.mf6.dis import StructuredDiscretization from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.msw.fixed_format import format_fixed_width @@ -27,17 +26,6 @@ DataDictType: TypeAlias = dict[str, IntArray | int | str] -def _is_scalar_nan(da: GridDataArray): - """ - Test if is_scalar_nan, carefully avoid loading grids in memory - """ - scalar_data: bool = is_scalar(da) - if scalar_data: - stripped_value = da.to_numpy()[()] - return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value) # type: ignore[call-overload] - return False - - class MetaSwapPackage(abc.ABC): """ MetaSwapPackage is used to share methods for Metaswap packages. @@ -448,7 +436,7 @@ def from_file(cls, path: str | Path, **kwargs) -> Self: # Replace NaNs by None for key, value in dataset.items(): - if _is_scalar_nan(value): + if is_scalar_nan(value): dataset[key] = None return cls._from_dataset(dataset) diff --git a/imod/tests/test_mf6/test_mf6_model.py b/imod/tests/test_mf6/test_mf6_model.py index 345dfd6b9..7495cbe51 100644 --- a/imod/tests/test_mf6/test_mf6_model.py +++ b/imod/tests/test_mf6/test_mf6_model.py @@ -11,7 +11,6 @@ from xugrid.core.wrap import UgridDataArray import imod -from imod.common.utilities.dump_model import _dump_model from imod.mf6 import ConstantConcentration, ConstantHead from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.model import Modflow6Model @@ -92,7 +91,7 @@ def test_key_assign(): def roundtrip(model, tmp_path): - _dump_model(model, tmp_path, "test") + model.dump(tmp_path, "test") back = type(model).from_file(tmp_path / "test/test.toml") assert isinstance(back, type(model))