Merge pull request #281 from gwmod/dev

Minor version update
gwmod · Nov 1, 2023 · e8b1ebb · e8b1ebb
2 parents bf21f20 + 4ad10b0
commit e8b1ebb
Show file tree

Hide file tree

Showing 10 changed files with 138 additions and 107 deletions.
diff --git a/README.md b/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)
 
@@ -49,6 +49,7 @@ Install the module with pip:
 * `matplotlib`
 * `dask`
 * `colorama`
+* `joblib`
 
 There are some optional dependecies, only needed (and imported) in a single method.
 Examples of this are `bottleneck` (used in calculate_gxg), `geocube` (used in

diff --git a/docs/conf.py b/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,

diff --git a/docs/getting_started.rst b/docs/getting_started.rst
@@ -117,6 +117,7 @@ potential solutions.
 - matplotlib
 - dask
 - colorama
+- joblib
 
 On top of that there are some optional dependecies:
 
@@ -126,4 +127,4 @@ On top of that there are some optional dependecies:
 - scikit-image (used in calculate_sea_coverage)
 
 These dependencies are only needed (and imported) in a single method or function.
-They can be installed using ``pip install nlmod[full]`` or ``pip install -e .[full]``.
+They can be installed using ``pip install nlmod[full]`` or ``pip install -e .[full]``.
diff --git a/nlmod/gwf/output.py b/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(

diff --git a/nlmod/gwf/recharge.py b/nlmod/gwf/recharge.py
@@ -33,11 +33,11 @@ def ds_to_rch(
         data array containing mask, recharge is only added where mask is True
     pname : str, optional
         package name. The default is 'rch'.
-    auxiliary : str or list of str
-        name(s) of data arrays to include as auxiliary data to reclist
     recharge : str, optional
         The name of the variable in ds that contains the recharge flux rate. The default
         is "recharge".
+    auxiliary : str or list of str
+        name(s) of data arrays to include as auxiliary data to reclist
 
     Returns
     -------
@@ -85,7 +85,7 @@ def ds_to_rch(
         **kwargs,
     )
     if (auxiliary is not None) and (ds.transport == 1):
-        logger.info("-> adding GHB to SSM sources list")
+        logger.info("-> adding RCH to SSM sources list")
         ssm_sources = list(ds.attrs["ssm_sources"])
         if rch.package_name not in ssm_sources:
             ssm_sources += [rch.package_name]
@@ -107,6 +107,7 @@ def ds_to_evt(
     nseg=1,
     surface=None,
     depth=None,
+    auxiliary=None,
     **kwargs,
 ):
     """Convert the evaporation data in the model dataset to a evt package with
@@ -133,6 +134,8 @@ def ds_to_evt(
     depth : str, float or xr.DataArray, optional
         The ET extinction depth. Set to 1 meter (below surface) when None. The default
         is None.
+    auxiliary : str or list of str
+        name(s) of data arrays to include as auxiliary data to reclist
     **kwargs : TYPE
         DESCRIPTION.
 
@@ -186,6 +189,7 @@ def ds_to_evt(
         col3=depth,
         first_active_layer=True,
         only_active_cells=False,
+        aux=auxiliary,
     )
 
     # create rch package
@@ -194,12 +198,20 @@ def ds_to_evt(
         filename=f"{gwf.name}.evt",
         pname=pname,
         fixed_cell=False,
+        auxiliary="CONCENTRATION" if auxiliary is not None else None,
         maxbound=len(spd),
         nseg=nseg,
         stress_period_data={0: spd},
         **kwargs,
     )
 
+    if (auxiliary is not None) and (ds.transport == 1):
+        logger.info("-> adding EVT to SSM sources list")
+        ssm_sources = list(ds.attrs["ssm_sources"])
+        if evt.package_name not in ssm_sources:
+            ssm_sources += [evt.package_name]
+            ds.attrs["ssm_sources"] = ssm_sources
+
     if use_ts:
         # create timeseries packages
         _add_time_series(evt, evt_unique_dic, ds)

diff --git a/nlmod/gwt/output.py b/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)