Issue #921 Update uzf

imod/mf6/out/__init__.py

@@ -227,6 +227,7 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
+    uzf: bool = False,
 ) -> Dict[str, Union[xr.DataArray, xu.UgridDataArray]]:
     """
     Open modflow6 cell-by-cell (.cbc) file.
@@ -287,4 +288,4 @@ def open_cbc(
     grb_content = read_grb(grb_path)
     distype = grb_content["distype"]
     _open = _get_function(_OPEN_CBC, distype)
-    return _open(cbc_path, grb_content, flowja, simulation_start_time, time_unit)
+    return _open(cbc_path, grb_content, flowja, simulation_start_time, time_unit, uzf)

imod/mf6/out/cbc.py

@@ -87,6 +87,7 @@ def read_imeth6_header(f: BinaryIO) -> Dict[str, Any]:
 
 def read_cbc_headers(
     cbc_path: FilePath,
+    uzf: bool = False,
 ) -> Dict[str, List[Union[Imeth1Header, Imeth6Header]]]:
     """
     Read all the header data from a cell-by-cell (.cbc) budget file.
@@ -129,10 +130,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"]
-                # 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-", "")
+                if uzf:
+                    # uzf package budget file
+                    key = header["text"]
+                else:
+                    key = imeth6_header["txt2id2"]
+                    # npf-key can be present multiple times in cases of saved saturation + specific discharge
+                    if header["text"].startswith("data-"):
+                        key = imeth6_header["txt2id2"] + "-" + header["text"].replace("data-", "")
+                    # uzf-key can be present multiple times for gwd and gwrch 
+                    if 'uzf' in header["text"]:
+                        key = key + "_" + header["text"]
                 headers[key].append(Imeth6Header(**header, **imeth6_header))
             else:
                 raise ValueError(
@@ -261,6 +269,7 @@ def read_imeth6_budgets_dense(
     size: int,
     shape: tuple,
     return_variable: str,
+    return_id: np.ndarray | None
 ) -> FloatArray:
     """
     Read the data for an imeth==6 budget section.
@@ -288,8 +297,9 @@ def read_imeth6_budgets_dense(
     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 return_id is None:
+        return_id = table["id1"] - 1  # Convert to 0 based index
+    out[return_id] = table[return_variable]
     return out.reshape(shape)

imod/mf6/out/dis.py

@@ -197,6 +197,7 @@ def open_imeth6_budgets(
     grb_content: dict,
     header_list: List[cbc.Imeth6Header],
     return_variable: str = "budget",
+    return_id: np.ndarray | None = None,
 ) -> xr.DataArray:
     """
     Open the data for an imeth==6 budget section.
@@ -227,7 +228,7 @@ 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)
@@ -408,12 +409,24 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
+    advanced_package: bool = False,
 ) -> Dict[str, xr.DataArray]:
-    headers = cbc.read_cbc_headers(cbc_path)
+    headers = cbc.read_cbc_headers(cbc_path, advanced_package)
+    return_id = None
+    if advanced_package:
+        # For advanced packages the id2 column of variable gwf contains the MF6 id's.
+        # Get id's from first stress period.
+        header = headers['gwf'][0]
+        dtype = np.dtype(
+            [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
+            + [(name, np.float64) for name in header.auxtxt]
+        )
+        table = cbc.read_imeth6_budgets(cbc_path, header.nlist, dtype, header.pos)
+        return_id = 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 not advanced_package:
             if flowja:
                 flowja, nm = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content
@@ -438,11 +451,11 @@ 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 = return_id
                         )
                 else:
                     cbc_content[key] = open_imeth6_budgets(
-                        cbc_path, grb_content, header_list
+                        cbc_path, grb_content, header_list, return_id = return_id
                     )
     if simulation_start_time is not None:
         for cbc_name, cbc_array in cbc_content.items():

imod/mf6/out/disv.py

@@ -229,6 +229,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.
@@ -259,7 +260,7 @@ 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)
@@ -463,11 +464,23 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
+    advanced_package: bool = False,
 ) -> Dict[str, xu.UgridDataArray]:
-    headers = cbc.read_cbc_headers(cbc_path)
+    headers = cbc.read_cbc_headers(cbc_path, advanced_package)
+    # For advanced packages the id2 column of variable gwf contains the MF6 id's.
+    # Get id's eager from first stress period.
+    return_id = None
+    if advanced_package:
+        header = headers['gwf'][0]
+        dtype = np.dtype(
+            [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
+            + [(name, np.float64) for name in header.auxtxt]
+        )
+        table = cbc.read_imeth6_budgets(cbc_path, header.nlist, dtype, header.pos)
+        return_id = 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 not advanced_package:
             if flowja:
                 flowja, ij = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content
@@ -490,11 +503,11 @@ 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 = return_id
                     )
             else:
                 cbc_content[key] = open_imeth6_budgets(
-                    cbc_path, grb_content, header_list
+                    cbc_path, grb_content, header_list, return_id = return_id
                 )
 
     if simulation_start_time is not None:

imod/mf6/uzf.py

@@ -91,6 +91,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.]
@@ -224,6 +228,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,
@@ -258,6 +264,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,
@@ -350,7 +358,7 @@ 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 = np.isfinite(self.dataset["landflag"].values)
         iuzno = self.dataset["iuzno"].values[notnull]
         landflag = self.dataset["landflag"].values[notnull]
         ivertcon = self.dataset["ivertcon"].values[notnull]
@@ -376,7 +384,7 @@ def _package_data_to_sparse(self):
         recarr_new[field_names] = recarr
 
         return recarr_new
-
+    
     def render(self, directory, pkgname, globaltimes, binary):
         """Render fills in the template only, doesn't write binary data"""
         d = {}
@@ -390,7 +398,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):

imod/templates/mf6/gwf-uzf.j2

@@ -11,8 +11,10 @@ begin options
   print_flows{% endif %}
 {%- if save_flows is defined %}
   save_flows{% endif %}
-{%- if budget_filerecord is defined %}
-  budget fileout {{budgetfile}}{% endif %}
+{%- if budget_fileout is defined %}
+  budget fileout {{budget_fileout}}{% endif %}
+{%- if budgetcsv_fileout is defined %}
+  budgetcsv fileout {{budgetcsv_fileout}}{% endif %}
 {%- if ts_filerecord is defined %}
   ts6 filein {{ts6_filename}}{% endif %}
 {%- if obs_filerecord is defined %}