From a64eb007c91742522e1527841fe49d897ea9dcc1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Thu, 22 Feb 2024 23:21:12 +0100
Subject: [PATCH 1/6] Add get_flow_residuals and get_flow_lower_face

---
 nlmod/gwf/output.py          | 143 ++++++++++++++++++++++++++++++++++-
 nlmod/mfoutput/mfoutput.py   |  13 +++-
 tests/test_015_gwf_output.py |  16 ++++
 3 files changed, 167 insertions(+), 5 deletions(-)

diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py
index ed4ed748..88c9c4ed 100644
--- a/nlmod/gwf/output.py
+++ b/nlmod/gwf/output.py
@@ -8,11 +8,13 @@
 from shapely.geometry import Point
 
 from ..dims.grid import modelgrid_from_ds
+from ..dims.resample import get_affine_world_to_mod
 from ..mfoutput.mfoutput import (
     _get_budget_da,
     _get_heads_da,
     _get_time_index,
     _get_flopy_data_object,
+    _get_grbfile,
 )
 
 logger = logging.getLogger(__name__)
@@ -232,7 +234,139 @@ def get_gwl_from_wet_cells(head, layer="layer", botm=None):
     return gwl
 
 
-def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True):
+def get_flow_residuals(ds, kstpkper=None, grb_file=None):
+    """
+    Get the flow residuals of a MODFLOW 6 simulation.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Xarray dataset with model data.
+    kstpkper : tuple of 2 ints, optional
+        The index of the timestep and the stress period to include in the result. Include
+        all data in the budget-file when None. The default is None.
+    grb_file : str, optional
+        The location of the grb-file. grb_file is determied from ds when None. The
+        default is None.
+
+    Returns
+    -------
+    da : xr.DataArray
+        The flow residual in each cell, in m3/d.
+
+    """
+    if grb_file is None:
+        grb_file = _get_grbfile(ds)
+    grb = flopy.mf6.utils.MfGrdFile(grb_file)
+    cbf = get_cellbudgetfile(ds)
+    dims = ds["botm"].dims
+    coords = ds["botm"].coords
+    flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper)
+    mask_active = np.diff(grb.ia) > 0
+    flowja_index = grb.ia[:-1][mask_active]
+    if kstpkper is None:
+        # loop over all timesteps/stress-periods
+        residuals = []
+        for iflowja in flowja:
+            # residuals.append(flopy.mf6.utils.get_residuals(iflowja, grb_file))
+            # use our own faster method instead of a for loop:
+            residual = np.full(grb.shape, np.NaN)
+            residual.ravel()[mask_active] = iflowja.flatten()[flowja_index]
+            residuals.append(residual)
+        dims = ("time",) + dims
+        coords = dict(coords) | {"time": _get_time_index(cbf, ds)}
+    else:
+        # residuals = flopy.mf6.utils.get_residuals(flowja[0], grb_file)
+        # use our own faster method instead of a for loop:
+        residuals = np.full(grb.shape, np.NaN)
+        residuals.ravel()[mask_active] = flowja[0].flatten()[flowja_index]
+    da = xr.DataArray(residuals, dims=dims, coords=coords)
+    return da
+
+
+def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
+    """
+    Get the flow over the lower face of all model cells
+
+    THe flow Lower Face (flf) used to be written to the budget file in previous versions
+    of MODFLOW.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Xarray dataset with model data.
+    kstpkper : tuple of 2 ints, optional
+        The index of the timestep and the stress period to include in the result. Include
+        all data in the budget-file when None. The default is None.
+    grb_file : str, optional
+        The location of the grb-file. grb_file is determied from ds when None. The
+        default is None.
+    lays : int or list of ints, optional
+        The layers to include in the result. When lays is None, all layers are included.
+        The default is None.
+
+    Returns
+    -------
+    da : xr.DataArray
+        The flow over the lower face of each cell, in m3/d.
+
+    """
+    if grb_file is None:
+        grb_file = _get_grbfile(ds)
+    cbf = get_cellbudgetfile(ds)
+    flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper)
+
+    if ds.gridtype == "vertex":
+        # detemine flf_index first
+        grb = flopy.mf6.utils.MfGrdFile(grb_file)
+
+        if lays is None:
+            lays = range(grb.nlay)
+        if isinstance(lays, int):
+            lays = [lays]
+        shape = (len(lays), len(ds.icell2d))
+
+        flf_index = np.full(shape, -1)
+        # get these properties outside of the for loop to increase speed
+        grb_ia = grb.ia
+        grb_ja = grb.ja
+        for ilay, lay in enumerate(lays):
+            ja_start_next_layer = (lay + 1) * grb.ncpl
+            for icell2d in range(grb.ncpl):
+                node = lay * grb.ncpl + icell2d
+                ia = np.arange(grb_ia[node], grb_ia[node + 1])
+                mask = grb_ja[ia] >= ja_start_next_layer
+                if mask.any():
+                    # assert mask.sum() == 1
+                    flf_index[ilay, icell2d] = int(ia[mask])
+
+    dims = ds["botm"].dims
+    coords = ds["botm"][lays].coords
+    if kstpkper is None:
+        # loop over all tiesteps/stress-periods
+        flfs = []
+        for iflowja in flowja:
+            if ds.gridtype == "vertex":
+                flf = np.full(shape, np.NaN)
+                mask = flf_index >= 0
+                flf[mask] = iflowja[0, 0, flf_index[mask]]
+            else:
+                _, _, flf = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
+            flfs.append(flf)
+        dims = ("time",) + dims
+        coords = dict(coords) | {"time": _get_time_index(cbf, ds)}
+    else:
+        if ds.gridtype == "vertex":
+            flfs = np.full(shape, np.NaN)
+            mask = flf_index >= 0
+            flfs[mask] = flowja[0][0, 0, flf_index[mask]]
+        else:
+            _, _, flfs = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
+    da = xr.DataArray(flfs, dims=dims, coords=coords)
+    return da
+
+
+def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True, rotated=True):
     """Get the head at a certain point from a head DataArray for all cells.
 
     Parameters
@@ -253,12 +387,19 @@ def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True):
         None.
     drop_nan_layers : bool, optional
         Drop layers that are NaN at all timesteps. The default is True.
+    rotated : bool, optional
+        If rotated is True, x and y are in real world coordinates. If rotated is False,
+        x and y are in model coordinates.
 
     Returns
     -------
     head_point : xarray.DataArray
         A DataArray with dimensions (time, layer).
     """
+    if rotated:
+        # calculate real-world coordinates to model coordinates
+        x, y = get_affine_world_to_mod(ds) * (x, y)
+
     if "icell2d" in head.dims:
         if gi is None:
             if ds is None:
diff --git a/nlmod/mfoutput/mfoutput.py b/nlmod/mfoutput/mfoutput.py
index 64158e3d..9a1db186 100644
--- a/nlmod/mfoutput/mfoutput.py
+++ b/nlmod/mfoutput/mfoutput.py
@@ -285,10 +285,7 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
         fname = os.path.join(ds.model_ws, ds.model_name + extension)
     if grbfile is None and ds is not None:
         # get grb file
-        if ds.gridtype == "vertex":
-            grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
-        elif ds.gridtype == "structured":
-            grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
+        grbfile = _get_grbfile(ds)
     if grbfile is not None and os.path.exists(grbfile):
         modelgrid = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
     elif ds is not None:
@@ -307,3 +304,11 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
             logger.warning(msg)
             warnings.warn(msg)
         return flopy.utils.HeadFile(fname, text=var, modelgrid=modelgrid)
+
+
+def _get_grbfile(ds):
+    if ds.gridtype == "vertex":
+        grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
+    elif ds.gridtype == "structured":
+        grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
+    return grbfile
diff --git a/tests/test_015_gwf_output.py b/tests/test_015_gwf_output.py
index 074e920b..aafba45b 100644
--- a/tests/test_015_gwf_output.py
+++ b/tests/test_015_gwf_output.py
@@ -199,3 +199,19 @@ def test_postprocess_head():
     nlmod.gwf.get_gwl_from_wet_cells(head, botm=ds["botm"])
 
     nlmod.gwf.get_head_at_point(head, float(ds.x.mean()), float(ds.y.mean()), ds=ds)
+
+
+def test_get_flow_residuals():
+    ds = test_001_model.get_ds_from_cache("basic_sea_model")
+    da = nlmod.gwf.output.get_flow_residuals(ds)
+    assert "time" in da.dims
+    da = nlmod.gwf.output.get_flow_residuals(ds, kstpkper=(0, 0))
+    assert "time" not in da.dims
+
+
+def test_get_flow_lower_face():
+    ds = test_001_model.get_ds_from_cache("basic_sea_model")
+    da = nlmod.gwf.output.get_flow_lower_face(ds)
+    assert "time" in da.dims
+    da = nlmod.gwf.output.get_flow_lower_face(ds, kstpkper=(0, 0))
+    assert "time" not in da.dims

From be078d169f4d97a40298addbda53458f668c1840 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Thu, 22 Feb 2024 23:36:03 +0100
Subject: [PATCH 2/6] Minor changes

Fix lays parameter for structured grids
---
 nlmod/gwf/output.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py
index 88c9c4ed..fd62a7c9 100644
--- a/nlmod/gwf/output.py
+++ b/nlmod/gwf/output.py
@@ -288,8 +288,8 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
     """
     Get the flow over the lower face of all model cells
 
-    THe flow Lower Face (flf) used to be written to the budget file in previous versions
-    of MODFLOW.
+    The flow Lower Face (flf) used to be written to the budget file in previous versions
+    of MODFLOW. In MODFLOW 6 we determine these flows from the flow-ja-face-records.
 
     Parameters
     ----------
@@ -352,6 +352,8 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
                 flf[mask] = iflowja[0, 0, flf_index[mask]]
             else:
                 _, _, flf = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
+                if lays is not None:
+                    flf = flf[lays]
             flfs.append(flf)
         dims = ("time",) + dims
         coords = dict(coords) | {"time": _get_time_index(cbf, ds)}
@@ -362,6 +364,8 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
             flfs[mask] = flowja[0][0, 0, flf_index[mask]]
         else:
             _, _, flfs = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
+            if lays is not None:
+                flf = flf[lays]
     da = xr.DataArray(flfs, dims=dims, coords=coords)
     return da
 

From 87444225ee95877371a0bda2ddea325e8a2be4a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Fri, 23 Feb 2024 00:02:15 +0100
Subject: [PATCH 3/6] Fix failing tests

---
 nlmod/gwf/output.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py
index fd62a7c9..5a68e252 100644
--- a/nlmod/gwf/output.py
+++ b/nlmod/gwf/output.py
@@ -339,9 +339,11 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
                 if mask.any():
                     # assert mask.sum() == 1
                     flf_index[ilay, icell2d] = int(ia[mask])
-
+        coords = ds["botm"][lays].coords
+    else:
+        coords = ds["botm"].coords
     dims = ds["botm"].dims
-    coords = ds["botm"][lays].coords
+
     if kstpkper is None:
         # loop over all tiesteps/stress-periods
         flfs = []
@@ -352,8 +354,6 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
                 flf[mask] = iflowja[0, 0, flf_index[mask]]
             else:
                 _, _, flf = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
-                if lays is not None:
-                    flf = flf[lays]
             flfs.append(flf)
         dims = ("time",) + dims
         coords = dict(coords) | {"time": _get_time_index(cbf, ds)}
@@ -364,13 +364,15 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
             flfs[mask] = flowja[0][0, 0, flf_index[mask]]
         else:
             _, _, flfs = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
-            if lays is not None:
-                flf = flf[lays]
     da = xr.DataArray(flfs, dims=dims, coords=coords)
+    if ds.gridtype != "vertex" and lays is not None:
+        da = da.isel(layer=lays)
     return da
 
 
-def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True, rotated=True):
+def get_head_at_point(
+    head, x, y, ds=None, gi=None, drop_nan_layers=True, rotated=False
+):
     """Get the head at a certain point from a head DataArray for all cells.
 
     Parameters
@@ -392,16 +394,17 @@ def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True, rotate
     drop_nan_layers : bool, optional
         Drop layers that are NaN at all timesteps. The default is True.
     rotated : bool, optional
-        If rotated is True, x and y are in real world coordinates. If rotated is False,
-        x and y are in model coordinates.
+        If the model grid has a rotation, and rotated is False, x and y are in model
+        coordinates. Otherwise x and y are in real world coordinates. The defaults is
+        False.
 
     Returns
     -------
     head_point : xarray.DataArray
         A DataArray with dimensions (time, layer).
     """
-    if rotated:
-        # calculate real-world coordinates to model coordinates
+    if rotated and "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
+        # calculate model coordinates from the specified real-world coordinates
         x, y = get_affine_world_to_mod(ds) * (x, y)
 
     if "icell2d" in head.dims:

From 33f5e327ac3288843cb64e9eb43bc949a3e44230 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Fri, 23 Feb 2024 08:53:47 +0100
Subject: [PATCH 4/6] Fix last failing test

---
 nlmod/gwf/output.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py
index 5a68e252..460463b9 100644
--- a/nlmod/gwf/output.py
+++ b/nlmod/gwf/output.py
@@ -363,7 +363,7 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
             mask = flf_index >= 0
             flfs[mask] = flowja[0][0, 0, flf_index[mask]]
         else:
-            _, _, flfs = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file)
+            _, _, flfs = flopy.mf6.utils.get_structured_faceflows(flowja[0], grb_file)
     da = xr.DataArray(flfs, dims=dims, coords=coords)
     if ds.gridtype != "vertex" and lays is not None:
         da = da.isel(layer=lays)

From a5edfe4d6c4eaee40e479b0a44b6737d1b72cf72 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Fri, 23 Feb 2024 14:17:08 +0100
Subject: [PATCH 5/6] Improve  nlmod.gwf.output.get_budget_da()

So it also works for multiple packages of the same type (for example 3 DRN-packages), by summing the budgets
---
 nlmod/mfoutput/binaryfile.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/nlmod/mfoutput/binaryfile.py b/nlmod/mfoutput/binaryfile.py
index a3197f0c..9d228a1b 100644
--- a/nlmod/mfoutput/binaryfile.py
+++ b/nlmod/mfoutput/binaryfile.py
@@ -138,13 +138,24 @@ def _get_binary_budget_data(kstpkper, fobj, text, column="q"):
         idx = np.array([idx])
 
     header = fobj.recordarray[idx]
-    ipos = fobj.iposarray[idx].item()
-    imeth = header["imeth"][0]
 
     t = header["text"][0]
     if isinstance(t, bytes):
         t = t.decode("utf-8")
 
+    data = []
+    for ipos in fobj.iposarray[idx]:
+        data.append(_get_binary_budget_record(fobj, ipos, header, column))
+
+    if len(data) == 1:
+        return data[0]
+    else:
+        return np.ma.sum(data, axis=0)
+
+
+def _get_binary_budget_record(fobj, ipos, header, column):
+    """Get a single data record from the budget file."""
+    imeth = header["imeth"][0]
     nlay = abs(header["nlay"][0])
     nrow = header["nrow"][0]
     ncol = header["ncol"][0]

From 33781bef8f34bd5f77436ca7edb82fb5b3b4ab29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Tue, 27 Feb 2024 10:05:50 +0100
Subject: [PATCH 6/6] Handle comments from @OnnoEbbens and replace grbfile by
 grb_file

Like in flopy
---
 nlmod/gwf/output.py          | 68 +++++++++++++++++++++---------------
 nlmod/mfoutput/mfoutput.py   | 20 +++++------
 tests/test_015_gwf_output.py | 32 ++++++++---------
 3 files changed, 66 insertions(+), 54 deletions(-)

diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py
index 460463b9..4b4e4ea2 100644
--- a/nlmod/gwf/output.py
+++ b/nlmod/gwf/output.py
@@ -14,13 +14,13 @@
     _get_heads_da,
     _get_time_index,
     _get_flopy_data_object,
-    _get_grbfile,
+    _get_grb_file,
 )
 
 logger = logging.getLogger(__name__)
 
 
-def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
+def get_headfile(ds=None, gwf=None, fname=None, grb_file=None):
     """Get flopy HeadFile object.
 
     Provide one of ds, gwf or fname.
@@ -33,7 +33,7 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
         groundwater flow model, by default None
     fname : str, optional
         path to heads file, by default None
-    grbfile : str
+    grb_file : str
         path to file containing binary grid information
 
     Returns
@@ -41,14 +41,14 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
     flopy.utils.HeadFile
         HeadFile object handle
     """
-    return _get_flopy_data_object("head", ds, gwf, fname, grbfile)
+    return _get_flopy_data_object("head", ds, gwf, fname, grb_file)
 
 
 def get_heads_da(
     ds=None,
     gwf=None,
     fname=None,
-    grbfile=None,
+    grb_file=None,
     delayed=False,
     chunked=False,
     **kwargs,
@@ -64,7 +64,7 @@ def get_heads_da(
         Flopy groundwaterflow object.
     fname : path, optional
         path to a binary heads file
-    grbfile : str, optional
+    grb_file : str, optional
         path to file containing binary grid information, only needed if reading
         output from file using fname
     delayed : bool, optional
@@ -77,7 +77,7 @@ def get_heads_da(
     da : xarray.DataArray
         heads data array.
     """
-    hobj = get_headfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile)
+    hobj = get_headfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file)
     # gwf.output.head() defaults to a structured grid
     if gwf is not None and ds is None and fname is None:
         kwargs["modelgrid"] = gwf.modelgrid
@@ -101,7 +101,7 @@ def get_heads_da(
     return da
 
 
-def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None):
+def get_cellbudgetfile(ds=None, gwf=None, fname=None, grb_file=None):
     """Get flopy CellBudgetFile object.
 
     Provide one of ds, gwf or fname.
@@ -113,9 +113,9 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None):
     gwf : flopy.mf6.ModflowGwf, optional
         groundwater flow model, by default None
     fname_cbc : str, optional
-        path to cell budget file, by default None\
-    grbfile : str, optional
-        path to file containing binary grid information, only needed if 
+        path to cell budget file, by default None
+    grb_file : str, optional
+        path to file containing binary grid information, only needed if
         fname_cbc is passed as only argument.
 
     Returns
@@ -123,7 +123,7 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None):
     flopy.utils.CellBudgetFile
         CellBudgetFile object handle
     """
-    return _get_flopy_data_object("budget", ds, gwf, fname, grbfile)
+    return _get_flopy_data_object("budget", ds, gwf, fname, grb_file)
 
 
 def get_budget_da(
@@ -131,7 +131,7 @@ def get_budget_da(
     ds=None,
     gwf=None,
     fname=None,
-    grbfile=None,
+    grb_file=None,
     column="q",
     delayed=False,
     chunked=False,
@@ -150,7 +150,7 @@ def get_budget_da(
     fname : path, optional
         specify the budget file to load, if not provided budget file will
         be obtained from ds or gwf.
-    grbfile : str
+    grb_file : str
         path to file containing binary grid information, only needed if reading
         output from file using fname
     column : str
@@ -166,7 +166,7 @@ def get_budget_da(
     da : xarray.DataArray
         budget data array.
     """
-    cbcobj = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile)
+    cbcobj = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file)
     da = _get_budget_da(cbcobj, text, column=column, **kwargs)
     da.attrs["units"] = "m3/d"
 
@@ -234,7 +234,7 @@ def get_gwl_from_wet_cells(head, layer="layer", botm=None):
     return gwl
 
 
-def get_flow_residuals(ds, kstpkper=None, grb_file=None):
+def get_flow_residuals(ds, gwf=None, fname=None, grb_file=None, kstpkper=None):
     """
     Get the flow residuals of a MODFLOW 6 simulation.
 
@@ -242,12 +242,17 @@ def get_flow_residuals(ds, kstpkper=None, grb_file=None):
     ----------
     ds : xarray.Dataset
         Xarray dataset with model data.
+    gwf : flopy ModflowGwf, optional
+        Flopy groundwaterflow object. One of ds or gwf must be provided.
+    fname : path, optional
+        specify the budget file to load, if not provided budget file will
+        be obtained from ds or gwf.
+    grb_file : str
+        The location of the grb-file. grb_file is determied from ds when None. The
+        default is None.
     kstpkper : tuple of 2 ints, optional
         The index of the timestep and the stress period to include in the result. Include
         all data in the budget-file when None. The default is None.
-    grb_file : str, optional
-        The location of the grb-file. grb_file is determied from ds when None. The
-        default is None.
 
     Returns
     -------
@@ -256,9 +261,9 @@ def get_flow_residuals(ds, kstpkper=None, grb_file=None):
 
     """
     if grb_file is None:
-        grb_file = _get_grbfile(ds)
+        grb_file = _get_grb_file(ds)
     grb = flopy.mf6.utils.MfGrdFile(grb_file)
-    cbf = get_cellbudgetfile(ds)
+    cbf = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file)
     dims = ds["botm"].dims
     coords = ds["botm"].coords
     flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper)
@@ -284,7 +289,9 @@ def get_flow_residuals(ds, kstpkper=None, grb_file=None):
     return da
 
 
-def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
+def get_flow_lower_face(
+    ds, gwf=None, fname=None, grb_file=None, kstpkper=None, lays=None
+):
     """
     Get the flow over the lower face of all model cells
 
@@ -295,12 +302,17 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
     ----------
     ds : xarray.Dataset
         Xarray dataset with model data.
-    kstpkper : tuple of 2 ints, optional
-        The index of the timestep and the stress period to include in the result. Include
-        all data in the budget-file when None. The default is None.
+    gwf : flopy ModflowGwf, optional
+        Flopy groundwaterflow object. One of ds or gwf must be provided.
+    fname : path, optional
+        specify the budget file to load, if not provided budget file will
+        be obtained from ds or gwf.
     grb_file : str, optional
         The location of the grb-file. grb_file is determied from ds when None. The
         default is None.
+    kstpkper : tuple of 2 ints, optional
+        The index of the timestep and the stress period to include in the result. Include
+        all data in the budget-file when None. The default is None.
     lays : int or list of ints, optional
         The layers to include in the result. When lays is None, all layers are included.
         The default is None.
@@ -312,12 +324,12 @@ def get_flow_lower_face(ds, kstpkper=None, grb_file=None, lays=None):
 
     """
     if grb_file is None:
-        grb_file = _get_grbfile(ds)
-    cbf = get_cellbudgetfile(ds)
+        grb_file = _get_grb_file(ds)
+    cbf = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file)
     flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper)
 
     if ds.gridtype == "vertex":
-        # detemine flf_index first
+        # determine flf_index first
         grb = flopy.mf6.utils.MfGrdFile(grb_file)
 
         if lays is None:
diff --git a/nlmod/mfoutput/mfoutput.py b/nlmod/mfoutput/mfoutput.py
index 9a1db186..d26fc73c 100644
--- a/nlmod/mfoutput/mfoutput.py
+++ b/nlmod/mfoutput/mfoutput.py
@@ -239,7 +239,7 @@ def _get_budget_da(
     return da
 
 
-def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
+def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grb_file=None):
     """Get modflow HeadFile or CellBudgetFile object, containg heads, budgets or
     concentrations
 
@@ -255,7 +255,7 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
         groundwater flow or transport model, by default None
     fname : str, optional
         path to Head- or CellBudgetFile, by default None
-    grbfile : str, optional
+    grb_file : str, optional
         path to file containing binary grid information, if None modelgrid
         information is obtained from ds. By default None
 
@@ -283,11 +283,11 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
             # return gwf.output.head(), gwf.output.budget() or gwt.output.concentration()
             return getattr(gwml.output, var)()
         fname = os.path.join(ds.model_ws, ds.model_name + extension)
-    if grbfile is None and ds is not None:
+    if grb_file is None and ds is not None:
         # get grb file
-        grbfile = _get_grbfile(ds)
-    if grbfile is not None and os.path.exists(grbfile):
-        modelgrid = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
+        grb_file = _get_grb_file(ds)
+    if grb_file is not None and os.path.exists(grb_file):
+        modelgrid = flopy.mf6.utils.MfGrdFile(grb_file).modelgrid
     elif ds is not None:
         modelgrid = modelgrid_from_ds(ds)
     else:
@@ -306,9 +306,9 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
         return flopy.utils.HeadFile(fname, text=var, modelgrid=modelgrid)
 
 
-def _get_grbfile(ds):
+def _get_grb_file(ds):
     if ds.gridtype == "vertex":
-        grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
+        grb_file = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
     elif ds.gridtype == "structured":
-        grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
-    return grbfile
+        grb_file = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
+    return grb_file
diff --git a/tests/test_015_gwf_output.py b/tests/test_015_gwf_output.py
index aafba45b..fc3608ee 100644
--- a/tests/test_015_gwf_output.py
+++ b/tests/test_015_gwf_output.py
@@ -67,17 +67,17 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"):
     assert np.array_equal(da.values, heads_correct, equal_nan=True)
 
     fname_hds = os.path.join(ds.model_ws, ds.model_name + ".hds")
-    grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
-    da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grbfile=grbfile)  # fname
+    grb_file = os.path.join(ds.model_ws, ds.model_name + ".dis.grb")
+    da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grb_file=grb_file)  # fname
     assert np.array_equal(da.values, heads_correct, equal_nan=True)
 
     # budget
     da = get_budget_da("CHD", ds=ds, gwf=None, fname=None)  # ds
     da = get_budget_da("CHD", ds=None, gwf=gwf, fname=None)  # gwf
     fname_cbc = os.path.join(ds.model_ws, ds.model_name + ".cbc")
-    get_budget_da("CHD", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile)  # fname
+    get_budget_da("CHD", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file)  # fname
     get_budget_da(
-        "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile
+        "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file
     )  # fname
 
     # unstructured
@@ -127,18 +127,18 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"):
     assert np.array_equal(da.values, heads_correct, equal_nan=True)
 
     fname_hds = os.path.join(ds.model_ws, ds.model_name + ".hds")
-    grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
-    da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grbfile=grbfile)  # fname
+    grb_file = os.path.join(ds.model_ws, ds.model_name + ".disv.grb")
+    da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grb_file=grb_file)  # fname
     assert np.array_equal(da.values, heads_correct, equal_nan=True)
 
     # budget
     da = get_budget_da("CHD", ds=ds_unstr, gwf=None, fname=None)  # ds
     da = get_budget_da("CHD", ds=None, gwf=gwf_unstr, fname=None)  # gwf
     da = get_budget_da(
-        "CHD", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile
+        "CHD", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file
     )  # fname
     _ = get_budget_da(
-        "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile
+        "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file
     )  # fname
 
 
@@ -150,8 +150,8 @@ def test_get_heads_da_from_file_structured_no_grb():
 
 def test_get_heads_da_from_file_structured_with_grb():
     fname_hds = "./tests/data/mf6output/structured/test.hds"
-    grbfile = "./tests/data/mf6output/structured/test.dis.grb"
-    nlmod.gwf.output.get_heads_da(fname=fname_hds, grbfile=grbfile)
+    grb_file = "./tests/data/mf6output/structured/test.dis.grb"
+    nlmod.gwf.output.get_heads_da(fname=fname_hds, grb_file=grb_file)
 
 
 def test_get_budget_da_from_file_structured_no_grb():
@@ -162,8 +162,8 @@ def test_get_budget_da_from_file_structured_no_grb():
 
 def test_get_budget_da_from_file_structured_with_grb():
     fname_cbc = "./tests/data/mf6output/structured/test.cbc"
-    grbfile = "./tests/data/mf6output/structured/test.dis.grb"
-    nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grbfile=grbfile)
+    grb_file = "./tests/data/mf6output/structured/test.dis.grb"
+    nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grb_file=grb_file)
 
 
 def test_get_heads_da_from_file_vertex_no_grb():
@@ -174,8 +174,8 @@ def test_get_heads_da_from_file_vertex_no_grb():
 
 def test_get_heads_da_from_file_vertex_with_grb():
     fname_hds = "./tests/data/mf6output/vertex/test.hds"
-    grbfile = "./tests/data/mf6output/vertex/test.disv.grb"
-    nlmod.gwf.output.get_heads_da(fname=fname_hds, grbfile=grbfile)
+    grb_file = "./tests/data/mf6output/vertex/test.disv.grb"
+    nlmod.gwf.output.get_heads_da(fname=fname_hds, grb_file=grb_file)
 
 
 def test_get_budget_da_from_file_vertex_no_grb():
@@ -186,8 +186,8 @@ def test_get_budget_da_from_file_vertex_no_grb():
 
 def test_get_budget_da_from_file_vertex_with_grb():
     fname_cbc = "./tests/data/mf6output/vertex/test.cbc"
-    grbfile = "./tests/data/mf6output/vertex/test.disv.grb"
-    nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grbfile=grbfile)
+    grb_file = "./tests/data/mf6output/vertex/test.disv.grb"
+    nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grb_file=grb_file)
 
 
 def test_postprocess_head():