Issue #921 Update uzf (#867)

Fixes #921

# Description
This branch contains all fixes that I needed to work with the
UZF-package. The branch contains:

1- Changes to UZF- package methods
2- Changes to the open_cbc methods to deal with cbc-files of advanced
packages
3- Changes to the UZF-tests + all tests that deal with budget output

There are some breaking changes in budget output names:
1- The new names support now support output of multiple equal boundary
packages, since the package name is used.
2- The budget arrays now contain nans if a package is not active, since
zero is a valid budget output.

# 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``
- [x] Unit tests were added
- [ ] **If feature added**: Added/extended example

docs/api/changelog.rst

@@ -32,6 +32,14 @@ Added
 - Added validation for ``linear_acceleration``, ``rclose_option``,
   ``scaling_method``, ``reordering_method``, ``print_option`` and ``no_ptc``
   entries in :class:`imod.mf6.Solution`.
+- :func:`imod.mf6.open_cbc` now returns arrays which contain np.nan for cells where 
+  budget variables are not defined. Based on new budget output a disquisition between 
+  active cells but zero flow and inactive cells can be made.
+- :func:`imod.mf6.open_cbc` now returns package type in return budget names. New format 
+  is "package type"-"optional package variable"_"package name". E.g. a River package 
+  named ``primary-sys`` will get a budget name ``riv_primary-sys``. An UZF package 
+  with name ``uzf-sys1`` will get a budget name ``uzf-gwrch_uzf-sys1`` for the 
+  groundwater recharge budget from the UZF-CBC.
 
 Fixed
 ~~~~~
@@ -42,7 +50,7 @@ Fixed
 - When importing data from a .prj file, the multipliers and additions specified for
   ipf and idf files are now applied
 - Fix bug where y-coords were flipped in :class:`imod.msw.MeteoMapping`
-  
+
 Changed
 ~~~~~~~
 - Replaced csv_output by outer_csvfile and inner_csvfile in

imod/mf6/out/cbc.py

@@ -129,10 +129,17 @@ def read_cbc_headers(
                 imeth6_header = read_imeth6_header(f)
                 datasize = imeth6_header["nlist"] * (8 + imeth6_header["ndat"] * 8)
                 header["pos"] = f.tell()
-                key = imeth6_header["txt2id2"]
+                # key-format: "package type"-"optional_package_variable"_"package name"
+                # for river output: riv_sys1
+                # for uzf output: uzf-gwrch_uzf_sys1
+                key = header["text"] + "_" + imeth6_header["txt2id2"]
                 # npf-key can be present multiple times in cases of saved saturation + specific discharge
                 if header["text"].startswith("data-"):
-                    key = key + "-" + header["text"].replace("data-", "")
+                    key = (
+                        imeth6_header["txt2id2"]
+                        + "_"
+                        + header["text"].replace("data-", "")
+                    )
                 headers[key].append(Imeth6Header(**header, **imeth6_header))
             else:
                 raise ValueError(
@@ -261,6 +268,7 @@ def read_imeth6_budgets_dense(
     size: int,
     shape: tuple,
     return_variable: str,
+    indices: np.ndarray | None,
 ) -> FloatArray:
     """
     Read the data for an imeth==6 budget section.
@@ -282,14 +290,19 @@ def read_imeth6_budgets_dense(
         size of the entire model domain
     shape: tuple[int, int, int]
         Shape (nlayer, nrow, ncolumn) of entire model domain.
+    return_variable: str
+        variable name to return from budget table
+    indices: np.ndarray | None
+        optional array that contains the indices to map return_variable to model topology
 
     Returns
     -------
     Three-dimensional array of floats
     """
     # Allocates a dense array for the entire domain
-    out = np.zeros(size, dtype=np.float64)
+    out = np.full(size, np.nan, dtype=np.float64)
     table = read_imeth6_budgets(cbc_path, count, dtype, pos)
-    id1 = table["id1"] - 1  # Convert to 0 based index
-    out[id1] = table[return_variable]
+    if indices is None:
+        indices = table["id1"] - 1  # Convert to 0 based index
+    out[indices] = table[return_variable]
     return out.reshape(shape)

imod/mf6/out/common.py

@@ -1,5 +1,5 @@
 import pathlib
-from typing import BinaryIO, Union
+from typing import Any, BinaryIO, Dict, List, Union
 
 import numpy as np
 
@@ -19,3 +19,13 @@ def _to_nan(a: FloatArray, dry_nan: bool) -> FloatArray:
     if dry_nan:
         a[a == -1e30] = np.nan
     return a
+
+
+def get_first_header_advanced_package(
+    headers: Dict[str, List[Any]],
+) -> Any:
+    for key, header_list in headers.items():
+        # multimodels have a gwf-gwf budget for flow-ja-face between domains
+        if "flow-ja-face" not in key and "gwf_" in key:
+            return header_list[0]
+    return None

imod/mf6/out/dis.py

@@ -11,7 +11,13 @@
 from imod.mf6.utilities.dataset import assign_datetime_coords
 
 from . import cbc
-from .common import FilePath, FloatArray, IntArray, _to_nan
+from .common import (
+    FilePath,
+    FloatArray,
+    IntArray,
+    _to_nan,
+    get_first_header_advanced_package,
+)
 
 
 # Binary Grid File / DIS Grids
@@ -197,6 +203,7 @@ def open_imeth6_budgets(
     grb_content: dict,
     header_list: List[cbc.Imeth6Header],
     return_variable: str = "budget",
+    indices: np.ndarray | None = None,
 ) -> xr.DataArray:
     """
     Open the data for an imeth==6 budget section.
@@ -210,6 +217,8 @@ def open_imeth6_budgets(
     cbc_path: str, pathlib.Path
     grb_content: dict
     header_list: List[Imeth1Header]
+    return_variable: str
+    return_id: np.ndarray | None
 
     Returns
     -------
@@ -227,7 +236,14 @@ def open_imeth6_budgets(
     for i, header in enumerate(header_list):
         time[i] = header.totim
         a = dask.delayed(cbc.read_imeth6_budgets_dense)(
-            cbc_path, header.nlist, dtype, header.pos, size, shape, return_variable
+            cbc_path,
+            header.nlist,
+            dtype,
+            header.pos,
+            size,
+            shape,
+            return_variable,
+            indices,
         )
         x = dask.array.from_delayed(a, shape=shape, dtype=np.float64)
         dask_list.append(x)
@@ -410,10 +426,23 @@ def open_cbc(
     time_unit: Optional[str] = "d",
 ) -> Dict[str, xr.DataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
+    indices = None
+    header_advanced_package = get_first_header_advanced_package(headers)
+    if header_advanced_package is not None:
+        # For advanced packages the id2 column of variable gwf contains the MF6 id's.
+        # Get id's eager from first stress period.
+        dtype = np.dtype(
+            [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
+            + [(name, np.float64) for name in header_advanced_package.auxtxt]
+        )
+        table = cbc.read_imeth6_budgets(
+            cbc_path, header_advanced_package.nlist, dtype, header_advanced_package.pos
+        )
+        indices = table["id2"] - 1  # Convert to 0 based index
     cbc_content = {}
     for key, header_list in headers.items():
         # TODO: validate homogeneity of header_list, ndat consistent, nlist consistent etc.
-        if key == "flow-ja-face":
+        if key == "flow-ja-face" and isinstance(header_list[0], cbc.Imeth1Header):
             if flowja:
                 flowja, nm = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content
@@ -438,11 +467,15 @@ def open_cbc(
                     for return_variable in header_list[0].auxtxt:
                         key_aux = header_list[0].txt2id1 + "-" + return_variable
                         cbc_content[key_aux] = open_imeth6_budgets(
-                            cbc_path, grb_content, header_list, return_variable
+                            cbc_path,
+                            grb_content,
+                            header_list,
+                            return_variable=return_variable,
+                            indices=indices,
                         )
                 else:
                     cbc_content[key] = open_imeth6_budgets(
-                        cbc_path, grb_content, header_list
+                        cbc_path, grb_content, header_list, indices=indices
                     )
     if simulation_start_time is not None:
         for cbc_name, cbc_array in cbc_content.items():

imod/mf6/out/disv.py

@@ -12,7 +12,13 @@
 from imod.mf6.utilities.dataset import assign_datetime_coords
 
 from . import cbc
-from .common import FilePath, FloatArray, IntArray, _to_nan
+from .common import (
+    FilePath,
+    FloatArray,
+    IntArray,
+    _to_nan,
+    get_first_header_advanced_package,
+)
 
 
 def _ugrid_iavert_javert(
@@ -229,6 +235,7 @@ def open_imeth6_budgets(
     grb_content: dict,
     header_list: List[cbc.Imeth6Header],
     return_variable: str = "budget",
+    return_id: np.ndarray | None = None,
 ) -> xu.UgridDataArray:
     """
     Open the data for an imeth==6 budget section.
@@ -242,6 +249,8 @@ def open_imeth6_budgets(
     cbc_path: str, pathlib.Path
     grb_content: dict
     header_list: List[Imeth1Header]
+    return_variable: str
+    return_id: np.ndarray | None
 
     Returns
     -------
@@ -259,7 +268,14 @@ def open_imeth6_budgets(
     for i, header in enumerate(header_list):
         time[i] = header.totim
         a = dask.delayed(cbc.read_imeth6_budgets_dense)(
-            cbc_path, header.nlist, dtype, header.pos, size, shape, return_variable
+            cbc_path,
+            header.nlist,
+            dtype,
+            header.pos,
+            size,
+            shape,
+            return_variable,
+            return_id,
         )
         x = dask.array.from_delayed(a, shape=shape, dtype=np.float64)
         dask_list.append(x)
@@ -465,9 +481,22 @@ def open_cbc(
     time_unit: Optional[str] = "d",
 ) -> Dict[str, xu.UgridDataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
+    indices = None
+    header_advanced_package = get_first_header_advanced_package(headers)
+    if header_advanced_package is not None:
+        # For advanced packages the id2 column of variable gwf contains the MF6 id's.
+        # Get id's eager from first stress period.
+        dtype = np.dtype(
+            [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
+            + [(name, np.float64) for name in header_advanced_package.auxtxt]
+        )
+        table = cbc.read_imeth6_budgets(
+            cbc_path, header_advanced_package.nlist, dtype, header_advanced_package.pos
+        )
+        indices = table["id2"] - 1  # Convert to 0 based index
     cbc_content = {}
     for key, header_list in headers.items():
-        if key == "flow-ja-face":
+        if key == "flow-ja-face" and isinstance(header_list[0], cbc.Imeth1Header):
             if flowja:
                 flowja, ij = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content
@@ -490,11 +519,15 @@ def open_cbc(
                 for return_variable in header_list[0].auxtxt:
                     key_aux = header_list[0].txt2id1 + "-" + return_variable
                     cbc_content[key_aux] = open_imeth6_budgets(
-                        cbc_path, grb_content, header_list, return_variable
+                        cbc_path,
+                        grb_content,
+                        header_list,
+                        return_variable=return_variable,
+                        return_id=indices,
                     )
             else:
                 cbc_content[key] = open_imeth6_budgets(
-                    cbc_path, grb_content, header_list
+                    cbc_path, grb_content, header_list, return_id=indices
                 )
 
     if simulation_start_time is not None:

imod/mf6/simulation.py

@@ -630,7 +630,7 @@ def _merge_and_assign_exchange_budgets(self, cbc: GridDataset) -> GridDataset:
         exchange_budgets = cbc[exchange_names].to_array().sum(dim="variable")
         cbc = cbc.drop_vars(exchange_names)
         # "gwf-gwf" or "gwt-gwt"
-        exchange_key = exchange_names[0].split("_")[0]
+        exchange_key = exchange_names[0].split("_")[1]
         cbc[exchange_key] = exchange_budgets
         return cbc
 

imod/mf6/uzf.py

@@ -4,6 +4,7 @@
 from imod.logging import init_log_decorator
 from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
+from imod.prepare.layer import get_upper_active_grid_cells
 from imod.schemata import (
     AllInsideNoDataSchema,
     AllNoDataSchema,
@@ -92,6 +93,10 @@ class UnsaturatedZoneFlow(AdvancedBoundaryCondition):
         keyword to indicate that UZF flow terms will be written to the file specified with "BUDGET
         FILEOUT" in Output Control.
         Default is False.
+    budget_fileout: ({"str"}, optional)
+        path to output cbc-file for UZF budgets
+    budgetcsv_fileout: ({"str"}, optional)
+        path to output csv-file for summed budgets
     observations: [Not yet supported.]
         Default is None.
     water_mover: [Not yet supported.]
@@ -183,12 +188,12 @@ class UnsaturatedZoneFlow(AdvancedBoundaryCondition):
         "theta_sat": [IdentityNoDataSchema("kv_sat"), AllValueSchema(">=", 0.0)],
         "theta_init": [IdentityNoDataSchema("kv_sat"), AllValueSchema(">=", 0.0)],
         "epsilon": [IdentityNoDataSchema("kv_sat")],
-        "infiltration_rate": [IdentityNoDataSchema("kv_sat")],
-        "et_pot": [IdentityNoDataSchema("kv_sat")],
-        "extinction_depth": [IdentityNoDataSchema("kv_sat")],
-        "extinction_theta": [IdentityNoDataSchema("kv_sat")],
-        "root_potential": [IdentityNoDataSchema("kv_sat")],
-        "root_activity": [IdentityNoDataSchema("kv_sat")],
+        "infiltration_rate": [IdentityNoDataSchema("stress_period_active")],
+        "et_pot": [IdentityNoDataSchema("stress_period_active")],
+        "extinction_depth": [IdentityNoDataSchema("stress_period_active")],
+        "extinction_theta": [IdentityNoDataSchema("stress_period_active")],
+        "root_potential": [IdentityNoDataSchema("stress_period_active")],
+        "root_activity": [IdentityNoDataSchema("stress_period_active")],
     }
 
     _package_data = (
@@ -226,6 +231,8 @@ def __init__(
         print_input=False,
         print_flows=False,
         save_flows=False,
+        budget_fileout=None,
+        budgetcsv_fileout=None,
         observations=None,
         water_mover=None,
         timeseries=None,
@@ -234,6 +241,7 @@ def __init__(
         landflag = self._determine_landflag(kv_sat)
         iuzno = self._create_uzf_numbers(landflag)
         ivertcon = self._determine_vertical_connection(iuzno)
+        stress_period_active = landflag.where(landflag == 1)
 
         dict_dataset = {
             # Package data
@@ -244,6 +252,7 @@ def __init__(
             "theta_init": theta_init,
             "epsilon": epsilon,
             # Stress period data
+            "stress_period_active": stress_period_active,
             "infiltration_rate": infiltration_rate,
             "et_pot": et_pot,
             "extinction_depth": extinction_depth,
@@ -260,6 +269,8 @@ def __init__(
             "print_input": print_input,
             "print_flows": print_flows,
             "save_flows": save_flows,
+            "budget_fileout": budget_fileout,
+            "budgetcsv_fileout": budgetcsv_fileout,
             "observations": observations,
             "water_mover": water_mover,
             "timeseries": timeseries,
@@ -343,16 +354,19 @@ def _check_options(
 
     def _create_uzf_numbers(self, landflag):
         """Create unique UZF ID's. Inactive cells equal 0"""
-        return np.cumsum(np.ravel(landflag)).reshape(landflag.shape) * landflag
+        active_nodes = landflag.notnull().astype(np.int8)
+        return np.nancumsum(active_nodes).reshape(landflag.shape) * active_nodes
 
     def _determine_landflag(self, kv_sat):
-        return (np.isfinite(kv_sat)).astype(np.int32)
+        """returns the landflag for uzf-model. Landflag == 1 for top active UZF-nodes"""
+        land_nodes = get_upper_active_grid_cells(kv_sat).astype(np.int32)
+        return land_nodes.where(kv_sat.notnull())
 
     def _determine_vertical_connection(self, uzf_number):
         return uzf_number.shift(layer=-1, fill_value=0)
 
     def _package_data_to_sparse(self):
-        notnull = self.dataset["landflag"].values == 1
+        notnull = self.dataset["landflag"].notnull().to_numpy()
         iuzno = self.dataset["iuzno"].values[notnull]
         landflag = self.dataset["landflag"].values[notnull]
         ivertcon = self.dataset["ivertcon"].values[notnull]
@@ -392,7 +406,8 @@ def render(self, directory, pkgname, globaltimes, binary):
         d = self._get_options(d, not_options=not_options)
         path = directory / pkgname / f"{self._pkg_id}-pkgdata.dat"
         d["packagedata"] = path.as_posix()
-        d["nuzfcells"] = self._max_active_n()
+        # max uzf-cells for which time period data will be supplied
+        d["nuzfcells"] = np.count_nonzero(np.isfinite(d["landflag"]))
         return self._template.render(d)
 
     def _to_struct_array(self, arrdict, layer):
@@ -422,6 +437,7 @@ def _to_struct_array(self, arrdict, layer):
     def _validate(self, schemata, **kwargs):
         # Insert additional kwargs
         kwargs["kv_sat"] = self["kv_sat"]
+        kwargs["stress_period_active"] = self["stress_period_active"]
         errors = super()._validate(schemata, **kwargs)
 
         return errors