replace three methods by _get_flopy_data_object

nlmod/gwf/output.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import warnings
 
 import flopy
@@ -9,16 +8,20 @@
 from shapely.geometry import Point
 
 from ..dims.grid import modelgrid_from_ds
-from ..mfoutput.mfoutput import _get_budget_da, _get_heads_da, _get_time_index
+from ..mfoutput.mfoutput import (
+    _get_budget_da,
+    _get_heads_da,
+    _get_time_index,
+    _get_flopy_data_object,
+)
 
 logger = logging.getLogger(__name__)
 
 
 def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
-    """Get modflow HeadFile object.
+    """Get flopy HeadFile object.
 
-    Provide one of ds, gwf or fname_hds. Not that it really matters but if
-    all are provided hierarchy is as follows: fname_hds > ds > gwf
+    Provide one of ds, gwf or fname.
 
     Parameters
     ----------
@@ -36,31 +39,7 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
     headobj : flopy.utils.HeadFile
         HeadFile object handle
     """
-    msg = "Load the heads using either the ds, gwf or fname_hds"
-    assert ((ds is not None) + (gwf is not None) + (fname is not None)) >= 1, msg
-
-    if fname is None:
-        if ds is None:
-            return gwf.output.head()
-        else:
-            fname = os.path.join(ds.model_ws, ds.model_name + ".hds")
-            # 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")
-            else:
-                grbfile = None
-
-    if fname is not None:
-        if grbfile is not None:
-            mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
-        else:
-            logger.warning(msg)
-            warnings.warn(msg)
-            mg = None
-        headobj = flopy.utils.HeadFile(fname, modelgrid=mg)
-    return headobj
+    return _get_flopy_data_object("head", ds, gwf, fname, grbfile)
 
 
 def get_heads_da(
@@ -121,10 +100,9 @@ def get_heads_da(
 
 
 def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None):
-    """Get modflow CellBudgetFile object.
+    """Get flopy CellBudgetFile object.
 
-    Provide one of ds, gwf or fname_cbc. Not that it really matters but if
-    all are provided hierarchy is as follows: fname_cbc > ds > gwf
+    Provide one of ds, gwf or fname.
 
     Parameters
     ----------
@@ -140,35 +118,10 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None):
 
     Returns
     -------
-    cbc : flopy.utils.CellBudgetFile
+    cbc : flopy.utils.HeadFile, flopy.utils.CellBudgetFile or 
         CellBudgetFile object
     """
-    msg = "Load the budgets using either the ds or the gwf"
-    assert ((ds is not None) + (gwf is not None) + (fname is not None)) == 1, msg
-
-    if fname is None:
-        if ds is None:
-            return gwf.output.budget()
-        else:
-            fname = os.path.join(ds.model_ws, ds.model_name + ".cbc")
-            # 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")
-            else:
-                grbfile = None
-    if fname is not None:
-        if grbfile is not None:
-            mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
-        else:
-            logger.error("Cannot create budget data-array without grid information.")
-            raise ValueError(
-                "Please provide grid information by passing path to the "
-                "binary grid file with `grbfile=<path to file>`."
-            )
-        cbc = flopy.utils.CellBudgetFile(fname, modelgrid=mg)
-    return cbc
+    return _get_flopy_data_object("budget", ds, gwf, fname, grbfile)
 
 
 def get_budget_da(

nlmod/gwt/output.py

@@ -1,41 +1,36 @@
 import logging
-import os
-import warnings
 
-import flopy
 import numpy as np
 import xarray as xr
 
 from ..dims.layers import calculate_thickness
-from ..mfoutput.mfoutput import _get_heads_da, _get_time_index
+from ..mfoutput.mfoutput import _get_heads_da, _get_time_index, _get_flopy_data_object
 
 logger = logging.getLogger(__name__)
 
 
 def get_concentration_obj(ds=None, gwt=None, fname=None, grbfile=None):
-    msg = "Load the concentration using either the ds, gwt or a fname_conc"
-    assert ((ds is not None) + (gwt is not None) + (fname is not None)) == 1, msg
+    """Get flopy HeadFile object connected to the file with the concetration of cells.
 
-    if fname is None:
-        if ds is None:
-            return gwt.output.concentration()
-        else:
-            fname = os.path.join(ds.model_ws, f"{ds.model_name}_gwt.ucn")
-            # 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")
-            else:
-                grbfile = None
-    if fname is not None:
-        if grbfile is not None:
-            mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
-        else:
-            logger.warning(msg)
-            warnings.warn(msg)
-            mg = None
-        concobj = flopy.utils.HeadFile(fname, text="concentration", modelgrid=mg)
+    Provide one of ds, gwf or fname.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset, optional
+        model dataset, by default None
+    gwf : flopy.mf6.ModflowGwf, optional
+        groundwater flow model, by default None
+    fname : str, optional
+        path to heads file, by default None
+    grbfile : str
+        path to file containing binary grid information
+
+    Returns
+    -------
+    headobj : flopy.utils.HeadFile
+        HeadFile object handle
+    """
+    concobj = _get_flopy_data_object("concentration", ds, gwt, fname, grbfile)
     return concobj
 
 

nlmod/mfoutput/mfoutput.py

@@ -1,9 +1,12 @@
+import os
 import logging
 
 import dask
 import xarray as xr
 
-from ..dims.grid import get_dims_coords_from_modelgrid
+import flopy
+
+from ..dims.grid import get_dims_coords_from_modelgrid, modelgrid_from_ds
 from ..dims.resample import get_affine_mod_to_world
 from ..dims.time import ds_time_idx
 from .binaryfile import _get_binary_budget_data, _get_binary_head_data
@@ -217,3 +220,72 @@ def _get_budget_da(
     da = _create_da(stacked_arr, modelgrid, cbcobj.get_times())
 
     return da
+
+
+def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None):
+    """Get modflow HeadFile or CellBudgetFile object, containg heads, budgets or
+    concentrations
+
+    Provide one of ds, gwf or fname.
+
+    Parameters
+    ----------
+    var : str
+        The name of the variable. Can be 'head', 'budget' or 'concentration'.
+    ds : xarray.Dataset, optional
+        model dataset, by default 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 
+        fname_cbc is passed as only argument.
+
+    Returns
+    -------
+    flopy.utils.HeadFile or flopy.utils.CellBudgetFile
+    """
+    if var == "head":
+        ml_name = "gwf"
+        extension = ".hds"
+        flopy_class = flopy.utils.HeadFile
+    elif var == "budget":
+        ml_name = "gwf"
+        extension = ".cbc"
+        flopy_class = flopy.utils.CellBudgetFile
+    elif var == "concentration":
+        ml_name = "gwt"
+        extension = "_gwt.ucn"
+        flopy_class = flopy.utils.HeadFile
+    else:
+        raise (ValueError(f"Unknown variable {var}"))
+    msg = f"Load the {var}s using either ds, {ml_name} or fname"
+    assert ((ds is not None) + (gwml is not None) + (fname is not None)) == 1, msg
+    if fname is None:
+        if ds is 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:
+        # 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")
+        else:
+            grbfile = None
+    if grbfile is not None and os.path.exists(grbfile):
+        mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
+    elif ds is not None:
+        mg = modelgrid_from_ds(ds)
+    else:
+        logger.error(f"Cannot create {var} data-array without grid information.")
+        raise ValueError(
+            "Please provide grid information by passing path to the "
+            "binary grid file with `grbfile=<path to file>`."
+        )
+    if isinstance(flopy_class, flopy.utils.HeadFile):
+        return flopy_class(fname, modelgrid=mg, text=var)
+    else:
+        return flopy_class(fname, modelgrid=mg)