Issue #1743 fix clip_box with transient states_at_boundary

docs/api/changelog.rst

@@ -15,6 +15,21 @@ Fixed
 - Fixed bug where :class:`imod.mf6.Evapotranspiration` package would write files
   to binary, which could not be parsed by MODFLOW 6 when ``proportion_depth``
   and ``proportion_rate`` were provided without segments.
+- Bug where :class:`imod.mf6.ConstantConcentration` package could not be written
+  for multiple timesteps.
+- Bug where :meth:`imod.mf6.Modflow6Simulation.clip_box` where a ValidationError
+  was thrown when clipping a model with a :class:`imod.mf6.ConstantHead` or
+  :class:`imod.mf6.ConstantConcentration` package with a ``time`` dimension and
+  providing ``states_for_boundary``.
+- Bug where :meth:`imod.mf6.Modflow6Simulation.clip_box` would drop layers if
+  ``states_for_boundary`` were provided and the model already contained a
+  :class:`imod.mf6.ConstantHead` or :class:`imod.mf6.ConstantConcentration` with
+  less layers.
+- Bug where :meth:`imod.mf6.Modflow6Simulation.clip_box` would not properly
+  align timesteps and forward fill data if ``states_for_boundary`` were provided
+  and the model already contained a :class:`imod.mf6.ConstantHead` or
+  :class:`imod.mf6.ConstantConcentration`, both with timesteps, which were
+  unaligned.
 
 Changed
 ~~~~~~~

imod/mf6/model.py

@@ -40,6 +40,7 @@
 from imod.mf6.riv import River
 from imod.mf6.utilities.clipped_bc_creator import (
     StateClassType,
+    StateType,
     create_clipped_boundary,
 )
 from imod.mf6.utilities.mf6hfb import merge_hfb_packages
@@ -64,8 +65,8 @@ def pkg_has_cleanup(pkg: Package):
 def _create_boundary_condition_for_unassigned_boundary(
     model: Modflow6Model,
     state_for_boundary: Optional[GridDataArray],
-    additional_boundaries: Optional[list[StateClassType]] = None,
-):
+    additional_boundaries: list[Optional[StateType]] = [None],
+) -> Optional[StateType]:
     if state_for_boundary is None:
         return None
 
@@ -90,7 +91,7 @@ def _create_boundary_condition_clipped_boundary(
     clipped_model: Modflow6Model,
     state_for_boundary: Optional[GridDataArray],
     clip_box_args: tuple,
-):
+) -> Optional[StateType]:
     # Create temporary boundary condition for the original model boundary. This
     # is used later to see which boundaries can be ignored as they were already
     # present in the original model. We want to just end up with the boundary
@@ -120,10 +121,30 @@ def _create_boundary_condition_clipped_boundary(
     else:
         state_for_boundary_clipped = None
 
-    return _create_boundary_condition_for_unassigned_boundary(
+    bc_constant_pkg = _create_boundary_condition_for_unassigned_boundary(
         clipped_model, state_for_boundary_clipped, [unassigned_boundary_clipped]
     )
 
+    # Remove all indices before first timestep of state_for_clipped_boundary.
+    # This to prevent empty dataarrays unnecessarily being made for these
+    # indices, which can lead to them to be removed when purging empty packages
+    # with ignore_time=True. Unfortunately, this is needs to be handled here and
+    # not in _create_boundary_condition_for_unassigned_boundary, as otherwise
+    # this function is called twice which could result in broadcasting errors in
+    # the second call if the time domain of state_for_boundary and assigned
+    # packages have no overlap.
+    if (
+        (state_for_boundary is not None)
+        and (state_for_boundary.indexes.get("time") is not None)
+        and (bc_constant_pkg is not None)
+    ):
+        start_time = state_for_boundary.indexes["time"][0]
+        bc_constant_pkg.dataset = bc_constant_pkg.dataset.sel(
+            time=slice(start_time, None)
+        )
+
+    return bc_constant_pkg
+
 
 class Modflow6Model(collections.UserDict, IModel, abc.ABC):
     _mandatory_packages: tuple[str, ...] = ()

imod/mf6/utilities/clipped_bc_creator.py

@@ -1,8 +1,11 @@
-from typing import TypeAlias
+from typing import Optional, Tuple, TypeAlias
+
+import xarray as xr
 
 from imod.mf6 import ConstantConcentration, ConstantHead
 from imod.select.grid import active_grid_boundary_xy
 from imod.typing import GridDataArray
+from imod.util.dims import enforced_dim_order
 
 StateType: TypeAlias = ConstantHead | ConstantConcentration
 StateClassType: TypeAlias = type[ConstantHead] | type[ConstantConcentration]
@@ -12,7 +15,7 @@ def _find_unassigned_grid_boundaries(
     active_grid_boundary: GridDataArray,
     boundary_conditions: list[StateType],
 ) -> GridDataArray:
-    unassigned_grid_boundaries = active_grid_boundary
+    unassigned_grid_boundaries = active_grid_boundary.copy()
     for boundary_condition in boundary_conditions:
         # Fetch variable name from the first boundary condition, can be "head" or
         # "concentration".
@@ -24,6 +27,96 @@ def _find_unassigned_grid_boundaries(
     return unassigned_grid_boundaries
 
 
+def _align_time_indexes_boundaries(
+    state_for_clipped_boundary: GridDataArray,
+    unassigned_grid_boundaries: GridDataArray,
+) -> Optional[GridDataArray]:
+    """
+    Create an outer time index for aligning boundaries. Furthermore deal with
+    cases where one or both boundaries don't have a time dimension. In a graphic
+    way, we want this:
+
+    State
+       a-----b-----c
+    Unassigned
+    d-----e
+
+    Needs to align to:
+    d---a--e--b-----c
+    """
+    index_unassigned = unassigned_grid_boundaries.indexes
+    index_state = state_for_clipped_boundary.indexes
+    if "time" in index_unassigned and "time" in index_state:
+        return index_unassigned["time"].join(index_state["time"], how="outer")
+    elif "time" in index_unassigned:
+        return index_unassigned["time"]
+    elif "time" in index_state:
+        return index_state["time"]
+    else:
+        return None
+
+
+def _align_boundaries(
+    state_for_clipped_boundary: GridDataArray,
+    unassigned_grid_boundaries: Optional[GridDataArray],
+) -> Tuple[GridDataArray, Optional[GridDataArray]]:
+    """
+    Customly align the state_for_clipped_boundary and unassigned grid boundaries.
+    - "layer" dimension requires outer alignment
+    - "time" requires reindexing with ffill
+    - planar coordinates are expected to be aligned already
+    """
+    # Align dimensions
+    if unassigned_grid_boundaries is not None:
+        # Align along layer dimension with outer join, xarray API only supports
+        # excluding dims, not specifying dims to align along, so we have to do
+        # it this way.
+        dims_to_exclude = set(state_for_clipped_boundary.dims) | set(
+            unassigned_grid_boundaries.dims
+        )
+        dims_to_exclude.remove("layer")
+        state_for_clipped_boundary, unassigned_grid_boundaries = xr.align(
+            state_for_clipped_boundary,
+            unassigned_grid_boundaries,
+            join="outer",
+            exclude=dims_to_exclude,
+        )
+        # Align along time dimension by finding the outer time indexes and then
+        # reindexing with a ffill.
+        outer_time_index = _align_time_indexes_boundaries(
+            state_for_clipped_boundary, unassigned_grid_boundaries
+        )
+        if "time" in state_for_clipped_boundary.indexes:
+            state_for_clipped_boundary = state_for_clipped_boundary.reindex(
+                {"time": outer_time_index}, method="ffill"
+            )
+        if "time" in unassigned_grid_boundaries.indexes:
+            unassigned_grid_boundaries = unassigned_grid_boundaries.reindex(
+                {"time": outer_time_index}, method="ffill"
+            )
+
+    return state_for_clipped_boundary, unassigned_grid_boundaries
+
+
+@enforced_dim_order
+def _create_clipped_boundary_state(
+    idomain: GridDataArray,
+    state_for_clipped_boundary: GridDataArray,
+    original_constant_head_boundaries: list[StateType],
+):
+    """Helper function to make sure dimension order is enforced"""
+    active_grid_boundary = active_grid_boundary_xy(idomain > 0)
+    unassigned_grid_boundaries = _find_unassigned_grid_boundaries(
+        active_grid_boundary, original_constant_head_boundaries
+    )
+
+    state_for_clipped_boundary, unassigned_grid_boundaries = _align_boundaries(
+        state_for_clipped_boundary, unassigned_grid_boundaries
+    )
+
+    return state_for_clipped_boundary.where(unassigned_grid_boundaries)
+
+
 def create_clipped_boundary(
     idomain: GridDataArray,
     state_for_clipped_boundary: GridDataArray,
@@ -52,10 +145,10 @@ def create_clipped_boundary(
         packages
 
     """
-    active_grid_boundary = active_grid_boundary_xy(idomain > 0)
-    unassigned_grid_boundaries = _find_unassigned_grid_boundaries(
-        active_grid_boundary, original_constant_head_boundaries
+    constant_state = _create_clipped_boundary_state(
+        idomain,
+        state_for_clipped_boundary,
+        original_constant_head_boundaries,
     )
-    constant_state = state_for_clipped_boundary.where(unassigned_grid_boundaries)
 
     return pkg_type(constant_state, print_input=True, print_flows=True, save_flows=True)

imod/templates/mf6/gwf-cnc.j2

@@ -25,4 +25,5 @@ end dimensions
 
 {% for i, path in periods.items() %}begin period {{i}}
   open/close {{path}}{% if binary %} (binary){% endif %}
-end period{% endfor %}
+end period
+{% endfor %}

imod/tests/test_mf6/test_circle.py

@@ -205,15 +205,30 @@ def test_simulation_write_and_run_transport_vsc(circle_model_transport_vsc, tmp_
     assert concentration.shape == (52, 2, 216)
 
 
+def run_simulation(simulation, modeldir):
+    idomain = simulation["GWF_1"]["disv"]["idomain"].compute()
+    time = simulation["time_discretization"].dataset.coords["time"].values
+    simulation.write(modeldir)
+    simulation.run()
+    head = (
+        simulation.open_head(simulation_start_time=time[0])
+        .compute()
+        .reindex_like(idomain)
+    )
+    concentration = (
+        simulation.open_concentration()
+        .compute(simulation_start_time=time[0])
+        .reindex_like(idomain)
+    )
+    return head, concentration
+
+
 def test_simulation_clip_and_state_at_boundary(circle_model_transport, tmp_path):
     # Arrange
     simulation = circle_model_transport
     idomain = simulation["GWF_1"]["disv"]["idomain"].compute()
 
-    simulation.write(tmp_path / "full")
-    simulation.run()
-    head = simulation.open_head().compute().reindex_like(idomain)
-    concentration = simulation.open_concentration().compute().reindex_like(idomain)
+    head, concentration = run_simulation(simulation, tmp_path / "full")
 
     states_for_boundary = {
         "GWF_1": head.isel(time=-1, drop=True),
@@ -233,13 +248,99 @@ def test_simulation_clip_and_state_at_boundary(circle_model_transport, tmp_path)
         == 20
     )
     assert half_simulation["GWF_1"]["chd_clipped"]["head"].notnull().sum() == 20
+
+    # Test if model runs
+    head_half, concentration_half = run_simulation(half_simulation, tmp_path / "half")
+    assert head_half.shape == (52, 2, 108)
+    assert concentration_half.shape == (52, 2, 108)
+
+
+def test_simulation_clip_and_constant_state_at_boundary__transient_chd(
+    circle_model_transport, tmp_path
+):
+    # Arrange
+    simulation = circle_model_transport
+    idomain = simulation["GWF_1"]["disv"]["idomain"].compute()
+    time = simulation["time_discretization"].dataset.coords["time"].values
+
+    simulation["GWF_1"]["chd"].dataset["head"] = (
+        simulation["GWF_1"]["chd"].dataset["head"].expand_dims(time=time)
+    )
+
+    head, concentration = run_simulation(simulation, tmp_path / "full")
+
+    states_for_boundary = {
+        "GWF_1": head.isel(time=-1, drop=True),
+        "transport": concentration.isel(time=-1, drop=True),
+    }
+    # Act
+    half_simulation = simulation.clip_box(
+        x_max=0.1, states_for_boundary=states_for_boundary
+    )
+    # Assert
+    # Test if model dims halved
+    idomain_half = half_simulation["GWF_1"]["disv"]["idomain"]
+    dim = idomain_half.grid.face_dimension
+    np.testing.assert_array_equal(idomain_half.sizes[dim] / idomain.sizes[dim], 0.5)
+    # Test if the clipped boundary conditions have the correct dimension and values
+    conc_clipped = half_simulation["transport"]["cnc_clipped"]["concentration"]
+    assert conc_clipped.sizes["layer"] == 2
+    assert "time" not in conc_clipped.dims
+    assert conc_clipped.notnull().sum() == 20
+    head_clipped = half_simulation["GWF_1"]["chd_clipped"]["head"]
+    assert head_clipped.sizes["layer"] == 2
+    assert head_clipped.sizes["time"] == 52
+    assert head_clipped.notnull().sum() == 20 * 52
     # Test if model runs
-    half_simulation.write(tmp_path / "half")
-    half_simulation.run()
-    # Test if the clipped model output has the correct dimension
-    head_half = half_simulation.open_head().compute().reindex_like(idomain_half)
-    concentration_half = (
-        half_simulation.open_concentration().compute().reindex_like(idomain_half)
+    head_half, concentration_half = run_simulation(half_simulation, tmp_path / "half")
+    assert head_half.shape == (52, 2, 108)
+    assert concentration_half.shape == (52, 2, 108)
+
+
+def test_simulation_clip_and_transient_state_at_boundary__transient_chd(
+    circle_model_transport, tmp_path
+):
+    """
+    Test with a transient chd boundary condition and transient states for
+    boundary conditions. The time steps of the chd package are outside of the
+    time domain of the states_for_boundary.
+    """
+    # Arrange
+    simulation = circle_model_transport
+    idomain = simulation["GWF_1"]["disv"]["idomain"].compute()
+    time = simulation["time_discretization"].dataset.coords["time"].values
+
+    simulation["GWF_1"]["chd"].dataset["head"] = (
+        simulation["GWF_1"]["chd"].dataset["head"].expand_dims(time=time[:5])
     )
+
+    head, concentration = run_simulation(simulation, tmp_path / "full")
+
+    states_for_boundary = {
+        "GWF_1": head.isel(time=slice(11, -11)),
+        "transport": concentration.isel(time=slice(11, -11)),
+    }
+    # Act
+    half_simulation = simulation.clip_box(
+        x_max=0.1, states_for_boundary=states_for_boundary
+    )
+    # Assert
+    # Test if model dims halved
+    idomain_half = half_simulation["GWF_1"]["disv"]["idomain"]
+    dim = idomain_half.grid.face_dimension
+    np.testing.assert_array_equal(idomain_half.sizes[dim] / idomain.sizes[dim], 0.5)
+    # Test if the clipped boundary conditions have the correct dimension and values
+    # Transport data for just 30 time steps (from 12 up to 41)
+    conc_clipped = half_simulation["transport"]["cnc_clipped"]["concentration"]
+    assert conc_clipped.sizes["layer"] == 2
+    assert conc_clipped.sizes["time"] == 30
+    assert conc_clipped.notnull().sum() == 20 * 30
+    # Head data for chd for just 30 time steps (from 12 up to 41)
+    head_clipped = half_simulation["GWF_1"]["chd_clipped"]["head"]
+    assert head_clipped.sizes["layer"] == 2
+    assert head_clipped.sizes["time"] == 30
+    assert head_clipped.notnull().sum() == 20 * 30
+    # Test if model runs
+    head_half, concentration_half = run_simulation(half_simulation, tmp_path / "half")
     assert head_half.shape == (52, 2, 108)
     assert concentration_half.shape == (52, 2, 108)