fix, #3

examples/03_local_grid_refinement.ipynb

@@ -265,8 +265,8 @@
     "\n",
     "\n",
     "# add constant head cells at model boundaries\n",
-    "chd = nlmod.mfpackages.chd_at_model_edge_from_model_ds(\n",
-    "    model_ds, gwf, head='starting_head')"
+    "model_ds.update(nlmod.mfpackages.constant_head.get_chd_at_model_edge(model_ds, model_ds['idomain']))    \n",
+    "chd = nlmod.mfpackages.chd_from_model_ds(model_ds, gwf, head='starting_head')"
    ]
   },
   {

nlmod/mfpackages/__init__.py

@@ -1,2 +1,2 @@
-from . import recharge, surface_water
+from . import recharge, surface_water, constant_head
 from .mfpackages import *

nlmod/mfpackages/constant_head.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Dec  3 11:58:49 2021
+
+@author: oebbe
+"""
+
+import xarray as xr
+import numpy as np
+
+from .. import util, cache
+
+@cache.cache_netcdf
+def get_chd_at_model_edge(model_ds, idomain):
+    """get data array which is 1 at every active cell (defined by idomain) at 
+    the boundaries of the model (xmin, xmax, ymin, ymax). Other cells are 0.
+
+    Parameters
+    ----------
+    model_ds : xarray.Dataset
+        dataset with model data.
+    idomain : xarray.DataArray
+        idomain used to get active cells and shape of DataArray
+        
+    Returns
+    -------
+    model_ds_out : xarray.Dataset
+        dataset with chd array
+    """
+
+    # add constant head cells at model boundaries
+
+    # get mask with grid edges
+    xmin = model_ds['x'] == model_ds['x'].min()
+    xmax = model_ds['x'] == model_ds['x'].max()
+    ymin = model_ds['y'] == model_ds['y'].min()
+    ymax = model_ds['y'] == model_ds['y'].max()
+
+    model_ds_out = util.get_model_ds_empty(model_ds)
+
+    if model_ds.gridtype == 'structured':
+        mask2d = (ymin | ymax | xmin | xmax)
+
+        # assign 1 to cells that are on the edge and have an active idomain
+        model_ds_out['chd'] = xr.zeros_like(idomain)
+        for lay in model_ds.layer:
+            model_ds_out['chd'].loc[lay] = np.where(
+                mask2d & (idomain.loc[lay] == 1), 1, 0)
+
+    elif model_ds.gridtype == 'unstructured':
+        mask = np.where([xmin | xmax | ymin | ymax])[1]
+
+        # assign 1 to cells that are on the edge, have an active idomain
+        model_ds_out['chd'] = xr.zeros_like(idomain)
+        model_ds_out['chd'].loc[:, mask] = 1
+        model_ds_out['chd'] = xr.where(idomain == 1,
+                                       model_ds_out['chd'], 0)
+
+
+    return model_ds_out

nlmod/mfpackages/mfpackages.py

@@ -336,7 +336,8 @@ def sto_from_model_ds(model_ds, gwf,
         return sto
 
 
-def chd_at_model_edge_from_model_ds(model_ds, gwf, head='starting_head'):
+def chd_from_model_ds(model_ds, gwf, chd='chd', 
+                      head='starting_head'):
     """get constant head boundary at the model's edges from the model dataset.
 
     Parameters
@@ -345,6 +346,9 @@ def chd_at_model_edge_from_model_ds(model_ds, gwf, head='starting_head'):
         dataset with model data.
     gwf : flopy ModflowGwf
         groundwaterflow object.
+    chd : str, optional
+        name of data variable in model_ds that is 1 for cells with a constant
+        head and zero for all other cells. The default is 'chd'.
     head : str, optional
         name of data variable in model_ds that is used as the head in the chd
         cells. The default is 'starting_head'.
@@ -354,43 +358,18 @@ def chd_at_model_edge_from_model_ds(model_ds, gwf, head='starting_head'):
     chd : flopy ModflowGwfchd
         chd package
     """
-    # add constant head cells at model boundaries
-
-    # get mask with grid edges
-    xmin = model_ds['x'] == model_ds['x'].min()
-    xmax = model_ds['x'] == model_ds['x'].max()
-    ymin = model_ds['y'] == model_ds['y'].min()
-    ymax = model_ds['y'] == model_ds['y'].max()
-
+    # get the stress_period_data
     if model_ds.gridtype == 'structured':
-        mask2d = (ymin | ymax | xmin | xmax)
-
-        # assign 1 to cells that are on the edge and have an active idomain
-        model_ds['chd'] = xr.zeros_like(model_ds['idomain'])
-        for lay in model_ds.layer:
-            model_ds['chd'].loc[lay] = np.where(
-                mask2d & (model_ds['idomain'].loc[lay] == 1), 1, 0)
-
-        # get the stress_period_data
         chd_rec = mdims.data_array_3d_to_rec_list(model_ds,
-                                                  model_ds['chd'] != 0,
+                                                  model_ds[chd] != 0,
                                                   col1=head)
     elif model_ds.gridtype == 'unstructured':
-        mask = np.where([xmin | xmax | ymin | ymax])[1]
-
-        # assign 1 to cells that are on the edge, have an active idomain
-        model_ds['chd'] = xr.zeros_like(model_ds['idomain'])
-        model_ds['chd'].loc[:, mask] = 1
-        model_ds['chd'] = xr.where(model_ds['idomain'] == 1,
-                                   model_ds['chd'], 0)
-
-        # get the stress_period_data
-        cellids = np.where(model_ds['chd'])
+        cellids = np.where(model_ds[chd])
         chd_rec = list(zip(zip(cellids[0],
                                cellids[1]),
                            [1.0] * len(cellids[0])))
 
-    chd = flopy.mf6.ModflowGwfchd(gwf, pname='chd',
+    chd = flopy.mf6.ModflowGwfchd(gwf, pname=chd,
                                   maxbound=len(chd_rec),
                                   stress_period_data=chd_rec,
                                   save_flows=True)

tests/test_001_model.py

@@ -191,8 +191,8 @@ def test_create_sea_model(tmpdir):
     nlmod.mfpackages.surface_drain_from_model_ds(model_ds, gwf)
 
     # add constant head cells at model boundaries
-    nlmod.mfpackages.chd_at_model_edge_from_model_ds(
-        model_ds, gwf, head='starting_head')
+    model_ds.update(nlmod.mfpackages.constant_head.get_chd_at_model_edge(model_ds, model_ds['idomain']))    
+    nlmod.mfpackages.chd_from_model_ds(model_ds, gwf, head='starting_head')
 
     # add knmi recharge to the model datasets
     model_ds.update(nlmod.read.knmi.add_knmi_to_model_dataset(model_ds))
@@ -262,8 +262,8 @@ def test_create_sea_model_perlen_list(tmpdir):
     nlmod.mfpackages.surface_drain_from_model_ds(model_ds, gwf)
 
     # add constant head cells at model boundaries
-    nlmod.mfpackages.chd_at_model_edge_from_model_ds(
-        model_ds, gwf, head='starting_head')
+    model_ds.update(nlmod.mfpackages.constant_head.get_chd_at_model_edge(model_ds, model_ds['idomain']))    
+    nlmod.mfpackages.chd_from_model_ds(model_ds, gwf, head='starting_head')
 
     # add knmi recharge to the model datasets
     model_ds.update(nlmod.read.knmi.add_knmi_to_model_dataset(model_ds))
@@ -325,8 +325,8 @@ def test_create_sea_model_perlen_14(tmpdir):
     nlmod.mfpackages.surface_drain_from_model_ds(model_ds, gwf)
 
     # add constant head cells at model boundaries
-    nlmod.mfpackages.chd_at_model_edge_from_model_ds(
-        model_ds, gwf, head='starting_head')
+    model_ds.update(nlmod.mfpackages.constant_head.get_chd_at_model_edge(model_ds, model_ds['idomain']))    
+    nlmod.mfpackages.chd_from_model_ds(model_ds, gwf, head='starting_head')
 
     # add knmi recharge to the model datasets
     model_ds.update(nlmod.read.knmi.add_knmi_to_model_dataset(model_ds))

tests/test_003_mfpackages.py

@@ -113,6 +113,8 @@ def chd_from_model_ds(tmpdir):
     nlmod.mfpackages.ic_from_model_ds(model_ds, gwf,
                                       starting_head=1.0)
 
-    chd = nlmod.mfpackages.chd_at_model_edge_from_model_ds(model_ds, gwf)
+    # add constant head cells at model boundaries
+    model_ds.update(nlmod.mfpackages.constant_head.get_chd_at_model_edge(model_ds, model_ds['idomain']))    
+    chd = nlmod.mfpackages.chd_from_model_ds(model_ds, gwf, head='starting_head')
 
     return chd