Issue #1313 recharge from imod5 cap data (#1315)

Fixes #1313

# Description
Construct an rch-package from iMOD5 data in the CAP package, loaded with
the ``open_projectfile_data`` function. Package is
used to couple MODFLOW6 to MetaSWAP models. Active cells will have a
recharge rate of 0.0.

At the moment, MetaSWAP can only be coupled to the first layer, as this
is also the case for ``primod``. iMOD Coupler these days supports
coupling to other layers as well, but ``primod`` doesn't. Picking this
up for iMOD Python and ``primod`` is worthy a separate story.

In detail:
- Add ``Recharge.from_imod5_cap_data`` class method, to construct an
empty Recharge package with 0.0 rate.
- Minor refactor: Put ``GridData`` helper functions to
``msw/utilities/imod5_converter.py``, so that they can be reused.

# Checklist

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

docs/api/changelog.rst

@@ -14,6 +14,8 @@ Added
 
 - :class:`imod.msw.MeteoGridCopy` to copy existing `mete_grid.inp` files, so
   ASCII grids in large existing meteo databases do not have to be read.
+- :meth:`imod.mf6.Recharge.from_imod5_cap_data` to construct a recharge package
+  for coupling a MODFLOW6 model to MetaSWAP.
 
 Fixed
 ~~~~~

imod/mf6/model_gwf.py

@@ -330,6 +330,9 @@ def from_imod5_data(
                     wel_key, imod5_data, times
                 )
 
+        if "cap" in imod5_keys:
+            result["msw-rch"] = Recharge.from_imod5_cap_data(imod5_data)  # type: ignore
+
         # import ghb's
         imod5_keys = list(imod5_data.keys())
         ghb_keys = [key for key in imod5_keys if key[0:3] == "ghb"]

imod/mf6/rch.py

@@ -1,7 +1,8 @@
 from copy import deepcopy
-from typing import Optional
+from typing import Optional, cast
 
 import numpy as np
+import xarray as xr
 
 from imod.logging import init_log_decorator
 from imod.mf6.boundary_condition import BoundaryCondition
@@ -11,6 +12,10 @@
 from imod.mf6.utilities.imod5_converter import convert_unit_rch_rate
 from imod.mf6.utilities.regrid import RegridderWeightsCache, _regrid_package_data
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
+from imod.msw.utilities.imod5_converter import (
+    get_cell_area_from_imod5_data,
+    is_msw_active_cell,
+)
 from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_rch_cells
 from imod.schemata import (
     AllInsideNoDataSchema,
@@ -23,7 +28,7 @@
     IndexesSchema,
     OtherCoordsSchema,
 )
-from imod.typing import GridDataArray
+from imod.typing import GridDataArray, GridDataDict, Imod5DataDict
 from imod.typing.grid import (
     enforce_dim_order,
     is_planar_grid,
@@ -218,7 +223,7 @@ def from_imod5_data(
             # create an array indicating in which cells rch is active
             is_rch_cell = allocate_rch_cells(
                 ALLOCATION_OPTION.at_first_active,
-                new_idomain == 1,
+                new_idomain > 0,
                 planar_rate_regridded,
             )
 
@@ -227,7 +232,33 @@ def from_imod5_data(
             rch_rate = enforce_dim_order(rch_rate)
             new_package_data["rate"] = rch_rate
 
-        return Recharge(**new_package_data, validate=True, fixed_cell=False)
+        return cls(**new_package_data, validate=True, fixed_cell=False)
+
+    @classmethod
+    def from_imod5_cap_data(
+        cls,
+        imod5_data: Imod5DataDict,
+        target_dis: StructuredDiscretization,
+    ) -> "Recharge":
+        """
+        Construct an rch-package from iMOD5 data in the CAP package, loaded with
+        the :func:`imod.formats.prj.open_projectfile_data` function. Package is
+        used to couple MODFLOW6 to MetaSWAP models. Active cells will have a
+        recharge rate of 0.0.
+        """
+        cap_data = cast(GridDataDict, imod5_data["cap"])
+
+        msw_area = get_cell_area_from_imod5_data(cap_data)
+        active, _ = is_msw_active_cell(target_dis, cap_data, msw_area)
+
+        data = {}
+        layer_da = xr.full_like(
+            target_dis.dataset.coords["layer"], np.nan, dtype=np.float64
+        )
+        layer_da.loc[{"layer": 1}] = 0.0
+        data["rate"] = layer_da.where(active)
+
+        return cls(**data, validate=True, fixed_cell=False)
 
     @classmethod
     def get_regrid_methods(cls) -> RechargeRegridMethod:

imod/msw/grid_data.py

@@ -3,7 +3,6 @@
 import numpy as np
 import xarray as xr
 
-from imod.logging import LogLevel, logger
 from imod.mf6 import StructuredDiscretization
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.mf6.regrid.regrid_schemes import (
@@ -13,65 +12,16 @@
 from imod.msw.fixed_format import VariableMetaData
 from imod.msw.pkgbase import MetaSwapPackage
 from imod.msw.regrid.regrid_schemes import GridDataRegridMethod
-from imod.msw.utilities.common import concat_imod5
-from imod.typing import GridDataArray, GridDataDict
-from imod.typing.grid import ones_like
+from imod.msw.utilities.imod5_converter import (
+    get_cell_area_from_imod5_data,
+    get_landuse_from_imod5_data,
+    get_rootzone_depth_from_imod5_data,
+    is_msw_active_cell,
+)
+from imod.typing import GridDataDict
 from imod.util.spatial import get_cell_area, spatial_reference
 
 
-def get_cell_area_from_imod5_data(
-    imod5_cap: GridDataDict,
-) -> GridDataArray:
-    # area's per type of svats
-    mf6_area = get_cell_area(imod5_cap["wetted_area"])
-    wetted_area = imod5_cap["wetted_area"]
-    urban_area = imod5_cap["urban_area"]
-    rural_area = mf6_area - (wetted_area + urban_area)
-    if (wetted_area > mf6_area).any():
-        logger.log(
-            loglevel=LogLevel.WARNING,
-            message=f"wetted area was set to the max cell area of {mf6_area}",
-            additional_depth=0,
-        )
-        wetted_area = wetted_area.where(wetted_area <= mf6_area, other=mf6_area)
-    if (rural_area < 0.0).any():
-        logger.log(
-            loglevel=LogLevel.WARNING,
-            message="found urban area > than (cel-area - wetted area). Urban area was set to 0",
-            additional_depth=0,
-        )
-        urban_area = urban_area.where(rural_area > 0.0, other=0.0)
-    rural_area = mf6_area - (wetted_area + urban_area)
-    return concat_imod5(rural_area, urban_area)
-
-
-def get_landuse_from_imod5_data(
-    imod5_cap: GridDataDict,
-) -> GridDataArray:
-    """
-    Get landuse from imod5 capillary zone data. This adds two subunits, one
-    based on the landuse grid, which specifies rural landuse. The other
-    specifies urban landuse, which is coded to value 18.
-    """
-    rural_landuse = imod5_cap["landuse"]
-    # Urban landuse = 18
-    urban_landuse = ones_like(rural_landuse) * 18
-    return concat_imod5(rural_landuse, urban_landuse)
-
-
-def get_rootzone_depth_from_imod5_data(
-    imod5_cap: GridDataDict,
-) -> GridDataArray:
-    """
-    Get rootzone depth from imod5 capillary zone data. Also does a unit
-    conversion: iMOD5 specifies rootzone thickness in centimeters, whereas
-    MetaSWAP requires rootzone depth in meters.
-    """
-    rootzone_thickness = imod5_cap["rootzone_thickness"] * 0.01
-    # rootzone depth is equal for both svats.
-    return concat_imod5(rootzone_thickness, rootzone_thickness)
-
-
 class GridData(MetaSwapPackage, IRegridPackage):
     """
     This contains the grid data of MetaSWAP.
@@ -194,11 +144,8 @@ def from_imod5_data(
         data["surface_elevation"] = imod5_cap["surface_elevation"]
         data["soil_physical_unit"] = imod5_cap["soil_physical_unit"]
 
-        mf6_top_active = target_dis["idomain"].isel(layer=0, drop=True)
-        subunit_active = (
-            (imod5_cap["boundary"] == 1) & (data["area"] > 0) & (mf6_top_active >= 1)
-        )
-        active = subunit_active.all(dim="subunit")
+        active, subunit_active = is_msw_active_cell(target_dis, imod5_cap, data["area"])
+
         data_active = {
             key: (
                 griddata.where(subunit_active)

imod/msw/utilities/imod5_converter.py

@@ -0,0 +1,87 @@
+from imod.logging import LogLevel, logger
+from imod.mf6 import StructuredDiscretization
+from imod.msw.utilities.common import concat_imod5
+from imod.typing import GridDataArray, GridDataDict
+from imod.typing.grid import ones_like
+from imod.util.spatial import get_cell_area
+
+
+def get_cell_area_from_imod5_data(
+    imod5_cap: GridDataDict,
+) -> GridDataArray:
+    # area's per type of svats
+    mf6_area = get_cell_area(imod5_cap["wetted_area"])
+    wetted_area = imod5_cap["wetted_area"]
+    urban_area = imod5_cap["urban_area"]
+    rural_area = mf6_area - (wetted_area + urban_area)
+    if (wetted_area > mf6_area).any():
+        logger.log(
+            loglevel=LogLevel.WARNING,
+            message=f"wetted area was set to the max cell area of {mf6_area}",
+            additional_depth=0,
+        )
+        wetted_area = wetted_area.where(wetted_area <= mf6_area, other=mf6_area)
+    if (rural_area < 0.0).any():
+        logger.log(
+            loglevel=LogLevel.WARNING,
+            message="found urban area > than (cel-area - wetted area). Urban area was set to 0",
+            additional_depth=0,
+        )
+        urban_area = urban_area.where(rural_area > 0.0, other=0.0)
+    rural_area = mf6_area - (wetted_area + urban_area)
+    return concat_imod5(rural_area, urban_area)
+
+
+def get_landuse_from_imod5_data(
+    imod5_cap: GridDataDict,
+) -> GridDataArray:
+    """
+    Get landuse from imod5 capillary zone data. This adds two subunits, one
+    based on the landuse grid, which specifies rural landuse. The other
+    specifies urban landuse, which is coded to value 18.
+    """
+    rural_landuse = imod5_cap["landuse"]
+    # Urban landuse = 18
+    urban_landuse = ones_like(rural_landuse) * 18
+    return concat_imod5(rural_landuse, urban_landuse)
+
+
+def get_rootzone_depth_from_imod5_data(
+    imod5_cap: GridDataDict,
+) -> GridDataArray:
+    """
+    Get rootzone depth from imod5 capillary zone data. Also does a unit
+    conversion: iMOD5 specifies rootzone thickness in centimeters, whereas
+    MetaSWAP requires rootzone depth in meters.
+    """
+    rootzone_thickness = imod5_cap["rootzone_thickness"] * 0.01
+    # rootzone depth is equal for both svats.
+    return concat_imod5(rootzone_thickness, rootzone_thickness)
+
+
+def is_msw_active_cell(
+    target_dis: StructuredDiscretization,
+    imod5_cap: GridDataDict,
+    msw_area: GridDataArray,
+) -> tuple[GridDataArray, GridDataArray]:
+    """
+    Return grid of cells that are active in the coupled computation, based on
+    following criteria:
+
+    - Active in top layer MODFLOW6
+    - Active in boundary array in CAP package
+    - MetaSWAP area > 0
+
+    Returns
+    -------
+    active: xr.DataArray
+        Active cells in any of the subunits
+    subunit_active: xr.DataArray
+        Cells active per subunit
+    """
+    mf6_top_active = target_dis["idomain"].isel(layer=0, drop=True)
+    subunit_active = (
+        (imod5_cap["boundary"] == 1) & (msw_area > 0) & (mf6_top_active >= 1)
+    )
+    active = subunit_active.any(dim="subunit")
+    return active, subunit_active

imod/tests/test_mf6/test_mf6_rch.py

@@ -455,3 +455,30 @@ def test_non_planar_rch_from_imod5_transient(imod5_dataset, tmp_path):
     )
     assert rendered_rch.count("begin period") == 3
     assert "maxbound 33856" in rendered_rch
+
+
+@pytest.mark.usefixtures("imod5_dataset")
+def test_from_imod5_cap_data(imod5_dataset):
+    # Arrange
+    data = deepcopy(imod5_dataset[0])
+    target_discretization = StructuredDiscretization.from_imod5_data(data)
+    data["cap"] = {}
+    msw_bound = data["bnd"]["ibound"].isel(layer=0, drop=True)
+    data["cap"]["boundary"] = msw_bound
+    data["cap"]["wetted_area"] = xr.ones_like(msw_bound) * 100
+    data["cap"]["urban_area"] = xr.ones_like(msw_bound) * 200
+    # Set to total cellsize, cell needs to be deactivated.
+    data["cap"]["wetted_area"][100, 100] = 625.0
+    # Set to zero, cell needs to be deactivated.
+    data["cap"]["urban_area"][100, 100] = 0.0
+    # Act
+    rch = imod.mf6.Recharge.from_imod5_cap_data(data, target_discretization)
+    rate = rch.dataset["rate"]
+    # Assert
+    np.testing.assert_array_equal(np.unique(rate), np.array([0.0, np.nan]))
+    # Boundaries inactive in MetaSWAP
+    assert np.isnan(rate[0, :, 0]).all()
+    assert np.isnan(rate[0, :, -1]).all()
+    assert np.isnan(rate[0, 0, :]).all()
+    assert np.isnan(rate[0, -1, :]).all()
+    assert np.isnan(rate[:, 100, 100]).all()