replace three methods by _get_flopy_data_object (#278)

* replace three methods by _get_flopy_data_object

* Fix bug, imporve documentation and simplify code

* Update mfoutput.py

* only error for budgets without grid information

* Last improvements of documentation

* Update pyproject.toml

* some more replacements of ArtesiaWater to gwmod

* Update codacy badges

* fix docs  typo

---------

Co-authored-by: OnnoEbbens <onnoebbens@gmail.com>

README.md

@@ -2,9 +2,9 @@
 
 <img src="docs/_static/logo_10000_2.png" width="256"/>
 
-[![nlmod](https://github.com/ArtesiaWater/nlmod/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/ArtesiaWater/nlmod/actions/workflows/ci.yml)
-[![Codacy Badge](https://app.codacy.com/project/badge/Grade/6fadea550ea04ea28b6ccde88fc56f35)](https://www.codacy.com/gh/ArtesiaWater/nlmod/dashboard?utm_source=github.com&utm_medium=referral&utm_content=ArtesiaWater/nlmod&utm_campaign=Badge_Grade)
-[![Codacy Badge](https://app.codacy.com/project/badge/Coverage/6fadea550ea04ea28b6ccde88fc56f35)](https://www.codacy.com/gh/ArtesiaWater/nlmod/dashboard?utm_source=github.com&utm_medium=referral&utm_content=ArtesiaWater/nlmod&utm_campaign=Badge_Coverage)
+[![nlmod](https://github.com/gwmod/nlmod/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/gwmod/nlmod/actions/workflows/ci.yml)
+[![Codacy Badge](https://app.codacy.com/project/badge/Grade/f1797b66e98b42b294bc1c5fc233dbf3)](https://app.codacy.com/gh/gwmod/nlmod/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade)
+[![Codacy Badge](https://app.codacy.com/project/badge/Coverage/f1797b66e98b42b294bc1c5fc233dbf3)](https://app.codacy.com/gh/gwmod/nlmod/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_coverage)
 [![PyPI version](https://badge.fury.io/py/nlmod.svg)](https://badge.fury.io/py/nlmod)
 [![Documentation Status](https://readthedocs.org/projects/nlmod/badge/?version=stable)](https://nlmod.readthedocs.io/en/stable/?badge=stable)
 

docs/conf.py

@@ -77,7 +77,7 @@
     "navigation_depth": 4,
     "includehidden": True,
     "titles_only": False,
-    # "github_url": "https://github.com/ArtesiaWater/nlmod",
+    # "github_url": "https://github.com/gwmod/nlmod",
 }
 
 # Add any paths that contain custom static files (such as style sheets) here,

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
     ----------
@@ -33,34 +36,10 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None):
 
     Returns
     -------
-    headobj : flopy.utils.HeadFile
+    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(
@@ -93,7 +72,7 @@ def get_heads_da(
 
     Returns
     -------
-    head_da : xarray.DataArray
+    da : xarray.DataArray
         heads data array.
     """
     hobj = get_headfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile)
@@ -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
-        CellBudgetFile object
+    flopy.utils.CellBudgetFile
+        CellBudgetFile object handle
     """
-    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,42 +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)
-    return concobj
+    Provide one of ds, gwf or fname.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset, optional
+        model dataset, by default None
+    gwt : flopy.mf6.ModflowGwt, optional
+        groundwater transport model, by default None
+    fname : str, optional
+        path to heads file, by default None
+    grbfile : str
+        path to file containing binary grid information
+
+    Returns
+    -------
+    flopy.utils.HeadFile
+        HeadFile object handle
+    """
+    return _get_flopy_data_object("concentration", ds, gwt, fname, grbfile)
 
 
 def get_concentration_da(
@@ -69,7 +63,7 @@ def get_concentration_da(
 
     Returns
     -------
-    conc_da : xarray.DataArray
+    da : xarray.DataArray
         concentration data array.
     """
     cobj = get_concentration_obj(ds=ds, gwt=gwt, fname=fname, grbfile=grbfile)

nlmod/mfoutput/mfoutput.py

@@ -1,9 +1,13 @@
+import os
 import logging
+import warnings
 
 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 +221,71 @@ 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
+    gwml : flopy.mf6.ModflowGwf or flopy.mf6.ModflowGwt, optional
+        groundwater flow or transport model, by default None
+    fname : str, optional
+        path to Head- or CellBudgetFile, by default None
+    grbfile : str, optional
+        path to file containing binary grid information, if None modelgrid
+        information is obtained from ds. By default None
+
+    Returns
+    -------
+    flopy.utils.HeadFile or flopy.utils.CellBudgetFile
+    """
+    if var == "head":
+        ml_name = "gwf"
+        extension = ".hds"
+    elif var == "budget":
+        ml_name = "gwf"
+        extension = ".cbc"
+    elif var == "concentration":
+        ml_name = "gwt"
+        extension = "_gwt.ucn"
+    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 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")
+    if grbfile is not None and os.path.exists(grbfile):
+        modelgrid = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid
+    elif ds is not None:
+        modelgrid = modelgrid_from_ds(ds)
+    else:
+        modelgrid = None
+
+    msg = f"Cannot create {var} data-array without grid information."
+    if var == "budget":
+        if modelgrid is None:
+            logger.error(msg)
+            raise ValueError(msg)
+        return flopy.utils.CellBudgetFile(fname, modelgrid=modelgrid)
+    else:
+        if modelgrid is None:
+            logger.warning(msg)
+            warnings.warn(msg)
+        return flopy.utils.HeadFile(fname, text=var, modelgrid=modelgrid)

pyproject.toml

@@ -51,8 +51,8 @@ classifiers = [
 ]
 
 [project.urls]
-homepage = "https://github.com/ArtesiaWater/nlmod"
-repository = "https://github.com/ArtesiaWater/nlmod"
+homepage = "https://github.com/gwmod/nlmod"
+repository = "https://github.com/gwmod/nlmod"
 documentation = "https://nlmod.readthedocs.io/en/latest/"
 
 [project.optional-dependencies]

tests/test_015_gwf_output.py

@@ -12,10 +12,7 @@
 tmpdir = tempfile.gettempdir()
 tst_model_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")
 
-grberror = (
-    "Please provide grid information by passing path to the "
-    "binary grid file with `grbfile=<path to file>`."
-)
+grberror = "Cannot create budget data-array without grid information."
 
 
 def test_create_small_model_grid_only(tmpdir, model_name="test"):