From 0787948526a8391aa1d4fedc5d16011a9b7eb757 Mon Sep 17 00:00:00 2001
From: Hendrik Kok <hendrik.kok@deltares.nl>
Date: Wed, 21 Feb 2024 10:16:01 +0100
Subject: [PATCH 01/19] Various change to UZF package for correct rendering and
 support of additional options

---
 imod/mf6/uzf.py               | 15 ++++++++++++---
 imod/templates/mf6/gwf-uzf.j2 |  6 ++++--
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/imod/mf6/uzf.py b/imod/mf6/uzf.py
index 47f032181..a528eb148 100644
--- a/imod/mf6/uzf.py
+++ b/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):
diff --git a/imod/templates/mf6/gwf-uzf.j2 b/imod/templates/mf6/gwf-uzf.j2
index e486f1bfa..13dfd28a4 100644
--- a/imod/templates/mf6/gwf-uzf.j2
+++ b/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 %}

From 695b8a24ba9e2141f2aa6510f75709ff46d21095 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Wed, 21 Feb 2024 16:46:22 +0100
Subject: [PATCH 02/19] update read_cbc_headers method to deal with uzf cbc
 file

---
 imod/mf6/out/__init__.py |  3 ++-
 imod/mf6/out/cbc.py      | 16 ++++++++++++----
 imod/mf6/out/dis.py      |  5 +++--
 imod/mf6/out/disv.py     |  5 +++--
 4 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/imod/mf6/out/__init__.py b/imod/mf6/out/__init__.py
index d65653805..63a107279 100644
--- a/imod/mf6/out/__init__.py
+++ b/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)
diff --git a/imod/mf6/out/cbc.py b/imod/mf6/out/cbc.py
index 2a278cd99..0666dc902 100644
--- a/imod/mf6/out/cbc.py
+++ b/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(
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 6831258aa..0fce2186a 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -408,12 +408,13 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
+    uzf: bool = False,
 ) -> Dict[str, xr.DataArray]:
-    headers = cbc.read_cbc_headers(cbc_path)
+    headers = cbc.read_cbc_headers(cbc_path, uzf)
     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 uzf:
             if flowja:
                 flowja, nm = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index 5a38a974c..49e1a0484 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -463,11 +463,12 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
+    uzf: bool = False,
 ) -> Dict[str, xu.UgridDataArray]:
-    headers = cbc.read_cbc_headers(cbc_path)
+    headers = cbc.read_cbc_headers(cbc_path, uzf)
     cbc_content = {}
     for key, header_list in headers.items():
-        if key == "flow-ja-face":
+        if key == "flow-ja-face" and not uzf:
             if flowja:
                 flowja, ij = cbc.open_face_budgets_as_flowja(
                     cbc_path, header_list, grb_content

From 117e6509400e55cf4c1bad4640beb05fbdc1ef04 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 23 Feb 2024 09:53:34 +0100
Subject: [PATCH 03/19] Update open_cbc method to deal with budget files of
 advanced packages like the UZF

---
 imod/mf6/out/cbc.py  |  8 +++++---
 imod/mf6/out/dis.py  | 24 ++++++++++++++++++------
 imod/mf6/out/disv.py | 24 ++++++++++++++++++------
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/imod/mf6/out/cbc.py b/imod/mf6/out/cbc.py
index 0666dc902..8804eb5f0 100644
--- a/imod/mf6/out/cbc.py
+++ b/imod/mf6/out/cbc.py
@@ -269,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.
@@ -296,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)
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 0fce2186a..761c306c7 100644
--- a/imod/mf6/out/dis.py
+++ b/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,13 +409,24 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
-    uzf: bool = False,
+    advanced_package: bool = False,
 ) -> Dict[str, xr.DataArray]:
-    headers = cbc.read_cbc_headers(cbc_path, uzf)
+    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" and not uzf:
+        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
@@ -439,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():
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index 49e1a0484..b0f04bc99 100644
--- a/imod/mf6/out/disv.py
+++ b/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,12 +464,23 @@ def open_cbc(
     flowja: bool = False,
     simulation_start_time: Optional[np.datetime64] = None,
     time_unit: Optional[str] = "d",
-    uzf: bool = False,
+    advanced_package: bool = False,
 ) -> Dict[str, xu.UgridDataArray]:
-    headers = cbc.read_cbc_headers(cbc_path, uzf)
+    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" and not uzf:
+        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
@@ -491,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:

From cdce29519a90e6889718f24c855793d598259004 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 23 Feb 2024 10:44:11 +0100
Subject: [PATCH 04/19] update open_cbc method to more general form so no extra
 input is required for advanced packages budget files

---
 imod/mf6/out/__init__.py |  1 -
 imod/mf6/out/cbc.py      | 23 +++++++++++------------
 imod/mf6/out/dis.py      |  9 +++++----
 imod/mf6/out/disv.py     | 13 +++++++------
 4 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/imod/mf6/out/__init__.py b/imod/mf6/out/__init__.py
index 63a107279..e490da96c 100644
--- a/imod/mf6/out/__init__.py
+++ b/imod/mf6/out/__init__.py
@@ -227,7 +227,6 @@ 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.
diff --git a/imod/mf6/out/cbc.py b/imod/mf6/out/cbc.py
index 8804eb5f0..78c3c5d48 100644
--- a/imod/mf6/out/cbc.py
+++ b/imod/mf6/out/cbc.py
@@ -87,7 +87,6 @@ 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.
@@ -130,17 +129,13 @@ def read_cbc_headers(
                 imeth6_header = read_imeth6_header(f)
                 datasize = imeth6_header["nlist"] * (8 + imeth6_header["ndat"] * 8)
                 header["pos"] = f.tell()
-                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"]
+                # 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 = imeth6_header["txt2id2"] + "_" + header["text"].replace("data-", "")
                 headers[key].append(Imeth6Header(**header, **imeth6_header))
             else:
                 raise ValueError(
@@ -291,6 +286,10 @@ 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 
+    return_id: np.ndarray | None
+        optional array that contains the indices to map return_variable to model topology
 
     Returns
     -------
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 761c306c7..7eddd1175 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -211,6 +211,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
     -------
@@ -409,11 +411,10 @@ 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, advanced_package)
+    headers = cbc.read_cbc_headers(cbc_path)
     return_id = None
-    if advanced_package:
+    if 'gwf' in headers.keys():
         # 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]
@@ -426,7 +427,7 @@ def open_cbc(
     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" and not advanced_package:
+        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
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index b0f04bc99..f8923ff97 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -243,6 +243,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
     -------
@@ -464,13 +466,12 @@ 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, advanced_package)
-    # For advanced packages the id2 column of variable gwf contains the MF6 id's.
-    # Get id's eager from first stress period.
+    headers = cbc.read_cbc_headers(cbc_path)
     return_id = None
-    if advanced_package:
+    if 'gwf' in headers.keys():
+        # For advanced packages the id2 column of variable gwf contains the MF6 id's.
+        # Get id's eager from first stress period.
         header = headers['gwf'][0]
         dtype = np.dtype(
             [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
@@ -480,7 +481,7 @@ def open_cbc(
         return_id = table["id2"] - 1  # Convert to 0 based index
     cbc_content = {}
     for key, header_list in headers.items():
-        if key == "flow-ja-face" and not advanced_package:
+        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

From 546f7070eba1fd1c23fecb44e3b43c631346a961 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 23 Feb 2024 11:08:56 +0100
Subject: [PATCH 05/19] review

---
 imod/mf6/uzf.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/uzf.py b/imod/mf6/uzf.py
index a528eb148..67504d46c 100644
--- a/imod/mf6/uzf.py
+++ b/imod/mf6/uzf.py
@@ -358,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 = np.isfinite(self.dataset["landflag"].values)
+        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]

From ba60742b20807f332d92f06461945dd3c076a1b5 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Mon, 4 Mar 2024 10:38:17 +0100
Subject: [PATCH 06/19] Deal with new name prefixes

---
 imod/mf6/out/__init__.py |  2 +-
 imod/mf6/out/dis.py      | 26 ++++++++++++++++----------
 imod/mf6/out/disv.py     | 12 +++++++-----
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/imod/mf6/out/__init__.py b/imod/mf6/out/__init__.py
index e490da96c..d65653805 100644
--- a/imod/mf6/out/__init__.py
+++ b/imod/mf6/out/__init__.py
@@ -287,4 +287,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, uzf)
+    return _open(cbc_path, grb_content, flowja, simulation_start_time, time_unit)
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 7eddd1175..8b24f1654 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -1,6 +1,6 @@
 import os
 import struct
-from typing import Any, BinaryIO, Dict, List, Optional, Tuple
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
 
 import dask
 import numba
@@ -403,6 +403,11 @@ def dis_open_face_budgets(
     lower = dis_extract_face_budgets(budgets, lower_index)
     return right, front, lower
 
+def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
+    for key, header in headers.items():
+        if 'gwf_' in key:
+            return header[0]
+    return None
 
 # TODO: Currently assumes dis grb, can be checked & dispatched
 def open_cbc(
@@ -413,17 +418,18 @@ def open_cbc(
     time_unit: Optional[str] = "d",
 ) -> Dict[str, xr.DataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
-    return_id = None
-    if 'gwf' in headers.keys():
+    indices = None
+    header_advanced_package = get_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 from first stress period.
-        header = headers['gwf'][0]
+        # Get id's eager from first stress period.
+        header_advanced_package = headers['gwf'][0]
         dtype = np.dtype(
             [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
-            + [(name, np.float64) for name in header.auxtxt]
+            + [(name, np.float64) for name in header_advanced_package.auxtxt]
         )
-        table = cbc.read_imeth6_budgets(cbc_path, header.nlist, dtype, header.pos)
-        return_id = table["id2"] - 1  # Convert to 0 based index
+        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.
@@ -452,11 +458,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 = return_variable, return_id = return_id
+                            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, return_id = return_id
+                        cbc_path, grb_content, header_list, return_id = indices
                     )
     if simulation_start_time is not None:
         for cbc_name, cbc_array in cbc_content.items():
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index f8923ff97..ade7b9b66 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -1,5 +1,6 @@
 import os
 import struct
+from dis import get_header_advanced_package
 from typing import Any, BinaryIO, Dict, List, Optional, Tuple
 
 import dask
@@ -468,8 +469,9 @@ def open_cbc(
     time_unit: Optional[str] = "d",
 ) -> Dict[str, xu.UgridDataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
-    return_id = None
-    if 'gwf' in headers.keys():
+    indices = None
+    header_advanced_package = get_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.
         header = headers['gwf'][0]
@@ -478,7 +480,7 @@ def open_cbc(
             + [(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
+        indices = table["id2"] - 1  # Convert to 0 based index
     cbc_content = {}
     for key, header_list in headers.items():
         if key == "flow-ja-face" and isinstance(header_list[0], cbc.Imeth1Header):
@@ -504,11 +506,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 = return_variable, return_id = return_id
+                        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, return_id = return_id
+                    cbc_path, grb_content, header_list, return_id = indices
                 )
 
     if simulation_start_time is not None:

From eb232ff33a46652f4e8e5f6f1f70b9b9dff57341 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Mon, 4 Mar 2024 11:02:04 +0100
Subject: [PATCH 07/19] fix

---
 imod/mf6/out/dis.py  | 12 ++++++------
 imod/mf6/out/disv.py | 14 +++++++++-----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 8b24f1654..9bca3393e 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -403,11 +403,6 @@ def dis_open_face_budgets(
     lower = dis_extract_face_budgets(budgets, lower_index)
     return right, front, lower
 
-def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
-    for key, header in headers.items():
-        if 'gwf_' in key:
-            return header[0]
-    return None
 
 # TODO: Currently assumes dis grb, can be checked & dispatched
 def open_cbc(
@@ -423,7 +418,6 @@ def open_cbc(
     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.
-        header_advanced_package = headers['gwf'][0]
         dtype = np.dtype(
             [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
             + [(name, np.float64) for name in header_advanced_package.auxtxt]
@@ -484,3 +478,9 @@ def grid_info(like: xr.DataArray) -> Dict[str, Any]:
             "x": like["x"],
         },
     }
+
+def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
+    for key, header in headers.items():
+        if 'gwf_' in key:
+            return header[0]
+    return None
\ No newline at end of file
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index ade7b9b66..35773470d 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -1,7 +1,6 @@
 import os
 import struct
-from dis import get_header_advanced_package
-from typing import Any, BinaryIO, Dict, List, Optional, Tuple
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
 
 import dask
 import numba
@@ -474,12 +473,11 @@ def open_cbc(
     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.
-        header = headers['gwf'][0]
         dtype = np.dtype(
             [("id1", np.int32), ("id2", np.int32), ("budget", np.float64)]
-            + [(name, np.float64) for name in header.auxtxt]
+            + [(name, np.float64) for name in header_advanced_package.auxtxt]
         )
-        table = cbc.read_imeth6_budgets(cbc_path, header.nlist, dtype, header.pos)
+        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():
@@ -533,3 +531,9 @@ def grid_info(like: xu.UgridDataArray) -> Dict[str, Any]:
             facedim: like[facedim],
         },
     }
+
+def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
+    for key, header in headers.items():
+        if 'gwf_' in key:
+            return header[0]
+    return None
\ No newline at end of file

From 242c84edaf34269631d62e367d1eec249717b86c Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Wed, 6 Mar 2024 14:50:59 +0100
Subject: [PATCH 08/19] update tests

---
 imod/mf6/out/dis.py                           |  3 +-
 imod/mf6/simulation.py                        |  2 +-
 imod/mf6/uzf.py                               | 23 +++++----
 imod/tests/test_mf6/test_mf6_out.py           | 47 ++++++++++---------
 imod/tests/test_mf6/test_mf6_simulation.py    |  4 +-
 .../test_mf6/test_mf6_transport_model.py      |  4 +-
 imod/tests/test_mf6/test_mf6_uzf.py           | 10 ++--
 imod/tests/test_mf6/test_mf6_uzf_model.py     | 33 ++++---------
 ...6_partitioned_simulation_postprocessing.py |  6 +--
 9 files changed, 64 insertions(+), 68 deletions(-)

diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index f149bd231..c321bc6cb 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -481,6 +481,7 @@ def grid_info(like: xr.DataArray) -> Dict[str, Any]:
 
 def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
     for key, header in headers.items():
-        if 'gwf_' in key:
+        # multimodels have a gwf-gwf budget for flow-ja-face between domains
+        if not 'flow-ja-face' in key and 'gwf_' in key:
             return header[0]
     return None
\ No newline at end of file
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 895eaadcc..8850cedbe 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -638,7 +638,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
 
diff --git a/imod/mf6/uzf.py b/imod/mf6/uzf.py
index 67504d46c..17a76da85 100644
--- a/imod/mf6/uzf.py
+++ b/imod/mf6/uzf.py
@@ -1,6 +1,7 @@
 import numpy as np
 import xarray as xr
 
+from imod.prepare.layer import get_upper_active_layer_number
 from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
 from imod.schemata import (
@@ -186,12 +187,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 = (
@@ -238,6 +239,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
@@ -248,6 +250,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,
@@ -349,10 +352,13 @@ 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 = xr.ones_like(landflag).where(landflag.notnull(), other =0.0)
+        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 = xr.ones_like(kv_sat,dtype = np.int32).where(kv_sat.layer == get_upper_active_layer_number(kv_sat),other = 0)
+        return land_nodes.where(kv_sat.notnull())
 
     def _determine_vertical_connection(self, uzf_number):
         return uzf_number.shift(layer=-1, fill_value=0)
@@ -429,6 +435,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
diff --git a/imod/tests/test_mf6/test_mf6_out.py b/imod/tests/test_mf6/test_mf6_out.py
index 8c6f491dc..236172d07 100644
--- a/imod/tests/test_mf6/test_mf6_out.py
+++ b/imod/tests/test_mf6/test_mf6_out.py
@@ -64,14 +64,14 @@ def test_read_cbc_headers(twri_result):
         headers = imod.mf6.read_cbc_headers("GWF_1/GWF_1.cbc")
         assert isinstance(headers, dict)
         assert sorted(headers.keys()) == [
-            "chd",
-            "drn",
+            "chd_chd",
+            "drn_drn",
             "flow-ja-face",
-            "wel",
+            "wel_wel",
         ]
-        assert isinstance(headers["chd"], list)
+        assert isinstance(headers["chd_chd"], list)
         assert isinstance(headers["flow-ja-face"][0], imod.mf6.out.cbc.Imeth1Header)
-        assert isinstance(headers["chd"][0], imod.mf6.out.cbc.Imeth6Header)
+        assert isinstance(headers["chd_chd"][0], imod.mf6.out.cbc.Imeth6Header)
 
 
 @pytest.mark.usefixtures("transient_twri_result")
@@ -81,15 +81,15 @@ def test_read_cbc_headers__transient(transient_twri_result):
         headers = imod.mf6.read_cbc_headers("GWF_1/GWF_1.cbc")
         assert isinstance(headers, dict)
         assert sorted(headers.keys()) == [
-            "chd",
-            "drn",
+            "chd_chd",
+            "drn_drn",
             "flow-ja-face",
             "sto-ss",
-            "wel",
+            "wel_wel",
         ]
-        assert isinstance(headers["chd"], list)
+        assert isinstance(headers["chd_chd"], list)
         assert isinstance(headers["flow-ja-face"][0], imod.mf6.out.cbc.Imeth1Header)
-        assert isinstance(headers["chd"][0], imod.mf6.out.cbc.Imeth6Header)
+        assert isinstance(headers["chd_chd"][0], imod.mf6.out.cbc.Imeth6Header)
         assert isinstance(headers["sto-ss"][0], imod.mf6.out.cbc.Imeth1Header)
 
 
@@ -189,13 +189,14 @@ def test_open_cbc__dis(twri_result):
     with imod.util.cd(modeldir):
         cbc = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/dis.dis.grb")
         assert isinstance(cbc, dict)
+        # packagename_packagetype
         assert sorted(cbc.keys()) == [
-            "chd",
-            "drn",
+            "chd_chd",
+            "drn_drn",
             "flow-front-face",
             "flow-lower-face",
             "flow-right-face",
-            "wel",
+            "wel_wel",
         ]
         for array in cbc.values():
             assert array.shape == (1, 3, 15, 15)
@@ -213,13 +214,13 @@ def test_open_cbc__dis_transient(transient_twri_result):
         cbc = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/dis.dis.grb")
         assert isinstance(cbc, dict)
         assert sorted(cbc.keys()) == [
-            "chd",
-            "drn",
+            "chd_chd",
+            "drn_drn",
             "flow-front-face",
             "flow-lower-face",
             "flow-right-face",
             "sto-ss",
-            "wel",
+            "wel_wel",
         ]
         for array in cbc.values():
             assert array.shape == (30, 3, 15, 15)
@@ -252,8 +253,8 @@ def test_open_cbc__dis_transient_unconfined(transient_unconfined_twri_result):
         cbc = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/dis.dis.grb")
         assert isinstance(cbc, dict)
         assert sorted(cbc.keys()) == [
-            "chd",
-            "drn",
+            "chd_chd",
+            "drn_drn",
             "flow-front-face",
             "flow-lower-face",
             "flow-right-face",
@@ -263,7 +264,7 @@ def test_open_cbc__dis_transient_unconfined(transient_unconfined_twri_result):
             "npf-sat",
             "sto-ss",
             "sto-sy",
-            "wel",
+            "wel_wel",
         ]
         for array in cbc.values():
             assert array.shape == (30, 3, 15, 15)
@@ -281,14 +282,14 @@ def test_open_cbc__disv(circle_result):
         cbc = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/disv.disv.grb")
         assert isinstance(cbc, dict)
         assert sorted(cbc.keys()) == [
-            "chd",
+            "chd_chd",
             "flow-horizontal-face",
             "flow-horizontal-face-x",
             "flow-horizontal-face-y",
             "flow-lower-face",
         ]
         for key, array in cbc.items():
-            if key in ("chd", "flow-lower-face"):
+            if key in ("chd_chd", "flow-lower-face"):
                 assert array.shape == (1, 2, 216)
                 assert array.dims[-1] == array.ugrid.grid.face_dimension
             else:
@@ -324,7 +325,7 @@ def test_open_cbc__disv_sto(circle_result_sto):
         cbc = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/disv.disv.grb")
         assert isinstance(cbc, dict)
         assert sorted(cbc.keys()) == [
-            "chd",
+            "chd_chd",
             "flow-horizontal-face",
             "flow-horizontal-face-x",
             "flow-horizontal-face-y",
@@ -332,7 +333,7 @@ def test_open_cbc__disv_sto(circle_result_sto):
             "sto-ss",
         ]
         for key, array in cbc.items():
-            if key in ("chd", "flow-lower-face", "sto-ss"):
+            if key in ("chd_chd", "flow-lower-face", "sto-ss"):
                 assert array.shape == (1, 2, 216)
                 assert array.dims[-1] == array.ugrid.grid.face_dimension
             else:
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index 8218c70fe..d7fa6d41b 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -120,13 +120,13 @@ def test_simulation_open_flow_budget(circle_model, tmp_path):
 
     assert isinstance(budget, xu.UgridDataset)
     assert sorted(budget.keys()) == [
-        "chd",
+        "chd_chd",
         "flow-horizontal-face",
         "flow-horizontal-face-x",
         "flow-horizontal-face-y",
         "flow-lower-face",
     ]
-    assert isinstance(budget["chd"], xu.UgridDataArray)
+    assert isinstance(budget["chd_chd"], xu.UgridDataArray)
 
 def test_write_circle_model_twice(circle_model, tmp_path):  
     
diff --git a/imod/tests/test_mf6/test_mf6_transport_model.py b/imod/tests/test_mf6/test_mf6_transport_model.py
index 4bb5cf21e..2b7745e94 100644
--- a/imod/tests/test_mf6/test_mf6_transport_model.py
+++ b/imod/tests/test_mf6/test_mf6_transport_model.py
@@ -124,8 +124,8 @@ def test_transport_balance_loading(tmp_path, flow_transport_simulation):
     assert balance_time.coords["time"].dtype == np.dtype("datetime64[ns]")
 
     assert np.all(
-        balance_notime.sel(species="a")["ssm"].values
-        == balance_time.sel(species="a")["ssm"].values
+        balance_notime.sel(species="a")["source-sink mix_ssm"].values
+        == balance_time.sel(species="a")["source-sink mix_ssm"].values
     )
 
 
diff --git a/imod/tests/test_mf6/test_mf6_uzf.py b/imod/tests/test_mf6/test_mf6_uzf.py
index f06c89c12..885b9c07c 100644
--- a/imod/tests/test_mf6/test_mf6_uzf.py
+++ b/imod/tests/test_mf6/test_mf6_uzf.py
@@ -57,9 +57,10 @@ def test_wrong_dtype(test_data):
 
 def test_landflag(test_data):
     expected = np.ones((2, 2, 2))
-    expected[:, 0, 0] = 0
+    expected[1:2, :, :] = 0
+    expected[:, 0, 0] = np.nan
     uzf = imod.mf6.UnsaturatedZoneFlow(**test_data)
-    assert np.all(uzf["landflag"].values == expected)
+    np.testing.assert_equal(uzf["landflag"].values,expected)
 
 
 def test_iuzno(test_data):
@@ -92,12 +93,13 @@ def test_to_sparsedata(test_data):
     arrdict = uzf._ds_to_arrdict(bin_data.isel(time=0))
     layer = bin_data.isel(time=0)["layer"].values
     struct_array = uzf._to_struct_array(arrdict, layer)
-    expected_iuzno = np.array([1, 2, 3, 4, 5, 6])
+    land_nodes = uzf["landflag"].to_numpy().flatten() == 1
+    expected_iuzno = np.array([0, 1, 2, 3, 0, 4, 5, 6])[land_nodes]
     assert struct_array.dtype[0] == np.dtype("int32")  # pylint: disable=unsubscriptable-object
     assert struct_array.dtype[1] == np.dtype("float64")  # pylint: disable=unsubscriptable-object
     assert np.all(struct_array["iuzno"] == expected_iuzno)
     assert len(struct_array.dtype) == 8
-    assert len(struct_array) == 6
+    assert len(struct_array) == expected_iuzno.size
 
 
 def test_fill_perioddata(test_data):
diff --git a/imod/tests/test_mf6/test_mf6_uzf_model.py b/imod/tests/test_mf6/test_mf6_uzf_model.py
index e329468c5..121c5fe58 100644
--- a/imod/tests/test_mf6/test_mf6_uzf_model.py
+++ b/imod/tests/test_mf6/test_mf6_uzf_model.py
@@ -11,7 +11,7 @@
 @pytest.fixture(scope="module")
 def uzf_model():
     # Initiate model
-    gwf_model = imod.mf6.GroundwaterFlowModel()
+    gwf_model = imod.mf6.GroundwaterFlowModel(newton = True)
 
     # Create discretication
     shape = nlay, nrow, ncol = 7, 9, 9
@@ -64,7 +64,7 @@ def uzf_model():
         icelltype=icelltype,
         k=k,
         k33=k33,
-        variable_vertical_conductance=True,
+        variable_vertical_conductance=False, # cant be true when newton is active
         dewatered=False,
         perched=False,
         save_flows=True,
@@ -98,9 +98,9 @@ def uzf_model():
     uds["theta_res"] = uzf_units * 0.05
     uds["theta_init"] = uzf_units * 0.08
     uds["epsilon"] = ones_shape * 7.0
-    uds["surface_depression_depth"] = ones_shape * top + 0.1
+    uds["surface_depression_depth"] = xr.where(ones_shape.layer == 1,ones_shape * top + 0.1, ones_shape * 0.0)
 
-    uds["infiltration_rate"] = ones_shape_time * 0.003
+    uds["infiltration_rate"] = (ones_shape_time * 0.003).where(ones_shape_time.layer == 1)
     uds["et_pot"] = (
         xr.DataArray(
             (np.sin(np.linspace(0, 1, num=nper) * 2 * np.pi) + 1) * 0.5 * 0.003,
@@ -108,12 +108,12 @@ def uzf_model():
             dims=("time",),
         )
         * ones_shape_time
-    )
-    uds["extinction_depth"] = ones_shape_time * -10.0
+    ).where(ones_shape_time.layer == 1)
+    uds["extinction_depth"] = (ones_shape_time * -10.0).where(ones_shape_time.layer == 1)
 
     uds[
         "simulate_groundwater_seepage"
-    ] = False  # Model doesn't converge if set to True....
+    ] = True 
 
     gwf_model["uzf"] = imod.mf6.UnsaturatedZoneFlow(**uds)
 
@@ -132,22 +132,7 @@ def uzf_model():
     simulation = imod.mf6.Modflow6Simulation("test")
     simulation["GWF_1"] = gwf_model
     # Define solver settings
-    simulation["solver"] = imod.mf6.Solution(
-        modelnames=["GWF_1"],
-        print_option="summary",
-        csv_output=False,
-        no_ptc=True,
-        outer_dvclose=1.0e-4,
-        outer_maximum=500,
-        under_relaxation=None,
-        inner_dvclose=1.0e-4,
-        inner_rclose=0.001,
-        inner_maximum=100,
-        linear_acceleration="cg",
-        scaling_method=None,
-        reordering_method=None,
-        relaxation_factor=0.97,
-    )
+    simulation["solver"] = imod.mf6.SolutionPresetComplex(modelnames=["GWF_1"])
     # Collect time discretization
     simulation.create_time_discretization(additional_times=time)
 
@@ -165,5 +150,5 @@ def test_simulation_write(uzf_model, tmp_path):
         assert head.dims == ("time", "layer", "y", "x")
         assert head.shape == (47, 7, 9, 9)
         meanhead = head.mean().values
-        mean_answer = 17.32092902
+        mean_answer = -1.54998241
         assert np.allclose(meanhead, mean_answer)
diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioned_simulation_postprocessing.py b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioned_simulation_postprocessing.py
index b94f86853..2389b3460 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioned_simulation_postprocessing.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioned_simulation_postprocessing.py
@@ -131,11 +131,11 @@ def test_import_balances_structured(
     # Assert
     expected_keys = [
         "gwf-gwf",
-        "chd",
+        "chd_chd",
         "flow-right-face",
         "sto-ss",
         "flow-lower-face",
-        "drn",
+        "drn_drn",
         "flow-front-face",
     ]
     expected_coords = ["x", "y", "layer", "time", "dx", "dy"]
@@ -165,7 +165,7 @@ def test_import_balances_unstructured(
 
     # Assert
     expected_keys = [
-        "chd",
+        "chd_chd",
         "flow-horizontal-face",
         "gwf-gwf",
         "flow-horizontal-face-y",

From fc01f4c50d5e5fe438c1af10d9ec0abacabaa2f3 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Wed, 6 Mar 2024 14:52:54 +0100
Subject: [PATCH 09/19] ruff

---
 imod/mf6/uzf.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/uzf.py b/imod/mf6/uzf.py
index 17a76da85..881d756ac 100644
--- a/imod/mf6/uzf.py
+++ b/imod/mf6/uzf.py
@@ -1,9 +1,9 @@
 import numpy as np
 import xarray as xr
 
-from imod.prepare.layer import get_upper_active_layer_number
 from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
+from imod.prepare.layer import get_upper_active_layer_number
 from imod.schemata import (
     AllInsideNoDataSchema,
     AllNoDataSchema,

From 246742df74ea6ae1270d135200b0b8b52bb21ba8 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 8 Mar 2024 10:36:10 +0100
Subject: [PATCH 10/19] update test to deal with nodata

---
 imod/tests/test_mf6/test_mf6_transport_model.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/imod/tests/test_mf6/test_mf6_transport_model.py b/imod/tests/test_mf6/test_mf6_transport_model.py
index 2b7745e94..2e802b041 100644
--- a/imod/tests/test_mf6/test_mf6_transport_model.py
+++ b/imod/tests/test_mf6/test_mf6_transport_model.py
@@ -123,9 +123,12 @@ def test_transport_balance_loading(tmp_path, flow_transport_simulation):
     )
     assert balance_time.coords["time"].dtype == np.dtype("datetime64[ns]")
 
-    assert np.all(
-        balance_notime.sel(species="a")["source-sink mix_ssm"].values
-        == balance_time.sel(species="a")["source-sink mix_ssm"].values
+    np.testing.assert_allclose(
+        balance_notime.sel(species="a")["source-sink mix_ssm"].values,
+        balance_time.sel(species="a")["source-sink mix_ssm"].values,
+        rtol=7e-5,
+        atol=3e-3,
+        equal_nan=True,
     )
 
 

From 07dbefb437e3ab432d9cb6986e8f840802b1dc6f Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 8 Mar 2024 10:48:10 +0100
Subject: [PATCH 11/19] move get_header_advanced_package method

---
 imod/mf6/out/common.py | 14 +++++++++++++-
 imod/mf6/out/dis.py    | 11 ++---------
 imod/mf6/out/disv.py   | 10 ++--------
 3 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/imod/mf6/out/common.py b/imod/mf6/out/common.py
index 8266e7e43..e11a0b56a 100644
--- a/imod/mf6/out/common.py
+++ b/imod/mf6/out/common.py
@@ -1,8 +1,10 @@
 import pathlib
-from typing import BinaryIO, Union
+from typing import BinaryIO, Dict, List, Union
 
 import numpy as np
 
+from . import cbc
+
 # Type annotations
 IntArray = np.ndarray
 FloatArray = np.ndarray
@@ -19,3 +21,13 @@ def _to_nan(a: FloatArray, dry_nan: bool) -> FloatArray:
     if dry_nan:
         a[a == -1e30] = np.nan
     return a
+
+
+def get_header_advanced_package(
+    headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]],
+) -> cbc.Imeth6Header | None:
+    for key, header in headers.items():
+        # multimodels have a gwf-gwf budget for flow-ja-face between domains
+        if not "flow-ja-face" in key and "gwf_" in key:
+            return header[0]
+    return None
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index c321bc6cb..1680fe631 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -11,7 +11,7 @@
 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_header_advanced_package
 
 
 # Binary Grid File / DIS Grids
@@ -477,11 +477,4 @@ def grid_info(like: xr.DataArray) -> Dict[str, Any]:
             "y": like["y"],
             "x": like["x"],
         },
-    }
-
-def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
-    for key, header in headers.items():
-        # multimodels have a gwf-gwf budget for flow-ja-face between domains
-        if not 'flow-ja-face' in key and 'gwf_' in key:
-            return header[0]
-    return None
\ No newline at end of file
+    }
\ No newline at end of file
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index 35773470d..244aa6aa1 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -12,7 +12,7 @@
 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_header_advanced_package
 
 
 def _ugrid_iavert_javert(
@@ -530,10 +530,4 @@ def grid_info(like: xu.UgridDataArray) -> Dict[str, Any]:
             "layer": like["layer"],
             facedim: like[facedim],
         },
-    }
-
-def get_header_advanced_package(headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]]) -> cbc.Imeth6Header | None:
-    for key, header in headers.items():
-        if 'gwf_' in key:
-            return header[0]
-    return None
\ No newline at end of file
+    }
\ No newline at end of file

From 3d73429f5300ec3a8be154404a609ede1d89d5ba Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 8 Mar 2024 11:35:48 +0100
Subject: [PATCH 12/19] fix circular import

---
 imod/mf6/out/common.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/imod/mf6/out/common.py b/imod/mf6/out/common.py
index e11a0b56a..8475ef7df 100644
--- a/imod/mf6/out/common.py
+++ b/imod/mf6/out/common.py
@@ -1,9 +1,8 @@
 import pathlib
-from typing import BinaryIO, Dict, List, Union
+from typing import BinaryIO, Dict, List, Union, Any
 
 import numpy as np
 
-from . import cbc
 
 # Type annotations
 IntArray = np.ndarray
@@ -24,8 +23,8 @@ def _to_nan(a: FloatArray, dry_nan: bool) -> FloatArray:
 
 
 def get_header_advanced_package(
-    headers: Dict[str, List[Union[cbc.Imeth1Header, cbc.Imeth6Header]]],
-) -> cbc.Imeth6Header | None:
+    headers: Dict[str, List[Any]],
+) -> Any:
     for key, header in headers.items():
         # multimodels have a gwf-gwf budget for flow-ja-face between domains
         if not "flow-ja-face" in key and "gwf_" in key:

From cf0936d8023e96c6f29cc56253fd8e54d0d641de Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 8 Mar 2024 12:25:49 +0100
Subject: [PATCH 13/19] linter

---
 imod/mf6/out/common.py | 3 +--
 imod/mf6/out/dis.py    | 2 +-
 imod/mf6/out/disv.py   | 2 +-
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/imod/mf6/out/common.py b/imod/mf6/out/common.py
index 8475ef7df..2e86b7fc0 100644
--- a/imod/mf6/out/common.py
+++ b/imod/mf6/out/common.py
@@ -1,9 +1,8 @@
 import pathlib
-from typing import BinaryIO, Dict, List, Union, Any
+from typing import Any, BinaryIO, Dict, List, Union
 
 import numpy as np
 
-
 # Type annotations
 IntArray = np.ndarray
 FloatArray = np.ndarray
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index 1680fe631..db0441384 100644
--- a/imod/mf6/out/dis.py
+++ b/imod/mf6/out/dis.py
@@ -1,6 +1,6 @@
 import os
 import struct
-from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple
 
 import dask
 import numba
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index 244aa6aa1..c5ff29da6 100644
--- a/imod/mf6/out/disv.py
+++ b/imod/mf6/out/disv.py
@@ -1,6 +1,6 @@
 import os
 import struct
-from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple
 
 import dask
 import numba

From 15ff4fa4593475aa7cbb8cfdfc944dc456502e2d Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 3 May 2024 16:31:03 +0200
Subject: [PATCH 14/19] process review

---
 docs/api/changelog.rst                    | 11 ++++++-
 imod/mf6/out/cbc.py                       | 20 +++++++------
 imod/mf6/out/common.py                    |  6 ++--
 imod/mf6/out/dis.py                       | 35 +++++++++++++++++------
 imod/mf6/out/disv.py                      | 33 ++++++++++++++++-----
 imod/mf6/uzf.py                           | 16 +++++------
 imod/tests/test_mf6/test_mf6_uzf.py       |  2 +-
 imod/tests/test_mf6/test_mf6_uzf_model.py | 30 ++++++++++++-------
 8 files changed, 106 insertions(+), 47 deletions(-)

diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
index 79fbf4338..a4dcdc5e7 100644
--- a/docs/api/changelog.rst
+++ b/docs/api/changelog.rst
@@ -24,7 +24,16 @@ Added
   be selected, available in :func:`imod.prepare.DISTRIBUTION_OPTION`.
 - :func:`imod.prepare.celltable` supports an optional ``dtype`` argument. This
   can be used, for example, to create celltables of float values.
-
+- Budget arrays now 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. Changes affect:  
+  :func:`imod.mf6.open_cbc`
+- Package type added to return budget names of CBC-output. 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. Changes affect: 
+  :func:`imod.mf6.open_cbc`
 
 Fixed
 ~~~~~
diff --git a/imod/mf6/out/cbc.py b/imod/mf6/out/cbc.py
index 78c3c5d48..2dccb24d3 100644
--- a/imod/mf6/out/cbc.py
+++ b/imod/mf6/out/cbc.py
@@ -132,10 +132,14 @@ def read_cbc_headers(
                 # 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"]
+                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 = imeth6_header["txt2id2"] + "_" + header["text"].replace("data-", "")
+                    key = (
+                        imeth6_header["txt2id2"]
+                        + "_"
+                        + header["text"].replace("data-", "")
+                    )
                 headers[key].append(Imeth6Header(**header, **imeth6_header))
             else:
                 raise ValueError(
@@ -264,7 +268,7 @@ def read_imeth6_budgets_dense(
     size: int,
     shape: tuple,
     return_variable: str,
-    return_id: np.ndarray | None
+    indices: np.ndarray | None,
 ) -> FloatArray:
     """
     Read the data for an imeth==6 budget section.
@@ -287,8 +291,8 @@ def read_imeth6_budgets_dense(
     shape: tuple[int, int, int]
         Shape (nlayer, nrow, ncolumn) of entire model domain.
     return_variable: str
-        variable name to return from budget table 
-    return_id: np.ndarray | None
+        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
@@ -298,7 +302,7 @@ def read_imeth6_budgets_dense(
     # Allocates a dense array for the entire domain
     out = np.full(size, np.nan, dtype=np.float64)
     table = read_imeth6_budgets(cbc_path, count, dtype, pos)
-    if return_id is None:
-        return_id = table["id1"] - 1  # Convert to 0 based index
-    out[return_id] = 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)
diff --git a/imod/mf6/out/common.py b/imod/mf6/out/common.py
index 2e86b7fc0..8f4f1c2a0 100644
--- a/imod/mf6/out/common.py
+++ b/imod/mf6/out/common.py
@@ -21,11 +21,11 @@ def _to_nan(a: FloatArray, dry_nan: bool) -> FloatArray:
     return a
 
 
-def get_header_advanced_package(
+def get_first_header_advanced_package(
     headers: Dict[str, List[Any]],
 ) -> Any:
-    for key, header in headers.items():
+    for key, header_list in headers.items():
         # multimodels have a gwf-gwf budget for flow-ja-face between domains
         if not "flow-ja-face" in key and "gwf_" in key:
-            return header[0]
+            return header_list[0]
     return None
diff --git a/imod/mf6/out/dis.py b/imod/mf6/out/dis.py
index db0441384..19792897b 100644
--- a/imod/mf6/out/dis.py
+++ b/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, get_header_advanced_package
+from .common import (
+    FilePath,
+    FloatArray,
+    IntArray,
+    _to_nan,
+    get_first_header_advanced_package,
+)
 
 
 # Binary Grid File / DIS Grids
@@ -197,7 +203,7 @@ def open_imeth6_budgets(
     grb_content: dict,
     header_list: List[cbc.Imeth6Header],
     return_variable: str = "budget",
-    return_id: np.ndarray | None = None,
+    indices: np.ndarray | None = None,
 ) -> xr.DataArray:
     """
     Open the data for an imeth==6 budget section.
@@ -230,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, return_id
+            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)
@@ -414,7 +427,7 @@ def open_cbc(
 ) -> Dict[str, xr.DataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
     indices = None
-    header_advanced_package = get_header_advanced_package(headers)
+    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.
@@ -422,7 +435,9 @@ def open_cbc(
             [("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)
+        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():
@@ -452,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 = return_variable, return_id = indices
+                            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, return_id = indices
+                        cbc_path, grb_content, header_list, indices=indices
                     )
     if simulation_start_time is not None:
         for cbc_name, cbc_array in cbc_content.items():
@@ -477,4 +496,4 @@ def grid_info(like: xr.DataArray) -> Dict[str, Any]:
             "y": like["y"],
             "x": like["x"],
         },
-    }
\ No newline at end of file
+    }
diff --git a/imod/mf6/out/disv.py b/imod/mf6/out/disv.py
index c5ff29da6..628701f60 100644
--- a/imod/mf6/out/disv.py
+++ b/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, get_header_advanced_package
+from .common import (
+    FilePath,
+    FloatArray,
+    IntArray,
+    _to_nan,
+    get_first_header_advanced_package,
+)
 
 
 def _ugrid_iavert_javert(
@@ -262,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, return_id
+            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)
@@ -469,7 +482,7 @@ def open_cbc(
 ) -> Dict[str, xu.UgridDataArray]:
     headers = cbc.read_cbc_headers(cbc_path)
     indices = None
-    header_advanced_package = get_header_advanced_package(headers)
+    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.
@@ -477,7 +490,9 @@ def open_cbc(
             [("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)
+        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():
@@ -504,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 = return_variable, return_id = indices
+                        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, return_id = indices
+                    cbc_path, grb_content, header_list, return_id=indices
                 )
 
     if simulation_start_time is not None:
@@ -530,4 +549,4 @@ def grid_info(like: xu.UgridDataArray) -> Dict[str, Any]:
             "layer": like["layer"],
             facedim: like[facedim],
         },
-    }
\ No newline at end of file
+    }
diff --git a/imod/mf6/uzf.py b/imod/mf6/uzf.py
index b28077847..2376cc351 100644
--- a/imod/mf6/uzf.py
+++ b/imod/mf6/uzf.py
@@ -4,7 +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_layer_number
+from imod.prepare.layer import get_upper_active_grid_cells
 from imod.schemata import (
     AllInsideNoDataSchema,
     AllNoDataSchema,
@@ -231,8 +231,8 @@ def __init__(
         print_input=False,
         print_flows=False,
         save_flows=False,
-        budget_fileout = None,
-        budgetcsv_fileout = None,
+        budget_fileout=None,
+        budgetcsv_fileout=None,
         observations=None,
         water_mover=None,
         timeseries=None,
@@ -354,12 +354,12 @@ def _check_options(
 
     def _create_uzf_numbers(self, landflag):
         """Create unique UZF ID's. Inactive cells equal 0"""
-        active_nodes = xr.ones_like(landflag).where(landflag.notnull(), other =0.0)
+        active_nodes = landflag.notnull().astype(np.int8)
         return np.nancumsum(active_nodes).reshape(landflag.shape) * active_nodes
 
     def _determine_landflag(self, kv_sat):
-        """ returns the landflag for uzf-model. Landflag == 1 for top active UZF-nodes """
-        land_nodes = xr.ones_like(kv_sat,dtype = np.int32).where(kv_sat.layer == get_upper_active_layer_number(kv_sat),other = 0)
+        """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):
@@ -392,7 +392,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 = {}
@@ -407,7 +407,7 @@ def render(self, directory, pkgname, globaltimes, binary):
         path = directory / pkgname / f"{self._pkg_id}-pkgdata.dat"
         d["packagedata"] = path.as_posix()
         # max uzf-cells for which time period data will be supplied
-        d["nuzfcells"] =  np.count_nonzero(np.isfinite(d['landflag'])) 
+        d["nuzfcells"] = np.count_nonzero(np.isfinite(d["landflag"]))
         return self._template.render(d)
 
     def _to_struct_array(self, arrdict, layer):
diff --git a/imod/tests/test_mf6/test_mf6_uzf.py b/imod/tests/test_mf6/test_mf6_uzf.py
index 885b9c07c..1456d9d2b 100644
--- a/imod/tests/test_mf6/test_mf6_uzf.py
+++ b/imod/tests/test_mf6/test_mf6_uzf.py
@@ -60,7 +60,7 @@ def test_landflag(test_data):
     expected[1:2, :, :] = 0
     expected[:, 0, 0] = np.nan
     uzf = imod.mf6.UnsaturatedZoneFlow(**test_data)
-    np.testing.assert_equal(uzf["landflag"].values,expected)
+    np.testing.assert_equal(uzf["landflag"].values, expected)
 
 
 def test_iuzno(test_data):
diff --git a/imod/tests/test_mf6/test_mf6_uzf_model.py b/imod/tests/test_mf6/test_mf6_uzf_model.py
index 69d748045..2c4a5c27e 100644
--- a/imod/tests/test_mf6/test_mf6_uzf_model.py
+++ b/imod/tests/test_mf6/test_mf6_uzf_model.py
@@ -11,7 +11,7 @@
 @pytest.fixture(scope="module")
 def uzf_model():
     # Initiate model
-    gwf_model = imod.mf6.GroundwaterFlowModel(newton = True)
+    gwf_model = imod.mf6.GroundwaterFlowModel(newton=True)
 
     # Create discretication
     shape = nlay, nrow, ncol = 7, 9, 9
@@ -64,7 +64,7 @@ def uzf_model():
         icelltype=icelltype,
         k=k,
         k33=k33,
-        variable_vertical_conductance=False, # cant be true when newton is active
+        variable_vertical_conductance=False,  # cant be true when newton is active
         dewatered=False,
         perched=False,
         save_flows=True,
@@ -98,9 +98,13 @@ def uzf_model():
     uds["theta_res"] = uzf_units * 0.05
     uds["theta_init"] = uzf_units * 0.08
     uds["epsilon"] = ones_shape * 7.0
-    uds["surface_depression_depth"] = xr.where(ones_shape.layer == 1,ones_shape * top + 0.1, ones_shape * 0.0)
+    uds["surface_depression_depth"] = xr.where(
+        ones_shape.layer == 1, ones_shape * top + 0.1, ones_shape * 0.0
+    )
 
-    uds["infiltration_rate"] = (ones_shape_time * 0.003).where(ones_shape_time.layer == 1)
+    uds["infiltration_rate"] = (ones_shape_time * 0.003).where(
+        ones_shape_time.layer == 1
+    )
     uds["et_pot"] = (
         xr.DataArray(
             (np.sin(np.linspace(0, 1, num=nper) * 2 * np.pi) + 1) * 0.5 * 0.003,
@@ -109,13 +113,13 @@ def uzf_model():
         )
         * ones_shape_time
     ).where(ones_shape_time.layer == 1)
-    uds["extinction_depth"] = (ones_shape_time * -10.0).where(ones_shape_time.layer == 1)
+    uds["extinction_depth"] = (ones_shape_time * -10.0).where(
+        ones_shape_time.layer == 1
+    )
 
-    uds["simulate_groundwater_seepage"] = (
-        False  # Model doesn't converge if set to True....
-    uds[
-        "simulate_groundwater_seepage"
-    ] = True 
+    uds["simulate_groundwater_seepage"] = True
+    uds["save_flows"] = True
+    uds["budget_fileout"] = "GWF_1/uzf.cbc"
 
     gwf_model["uzf"] = imod.mf6.UnsaturatedZoneFlow(**uds)
 
@@ -137,7 +141,6 @@ def uzf_model():
     simulation["solver"] = imod.mf6.SolutionPresetComplex(modelnames=["GWF_1"])
     # Collect time discretization
     simulation.create_time_discretization(additional_times=time)
-
     return simulation
 
 
@@ -154,3 +157,8 @@ def test_simulation_write(uzf_model, tmp_path):
         meanhead = head.mean().values
         mean_answer = -1.54998241
         assert np.allclose(meanhead, mean_answer)
+        budget_mf6 = head = imod.mf6.open_cbc("GWF_1/GWF_1.cbc", "GWF_1/dis.dis.grb")
+        budgets_uzf = imod.mf6.open_cbc("GWF_1/uzf.cbc", "GWF_1/dis.dis.grb")
+        assert np.allclose(
+            budget_mf6["uzf-gwrch_uzf"], -budgets_uzf["gwf_gwf_1"], equal_nan=True
+        )

From c355535ffb8a47c51b302cba4370f9d7ada6a279 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 3 May 2024 21:55:55 +0200
Subject: [PATCH 15/19] fix labels in test

---
 .../test_multimodel/test_mf6_partitioning_unstructured.py   | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
index 00fd324e0..4c72d1454 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
@@ -427,7 +427,7 @@ def test_partition_transport(
     )
 
     for budget_term in (
-        "ssm",
+        "source-sink mix_ssm",
         "flow-lower-face",
         "storage-aqueous",
         "flow-horizontal-face",
@@ -489,7 +489,7 @@ def test_partition_transport_multispecies(
         actual_flow_budget, expected_flow_budget
     )
 
-    for key in ["flow-lower-face", "flow-horizontal-face", "sto-ss", "rch"]:
+    for key in ["flow-lower-face", "flow-horizontal-face", "sto-ss", "rch_rch"]:
         marker = is_exchange_cell
         if key == "flow-horizontal-face":
             marker = is_exchange_edge
@@ -502,7 +502,7 @@ def test_partition_transport_multispecies(
             rtol=rtol,
             atol=atol,
         )
-    for key in ["flow-lower-face", "flow-horizontal-face", "storage-aqueous", "ssm"]:
+    for key in ["flow-lower-face", "flow-horizontal-face", "storage-aqueous", "source-sink mix_ssm"]:
         marker = is_exchange_cell
         if key == "flow-horizontal-face":
             marker = is_exchange_edge

From b2db8ad22ecd0605ac657770ea5b84624977de97 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Fri, 3 May 2024 22:05:58 +0200
Subject: [PATCH 16/19] linter

---
 .../test_multimodel/test_mf6_partitioning_unstructured.py  | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
index 4c72d1454..e01658222 100644
--- a/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
+++ b/imod/tests/test_mf6/test_multimodel/test_mf6_partitioning_unstructured.py
@@ -502,7 +502,12 @@ def test_partition_transport_multispecies(
             rtol=rtol,
             atol=atol,
         )
-    for key in ["flow-lower-face", "flow-horizontal-face", "storage-aqueous", "source-sink mix_ssm"]:
+    for key in [
+        "flow-lower-face",
+        "flow-horizontal-face",
+        "storage-aqueous",
+        "source-sink mix_ssm",
+    ]:
         marker = is_exchange_cell
         if key == "flow-horizontal-face":
             marker = is_exchange_edge

From f93ea349d5d174f47a72a888648ffba171d8f3e4 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Mon, 13 May 2024 08:39:36 +0200
Subject: [PATCH 17/19] update changelog with master

---
 docs/api/changelog.rst | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
index a4dcdc5e7..f6ae4fd70 100644
--- a/docs/api/changelog.rst
+++ b/docs/api/changelog.rst
@@ -21,9 +21,17 @@ Added
   topsystem in :func:`imod.prepare.distribute_riv_conductance`,
   :func:`imod.prepare.distribute_drn_conductance`,
   :func:`imod.prepare.distribute_ghb_conductance`, for this multiple options can
-  be selected, available in :func:`imod.prepare.DISTRIBUTION_OPTION`.
+  be selected, available in :func:`imod.prepare.DISTRIBUTING_OPTION`.
 - :func:`imod.prepare.celltable` supports an optional ``dtype`` argument. This
   can be used, for example, to create celltables of float values.
+- Added ``fixed_cell`` option to :class:`imod.mf6.Recharge`. This option is
+  relevant for phreatic models, not using the Newton formulation and model cells
+  can become inactive. The prefered method for phreatic models is to use the
+  Newton formulation, where cells remain active, and this option irrelevant.
+- Added support for ``ats_outer_maximum_fraction`` in :class:`imod.mf6.Solution`.
+- Added validation for ``linear_acceleration``, ``rclose_option``,
+  ``scaling_method``, ``reordering_method``, ``print_option`` and ``no_ptc``
+  entries in :class:`imod.mf6.Solution`.
 - Budget arrays now 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. Changes affect:  
@@ -44,7 +52,16 @@ 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
+  :class:`imod.mf6.Solution` to match newer MODFLOW 6 releases.
+- Changed no_ptc from a bool to an option string in :class:`imod.mf6.Solution`.
+- Removed constructor arguments `source` and `target` from
+  :class:`imod.mf6.utilities.regrid.RegridderWeightsCache`, as they were not
+  used.
+
 [0.16.0] - 2024-03-29
 ---------------------
 

From 62819f85bc854ae0db579af297c4d2eb69cb0f6b Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Mon, 13 May 2024 08:44:35 +0200
Subject: [PATCH 18/19] linter

---
 imod/mf6/out/common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imod/mf6/out/common.py b/imod/mf6/out/common.py
index 8f4f1c2a0..b642ec57e 100644
--- a/imod/mf6/out/common.py
+++ b/imod/mf6/out/common.py
@@ -26,6 +26,6 @@ def get_first_header_advanced_package(
 ) -> Any:
     for key, header_list in headers.items():
         # multimodels have a gwf-gwf budget for flow-ja-face between domains
-        if not "flow-ja-face" in key and "gwf_" in key:
+        if "flow-ja-face" not in key and "gwf_" in key:
             return header_list[0]
     return None

From 0fb5c5d3f29a8905bcf3d848594ceaf818d88627 Mon Sep 17 00:00:00 2001
From: HendrikKok <hendrik.kok@deltares.nl>
Date: Mon, 13 May 2024 09:46:42 +0200
Subject: [PATCH 19/19] fix comment

---
 docs/api/changelog.rst | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
index f6ae4fd70..b5ebc2434 100644
--- a/docs/api/changelog.rst
+++ b/docs/api/changelog.rst
@@ -32,16 +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`.
-- Budget arrays now 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. Changes affect:  
-  :func:`imod.mf6.open_cbc`
-- Package type added to return budget names of CBC-output. New format is 
-  "package type"-"optional package variable"_"package name". E.g. a River package 
+- :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. Changes affect: 
-  :func:`imod.mf6.open_cbc`
+  groundwater recharge budget from the UZF-CBC.
 
 Fixed
 ~~~~~