Issue #1222 validation context (#1226)

Fixes #1222

# Description
- Add a ``ValidationContext`` dataclass
- Add a ``_validation_context`` attribute to ``Modflow6Simulation``;
this replaces the ``_is_from_imod5`` attribute.
- The ``ValidationContext`` contains an attribute for strict well
checks, turned on by default. This is set to False when calling
``from_imod5`` or for split simulations.
- Adds a ``_to_mf6_pkg`` method in a similar design as proposed in
#1223, this to preserve public API.
- Refactor ``WriteContext``, to make it a dataclass again. I had to
ignore type annotation for ``write_directory``, otherwise MyPy would
throw errors. The whole property shebang presumably started with MyPy
throwing errors. Reverting it back to a dataclass reduces the lines of
code considerably, which makes it more maintainable.
- Use jit for examples run, this speeds them up considerably. The
examples ran into a TimeOut on TeamCity, and this reduces the change of
that happening again.

# Checklist
<!---
Before requesting review, please go through this checklist:
-->

- [x] Links to correct issue
- [x] Update changelog, if changes affect users
- [x] PR title starts with ``Issue #nr``, e.g. ``Issue #737``
- [ ] Unit tests were added
- [ ] **If feature added**: Added/extended example

Author: JoerivanEngelen

docs/api/changelog.rst

@@ -46,12 +46,9 @@ Changed
   when elevations are above model top for all methods in
   :func:`imod.prepare.ALLOCATION_OPTION`.
 - :meth:`imod.mf6.Well.to_mf6_pkg` got a new argument:
-  ``error_on_well_removal``, which controls the behavior for when wells are
+  ``strict_well_validation``, which controls the behavior for when wells are
   removed entirely during their assignment to layers. This replaces the
-  ``is_partioned`` argument. Note: if ``is_partioned`` was set to True before,
-  this means you need to set ``error_on_removal`` to False now to get the same
-  behavior.
-
+  ``is_partitioned`` argument.
 
 Added
 ~~~~~

imod/mf6/model.py

@@ -27,6 +27,7 @@
 from imod.mf6.utilities.mf6hfb import merge_hfb_packages
 from imod.mf6.utilities.regrid import RegridderWeightsCache, _regrid_like
 from imod.mf6.validation import pkg_errors_to_status_info
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.wel import GridAgnosticWell
 from imod.mf6.write_context import WriteContext
 from imod.schemata import ValidationError
@@ -240,7 +241,11 @@ def validate(self, model_name: str = "") -> StatusInfoBase:
 
     @standard_log_decorator()
     def write(
-        self, modelname, globaltimes, validate: bool, write_context: WriteContext
+        self,
+        modelname,
+        globaltimes,
+        write_context: WriteContext,
+        validate_context: ValidationContext,
     ) -> StatusInfoBase:
         """
         Write model namefile
@@ -250,7 +255,7 @@ def write(
         workdir = write_context.simulation_directory
         modeldirectory = workdir / modelname
         Path(modeldirectory).mkdir(exist_ok=True, parents=True)
-        if validate:
+        if validate_context.validate:
             model_status_info = self.validate(modelname)
             if model_status_info.has_errors():
                 return model_status_info
@@ -271,16 +276,12 @@ def write(
                 if issubclass(type(pkg), GridAgnosticWell):
                     top, bottom, idomain = self.__get_domain_geometry()
                     k = self.__get_k()
-                    error_on_well_removal = (not pkg_write_context.is_partitioned) & (
-                        not pkg_write_context.is_from_imod5
-                    )
-                    mf6_well_pkg = pkg.to_mf6_pkg(
+                    mf6_well_pkg = pkg._to_mf6_pkg(
                         idomain,
                         top,
                         bottom,
                         k,
-                        validate,
-                        error_on_well_removal,
+                        validate_context,
                     )
                     mf6_well_pkg.write(
                         pkgname=pkg_name,

imod/mf6/simulation.py

@@ -43,6 +43,7 @@
 from imod.mf6.statusinfo import NestedStatusInfo
 from imod.mf6.utilities.mask import _mask_all_models
 from imod.mf6.utilities.regrid import _regrid_like
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.write_context import WriteContext
 from imod.prepare.topsystem.default_allocation_methods import (
     SimulationAllocationOptions,
@@ -97,7 +98,7 @@ def __init__(self, name):
         self.name = name
         self.directory = None
         self._initialize_template()
-        self._is_from_imod5 = False
+        self._validation_context = ValidationContext()
 
     def __setitem__(self, key, value):
         super().__setitem__(key, value)
@@ -254,10 +255,9 @@ def write(
         """
         # create write context
         write_context = WriteContext(directory, binary, use_absolute_paths)
+        self._validation_context.validate = validate
         if self.is_split():
-            write_context.is_partitioned = True
-        if self._is_from_imod5:
-            write_context.is_from_imod5 = True
+            self._validation_context.strict_well_validation = False
 
         # Check models for required content
         for key, model in self.items():
@@ -299,8 +299,8 @@ def write(
                     value.write(
                         modelname=key,
                         globaltimes=globaltimes,
-                        validate=validate,
                         write_context=model_write_context,
+                        validate_context=self._validation_context,
                     )
                 )
             elif isinstance(value, Package):
@@ -1367,7 +1367,7 @@ def from_imod5_data(
         -------
         """
         simulation = Modflow6Simulation("imported_simulation")
-        simulation._is_from_imod5 = True
+        simulation._validation_context.strict_well_validation = False
 
         # import GWF model,
         groundwaterFlowModel = GroundwaterFlowModel.from_imod5_data(

imod/mf6/validation_context.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ValidationContext:
+    validate: bool = True
+    strict_well_validation: bool = True

imod/mf6/wel.py

@@ -31,6 +31,7 @@
 from imod.mf6.utilities.dataset import remove_inactive
 from imod.mf6.utilities.grid import broadcast_to_full_domain
 from imod.mf6.validation import validation_pkg_error_message
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.write_context import WriteContext
 from imod.prepare import assign_wells
 from imod.prepare.cleanup import cleanup_wel
@@ -388,7 +389,7 @@ def to_mf6_pkg(
         bottom: GridDataArray,
         k: GridDataArray,
         validate: bool = False,
-        error_on_well_removal: bool = True,
+        strict_well_validation: bool = True,
     ) -> Mf6Wel:
         """
         Write package to Modflow 6 package.
@@ -415,9 +416,10 @@ def to_mf6_pkg(
             Grid with bottom of model layers.
         k: {xarry.DataArray, xugrid.UgridDataArray}
             Grid with hydraulic conductivities.
-        validate: bool
+        validate: bool, default True
             Run validation before converting
-        error_on_well_removal: bool
+        strict_well_validation: bool, default True
+            Set well validation strict:
             Throw error if well is removed entirely during its assignment to
             layers.
 
@@ -426,7 +428,20 @@ def to_mf6_pkg(
         Mf6Wel
             Object with wells as list based input.
         """
-        if validate:
+        validation_context = ValidationContext(
+            validate=validate, strict_well_validation=strict_well_validation
+        )
+        return self._to_mf6_pkg(active, top, bottom, k, validation_context)
+
+    def _to_mf6_pkg(
+        self,
+        active: GridDataArray,
+        top: GridDataArray,
+        bottom: GridDataArray,
+        k: GridDataArray,
+        validation_context: ValidationContext,
+    ) -> Mf6Wel:
+        if validation_context.validate:
             errors = self._validate(self._write_schemata)
             if len(errors) > 0:
                 message = validation_pkg_error_message(errors)
@@ -446,6 +461,7 @@ def to_mf6_pkg(
         message_assign = self.to_mf6_package_information(
             filtered_assigned_well_ids, reason_text="permeability/thickness constraints"
         )
+        error_on_well_removal = validation_context.strict_well_validation
         if error_on_well_removal and len(filtered_assigned_well_ids) > 0:
             logger.log(loglevel=LogLevel.ERROR, message=message_assign)
             raise ValidationError(message_assign)

imod/mf6/write_context.py

@@ -4,7 +4,6 @@
 from dataclasses import dataclass
 from os.path import relpath
 from pathlib import Path
-from typing import Optional, Union
 
 
 @dataclass
@@ -28,23 +27,18 @@ class WriteContext:
          it will be set to the simulation_directrory.
     """
 
-    def __init__(
-        self,
-        simulation_directory: Path = Path("."),
-        use_binary: bool = False,
-        use_absolute_paths: bool = False,
-        write_directory: Optional[Union[str, Path]] = None,
-    ):
-        self.__simulation_directory = Path(simulation_directory)
-        self.__use_binary = use_binary
-        self.__use_absolute_paths = use_absolute_paths
-        self.__write_directory = (
-            Path(write_directory)
-            if write_directory is not None
-            else self.__simulation_directory
+    simulation_directory: Path = Path(".")
+    use_binary: bool = False
+    use_absolute_paths: bool = False
+    write_directory: Path = None  # type: ignore
+
+    def __post_init__(self):
+        self.simulation_directory = Path(self.simulation_directory)
+        self.write_directory = (
+            Path(self.write_directory)
+            if self.write_directory is not None
+            else self.simulation_directory
         )
-        self.__is_partitioned = False
-        self.__is_from_imod5 = False
 
     def get_formatted_write_directory(self) -> Path:
         """
@@ -53,57 +47,21 @@ def get_formatted_write_directory(self) -> Path:
         be relative to the simulation directory, which makes it usable by MF6.
         """
         if self.use_absolute_paths:
-            return self.__write_directory
-        return Path(relpath(self.write_directory, self.__simulation_directory))
+            return self.write_directory
+        return Path(relpath(self.write_directory, self.simulation_directory))
 
     def copy_with_new_write_directory(self, new_write_directory: Path) -> WriteContext:
         new_context = deepcopy(self)
-        new_context.__write_directory = Path(new_write_directory)
+        new_context.write_directory = Path(new_write_directory)
         return new_context
 
-    @property
-    def simulation_directory(self) -> Path:
-        return self.__simulation_directory
-
-    @property
-    def use_binary(self) -> bool:
-        return self.__use_binary
-
-    @use_binary.setter
-    def use_binary(self, value) -> None:
-        self.__use_binary = value
-
-    @property
-    def use_absolute_paths(self) -> bool:
-        return self.__use_absolute_paths
-
-    @property
-    def write_directory(self) -> Path:
-        return self.__write_directory
-
     @property
     def root_directory(self) -> Path:
         """
         returns the simulation directory, or nothing, depending on use_absolute_paths; use this to compose paths
         that are in agreement with the use_absolute_paths setting.
         """
         if self.use_absolute_paths:
-            return self.__simulation_directory
+            return self.simulation_directory
         else:
             return Path("")
-
-    @property
-    def is_partitioned(self) -> bool:
-        return self.__is_partitioned
-
-    @is_partitioned.setter
-    def is_partitioned(self, value: bool) -> None:
-        self.__is_partitioned = value
-
-    @property
-    def is_from_imod5(self) -> bool:
-        return self.__is_from_imod5
-
-    @is_from_imod5.setter
-    def is_from_imod5(self, value: bool) -> None:
-        self.__is_from_imod5 = value

imod/tests/test_mf6/test_circle.py

@@ -9,6 +9,7 @@
 
 import imod
 from imod.logging import LoggerType, LogLevel
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.write_context import WriteContext
 
 
@@ -57,8 +58,8 @@ def test_gwfmodel_render(circle_model, tmp_path):
     simulation = circle_model
     globaltimes = simulation["time_discretization"]["time"].values
     gwfmodel = simulation["GWF_1"]
-    write_context = WriteContext()
-    actual = gwfmodel.render("GWF_1", write_context)
+    write_context1 = WriteContext()
+    actual = gwfmodel.render("GWF_1", write_context1)
     path = "GWF_1"
     expected = textwrap.dedent(
         f"""\
@@ -77,8 +78,9 @@ def test_gwfmodel_render(circle_model, tmp_path):
             """
     )
     assert actual == expected
-    context = WriteContext(tmp_path)
-    gwfmodel.write("GWF_1", globaltimes, True, context)
+    validation_context = ValidationContext(True)
+    write_context2 = WriteContext(tmp_path)
+    gwfmodel.write("GWF_1", globaltimes, write_context2, validation_context)
     assert (tmp_path / "GWF_1" / "GWF_1.nam").is_file()
     assert (tmp_path / "GWF_1").is_dir()
 
@@ -110,8 +112,8 @@ def test_gwfmodel_render_evt(circle_model_evt, tmp_path):
     simulation = circle_model_evt
     globaltimes = simulation["time_discretization"]["time"].values
     gwfmodel = simulation["GWF_1"]
-    write_context = WriteContext()
-    actual = gwfmodel.render("GWF_1", write_context)
+    write_context1 = WriteContext()
+    actual = gwfmodel.render("GWF_1", write_context1)
     path = "GWF_1"
     expected = textwrap.dedent(
         f"""\
@@ -131,7 +133,8 @@ def test_gwfmodel_render_evt(circle_model_evt, tmp_path):
             """
     )
     assert actual == expected
-    context = WriteContext(tmp_path)
-    gwfmodel.write("GWF_1", globaltimes, True, context)
+    validation_context = ValidationContext(True)
+    write_context2 = WriteContext(tmp_path)
+    gwfmodel.write("GWF_1", globaltimes, write_context2, validation_context)
     assert (tmp_path / "GWF_1" / "GWF_1.nam").is_file()
     assert (tmp_path / "GWF_1").is_dir()

imod/tests/test_mf6/test_ex01_twri.py

@@ -8,6 +8,7 @@
 import xarray as xr
 
 import imod
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.write_context import WriteContext
 from imod.schemata import ValidationError
 from imod.typing.grid import ones_like
@@ -351,6 +352,7 @@ def test_gwfmodel_render(twri_model, tmp_path):
     globaltimes = simulation["time_discretization"]["time"].values
     gwfmodel = simulation["GWF_1"]
     path = Path(tmp_path.stem).as_posix()
+    validation_context = ValidationContext(tmp_path)
     write_context = WriteContext(tmp_path)
     actual = gwfmodel.render(path, write_context)
     expected = textwrap.dedent(
@@ -373,7 +375,7 @@ def test_gwfmodel_render(twri_model, tmp_path):
             """
     )
     assert actual == expected
-    gwfmodel.write("GWF_1", globaltimes, True, write_context)
+    gwfmodel.write("GWF_1", globaltimes, write_context, validation_context)
     assert (tmp_path / "GWF_1" / "GWF_1.nam").is_file()
     assert (tmp_path / "GWF_1").is_dir()
 

imod/tests/test_mf6/test_mf6_logging.py

@@ -10,6 +10,7 @@
 
 import imod
 from imod.logging import LoggerType, LogLevel, standard_log_decorator
+from imod.mf6.validation_context import ValidationContext
 from imod.mf6.write_context import WriteContext
 
 out = StringIO()
@@ -80,6 +81,7 @@ def test_write_model_is_logged(
     # arrange
     logfile_path = tmp_path / "logfile.txt"
     transport_model = flow_transport_simulation["tpt_c"]
+    validation_context = ValidationContext()
     write_context = WriteContext(simulation_directory=tmp_path, use_binary=True)
     globaltimes = np.array(
         [
@@ -95,7 +97,9 @@ def test_write_model_is_logged(
             add_default_file_handler=False,
             add_default_stream_handler=True,
         )
-        transport_model.write("model.txt", globaltimes, True, write_context)
+        transport_model.write(
+            "model.txt", globaltimes, write_context, validation_context
+        )
 
     # assert
     with open(logfile_path, "r") as log_file: