Issue #708 Refactor dataset initialization (#722)

Fix #708

The changeset is incomplete, as it is not rolled out for all packages
yet. Therefore tests will fail, hence the draft status. This to make the
review process more focused, reviewers are asked to provide feedback on
the approach.

Changes:
- Variable names are inferred based on what is provided to locals().
This makes it possible to avoid having to manually assign variables in
the __init__ of each specific package.
- I had to add one function to work around this issue:
Deltares/xugrid#179
- Grids are merged with an exact join, this to avoid that xarray joins
variable coordinates in an unwanted way when dataarrays are
inconsistent, which causes issues like
#674
- Refactored the riv unit tests to use pytest cases, so that tests
(without the xr.merge) are run for both unstructured and structured
data.

Update:
- Variable names are not inferred with locals() anymore, but instead
with a ``pkg_init`` decorator.

docs/api/changelog.rst

@@ -14,6 +14,9 @@ Fixed
 - iMOD Python now supports versions of pandas >= 2
 - Fixed bugs with clipping :class:`imod.mf6.HorizontalFlowBarrier` for
   structured grids
+- Packages and boundary conditions in the ``imod.mf6`` module will now throw an
+  error upon initialization if coordinate labels are inconsistent amongst
+  variables 
 - Improved performance for merging structured multimodel Modflow 6 output
 - Bug where :function:`imod.formats.idf.open_subdomains` did not properly support custom
   patterns

imod/mf6/adv.py

@@ -18,8 +18,8 @@ class Advection(Package):
     _template = Package._initialize_template(_pkg_id)
 
     def __init__(self, scheme):
-        super().__init__()
-        self.dataset["scheme"] = scheme
+        dict_dataset = {"scheme": scheme}
+        super().__init__(dict_dataset)
 
     def render(self, directory, pkgname, globaltimes, binary):
         scheme = self.dataset["scheme"].item()

imod/mf6/boundary_condition.py

@@ -7,9 +7,13 @@
 import xarray as xr
 import xugrid as xu
 
-from imod.mf6.auxiliary_variables import get_variable_names
+from imod.mf6.auxiliary_variables import (
+    expand_transient_auxiliary_variables,
+    get_variable_names,
+)
 from imod.mf6.package import Package
 from imod.mf6.write_context import WriteContext
+from imod.typing.grid import GridDataArray
 
 
 def _dis_recarr(arrdict, layer, notnull):
@@ -64,6 +68,15 @@ class BoundaryCondition(Package, abc.ABC):
     not the array input which is used in :class:`Package`.
     """
 
+    def __init__(self, allargs: dict[str, GridDataArray | float | int | bool | str]):
+        super().__init__(allargs)
+        if "concentration" in allargs.keys() and allargs["concentration"] is None:
+            # Remove vars inplace
+            del self.dataset["concentration"]
+            del self.dataset["concentration_boundary_type"]
+        else:
+            expand_transient_auxiliary_variables(self)
+
     def _max_active_n(self):
         """
         Determine the maximum active number of cells that are active

imod/mf6/buy.py

@@ -110,17 +110,17 @@ def __init__(
         densityfile: str = None,
         validate: bool = True,
     ):
-        super().__init__(locals())
-        self.dataset["reference_density"] = reference_density
-        # Assign a shared index: this also forces equal lenghts
-        self.dataset["density_concentration_slope"] = assign_index(
-            density_concentration_slope
-        )
-        self.dataset["reference_concentration"] = assign_index(reference_concentration)
-        self.dataset["modelname"] = assign_index(modelname)
-        self.dataset["species"] = assign_index(species)
-        self.dataset["hhformulation_rhs"] = hhformulation_rhs
-        self.dataset["densityfile"] = densityfile
+        dict_dataset = {
+            "reference_density": reference_density,
+            # Assign a shared index: this also forces equal lenghts
+            "density_concentration_slope": assign_index(density_concentration_slope),
+            "reference_concentration": assign_index(reference_concentration),
+            "modelname": assign_index(modelname),
+            "species": assign_index(species),
+            "hhformulation_rhs": hhformulation_rhs,
+            "densityfile": densityfile,
+        }
+        super().__init__(dict_dataset)
         self.dependencies = []
         self._validate_init_schemata(validate)
 

imod/mf6/chd.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -126,17 +125,17 @@ def __init__(
         validate: bool = True,
         repeat_stress=None,
     ):
-        super().__init__(locals())
-        self.dataset["head"] = head
-        if concentration is not None:
-            self.dataset["concentration"] = concentration
-            self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            expand_transient_auxiliary_variables(self)
-        self.dataset["print_input"] = print_input
-        self.dataset["print_flows"] = print_flows
-        self.dataset["save_flows"] = save_flows
-        self.dataset["observations"] = observations
-        self.dataset["repeat_stress"] = repeat_stress
+        dict_dataset = {
+            "head": head,
+            "concentration": concentration,
+            "concentration_boundary_type": concentration_boundary_type,
+            "print_input": print_input,
+            "print_flows": print_flows,
+            "save_flows": save_flows,
+            "observations": observations,
+            "repeat_stress": repeat_stress,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def _validate(self, schemata, **kwargs):

imod/mf6/cnc.py

@@ -74,10 +74,12 @@ def __init__(
         observations=None,
         validate: bool = True,
     ):
-        super().__init__(locals())
-        self.dataset["concentration"] = concentration
-        self.dataset["print_input"] = print_input
-        self.dataset["print_flows"] = print_flows
-        self.dataset["save_flows"] = save_flows
-        self.dataset["observations"] = observations
+        dict_dataset = {
+            "concentration": concentration,
+            "print_input": print_input,
+            "print_flows": print_flows,
+            "save_flows": save_flows,
+            "observations": observations,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)

imod/mf6/dis.py

@@ -91,10 +91,12 @@ class StructuredDiscretization(Package):
     _skip_mask_arrays = ["bottom", "idomain"]
 
     def __init__(self, top, bottom, idomain, validate: bool = True):
-        super(__class__, self).__init__(locals())
-        self.dataset["idomain"] = idomain
-        self.dataset["top"] = top
-        self.dataset["bottom"] = bottom
+        dict_dataset = {
+            "idomain": idomain,
+            "top": top,
+            "bottom": bottom,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def _delrc(self, dx):

imod/mf6/disv.py

@@ -72,10 +72,12 @@ class VerticesDiscretization(Package):
     _skip_mask_arrays = ["bottom", "idomain"]
 
     def __init__(self, top, bottom, idomain, validate: bool = True):
-        super().__init__(locals())
-        self.dataset["idomain"] = idomain
-        self.dataset["top"] = top
-        self.dataset["bottom"] = bottom
+        dict_dataset = {
+            "idomain": idomain,
+            "top": top,
+            "bottom": bottom,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def render(self, directory, pkgname, binary):

imod/mf6/drn.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
@@ -115,18 +114,19 @@ def __init__(
         validate: bool = True,
         repeat_stress=None,
     ):
-        super().__init__(locals())
-        self.dataset["elevation"] = elevation
-        self.dataset["conductance"] = conductance
-        if concentration is not None:
-            self.dataset["concentration"] = concentration
-            self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            expand_transient_auxiliary_variables(self)
-        self.dataset["print_input"] = print_input
-        self.dataset["print_flows"] = print_flows
-        self.dataset["save_flows"] = save_flows
-        self.dataset["observations"] = observations
-        self.dataset["repeat_stress"] = repeat_stress
+        dict_dataset = {
+            "elevation": elevation,
+            "conductance": conductance,
+            "concentration": concentration,
+            "concentration_boundary_type": concentration_boundary_type,
+            "print_input": print_input,
+            "print_flows": print_flows,
+            "save_flows": save_flows,
+            "observations": observations,
+            "repeat_stress": repeat_stress,
+        }
+        super().__init__(dict_dataset)
+
         self._validate_init_schemata(validate)
 
     def _validate(self, schemata, **kwargs):

imod/mf6/dsp.py

@@ -145,13 +145,15 @@ def __init__(
         xt3d_rhs=False,
         validate: bool = True,
     ):
-        super().__init__(locals())
-        self.dataset["xt3d_off"] = xt3d_off
-        self.dataset["xt3d_rhs"] = xt3d_rhs
-        self.dataset["diffusion_coefficient"] = diffusion_coefficient
-        self.dataset["longitudinal_horizontal"] = longitudinal_horizontal
-        self.dataset["transversal_horizontal1"] = transversal_horizontal1
-        self.dataset["longitudinal_vertical"] = longitudinal_vertical
-        self.dataset["transversal_horizontal2"] = transversal_horizontal2
-        self.dataset["transversal_vertical"] = transversal_vertical
+        dict_dataset = {
+            "xt3d_off": xt3d_off,
+            "xt3d_rhs": xt3d_rhs,
+            "diffusion_coefficient": diffusion_coefficient,
+            "longitudinal_horizontal": longitudinal_horizontal,
+            "transversal_horizontal1": transversal_horizontal1,
+            "longitudinal_vertical": longitudinal_vertical,
+            "transversal_horizontal2": transversal_horizontal2,
+            "transversal_vertical": transversal_vertical,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)

imod/mf6/evt.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 
-from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
@@ -193,27 +192,27 @@ def __init__(
         validate: bool = True,
         repeat_stress=None,
     ):
-        super().__init__(locals())
-        self.dataset["surface"] = surface
-        self.dataset["rate"] = rate
-        self.dataset["depth"] = depth
         if ("segment" in proportion_rate.dims) ^ ("segment" in proportion_depth.dims):
             raise ValueError(
                 "Segment must be provided for both proportion_rate and"
                 " proportion_depth, or for none at all."
             )
-        self.dataset["proportion_rate"] = proportion_rate
-        self.dataset["proportion_depth"] = proportion_depth
-        if concentration is not None:
-            self.dataset["concentration"] = concentration
-            self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            expand_transient_auxiliary_variables(self)
-        self.dataset["fixed_cell"] = fixed_cell
-        self.dataset["print_input"] = print_input
-        self.dataset["print_flows"] = print_flows
-        self.dataset["save_flows"] = save_flows
-        self.dataset["observations"] = observations
-        self.dataset["repeat_stress"] = repeat_stress
+        dict_dataset = {
+            "surface": surface,
+            "rate": rate,
+            "depth": depth,
+            "proportion_rate": proportion_rate,
+            "proportion_depth": proportion_depth,
+            "concentration": concentration,
+            "concentration_boundary_type": concentration_boundary_type,
+            "fixed_cell": fixed_cell,
+            "print_input": print_input,
+            "print_flows": print_flows,
+            "save_flows": save_flows,
+            "observations": observations,
+            "repeat_stress": repeat_stress,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def _validate(self, schemata, **kwargs):

imod/mf6/ghb.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from imod.mf6.auxiliary_variables import expand_transient_auxiliary_variables
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.regridding_utils import RegridderType
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
@@ -133,18 +132,18 @@ def __init__(
         validate: bool = True,
         repeat_stress=None,
     ):
-        super().__init__(locals())
-        self.dataset["head"] = head
-        self.dataset["conductance"] = conductance
-        if concentration is not None:
-            self.dataset["concentration"] = concentration
-            self.dataset["concentration_boundary_type"] = concentration_boundary_type
-            expand_transient_auxiliary_variables(self)
-        self.dataset["print_input"] = print_input
-        self.dataset["print_flows"] = print_flows
-        self.dataset["save_flows"] = save_flows
-        self.dataset["observations"] = observations
-        self.dataset["repeat_stress"] = repeat_stress
+        dict_dataset = {
+            "head": head,
+            "conductance": conductance,
+            "concentration": concentration,
+            "concentration_boundary_type": concentration_boundary_type,
+            "print_input": print_input,
+            "print_flows": print_flows,
+            "save_flows": save_flows,
+            "observations": observations,
+            "repeat_stress": repeat_stress,
+        }
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def _validate(self, schemata, **kwargs):

imod/mf6/gwfgwf.py

@@ -34,16 +34,18 @@ def __init__(
         angldegx: Optional[xr.DataArray] = None,
         cdist: Optional[xr.DataArray] = None,
     ):
-        super().__init__(locals())
-        self.dataset["cell_id1"] = cell_id1
-        self.dataset["cell_id2"] = cell_id2
-        self.dataset["layer"] = layer
-        self.dataset["model_name_1"] = model_id1
-        self.dataset["model_name_2"] = model_id2
-        self.dataset["ihc"] = xr.DataArray(np.ones_like(cl1, dtype=int))
-        self.dataset["cl1"] = cl1
-        self.dataset["cl2"] = cl2
-        self.dataset["hwva"] = hwva
+        dict_dataset = {
+            "cell_id1": cell_id1,
+            "cell_id2": cell_id2,
+            "layer": layer,
+            "model_name_1": model_id1,
+            "model_name_2": model_id2,
+            "ihc": xr.DataArray(np.ones_like(cl1, dtype=int)),
+            "cl1": cl1,
+            "cl2": cl2,
+            "hwva": hwva,
+        }
+        super().__init__(dict_dataset)
 
         auxiliary_variables = [var for var in [angldegx, cdist] if var is not None]
         if auxiliary_variables:

imod/mf6/gwfgwt.py

@@ -14,9 +14,12 @@ class GWFGWT(ExchangeBase):
     _template = Package._initialize_template(_pkg_id)
 
     def __init__(self, model_id1: str, model_id2: str):
-        super().__init__(locals())
-        self.dataset["model_name_1"] = model_id1
-        self.dataset["model_name_2"] = model_id2
+        dict_dataset = {
+            "model_name_1": model_id1,
+            "model_name_2": model_id2,
+        }
+
+        super().__init__(dict_dataset)
 
     def clip_box(
         self,

imod/mf6/hfb.py

@@ -280,8 +280,8 @@ def __init__(
         geometry: gpd.GeoDataFrame,
         print_input: bool = False,
     ) -> None:
-        super().__init__(locals())
-        self.dataset["print_input"] = print_input
+        dict_dataset = {"print_input": print_input}
+        super().__init__(dict_dataset)
 
         self.line_data = geometry
 

imod/mf6/ic.py

@@ -67,7 +67,6 @@ class InitialConditions(Package):
     }
 
     def __init__(self, start=None, head=None, validate: bool = True):
-        super().__init__(locals())
         if start is None:
             start = head
             warnings.warn(
@@ -80,7 +79,8 @@ def __init__(self, start=None, head=None, validate: bool = True):
             if head is not None:
                 raise ValueError("start and head arguments cannot both be defined")
 
-        self.dataset["start"] = start
+        dict_dataset = {"start": start}
+        super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
     def render(self, directory, pkgname, globaltimes, binary):