fix(gridgen): fix add_refinement_feature() shapefile support (#2022)

* previously erroneously added ".shp" to caller-provided shapefile names
* raise error if features is not a filename or list of geometries
* correct/expand docstring
* add test case

autotest/test_gridgen.py

@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from shutil import which
+from shutil import copyfile, which
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -11,6 +11,7 @@
 from modflow_devtools.misc import has_pkg
 
 import flopy
+from flopy.discretization.vertexgrid import VertexGrid
 from flopy.utils.gridgen import Gridgen
 
 
@@ -25,6 +26,67 @@ def test_ctor_accepts_path_or_string(function_tmpdir):
     assert g.model_ws == function_tmpdir
 
 
+def base_grid():
+    Lx = 100.0
+    Ly = 100.0
+    nlay = 2
+    nrow = 51
+    ncol = 51
+    delr = Lx / ncol
+    delc = Ly / nrow
+    h0 = 10
+    h1 = 5
+    top = h0
+    botm = np.zeros((nlay, nrow, ncol), dtype=np.float32)
+    botm[1, :, :] = -10.0
+    ms = flopy.modflow.Modflow(rotation=-20.0)
+    dis = flopy.modflow.ModflowDis(
+        ms,
+        nlay=nlay,
+        nrow=nrow,
+        ncol=ncol,
+        delr=delr,
+        delc=delc,
+        top=top,
+        botm=botm,
+    )
+    return ms.modelgrid
+
+
+@requires_exe("gridgen")
+def test_add_refinement_feature_accepts_filename(function_tmpdir):
+    # same as 1st grid in .docs/Notebooks/gridgen_example.py
+    grid = base_grid()
+
+    # create initial refined grid (writing refinement
+    # feature shapefile to the workspace)
+    g1 = Gridgen(grid, model_ws=function_tmpdir)
+    g1.add_refinement_features(
+        [[[(0, 0), (0, 60), (40, 80), (60, 0), (0, 0)]]],
+        "polygon",
+        1,
+        range(grid.nlay),
+    )
+    g1.build()
+    g1 = VertexGrid(**g1.get_gridprops_vertexgrid())
+    # g1.plot()
+    # g1.show()
+
+    # recreate gridgen object, using shapefile from the
+    # first run
+    g2 = Gridgen(grid, model_ws=function_tmpdir)
+    g2.add_refinement_features("rf0.shp", "polygon", 1, range(grid.nlay))
+    g2.build()
+    g2 = VertexGrid(**g2.get_gridprops_vertexgrid())
+    # g2.plot()
+    # plt.show()
+
+    assert g1.nnodes > grid.nnodes
+    assert g1.ncpl != grid.ncpl
+    assert g1.ncpl == g2.ncpl
+    assert g1.nnodes == g2.nnodes
+
+
 @pytest.mark.slow
 @requires_exe("mf6", "gridgen")
 @requires_pkg("shapely")

autotest/test_mf6.py

@@ -257,42 +257,44 @@ def to_os_sep(s):
 def test_load_and_run_sim_when_namefile_uses_filenames(
     function_tmpdir, example_data_path
 ):
-    ws = function_tmpdir / "ws"
-    ml_name = "freyberg"
-    nam_name = "mfsim.nam"
-    nam_path = ws / nam_name
-    copytree(example_data_path / f"mf6-{ml_name}", ws)
+    # copy model input files to temp workspace
+    model_name = "mf6-freyberg"
+    workspace = function_tmpdir / model_name
+    copytree(example_data_path / model_name, workspace)
 
-    sim = MFSimulation.load(nam_name, sim_ws=ws)
+    # load, check and run simulation
+    sim = MFSimulation.load(sim_ws=workspace)
     sim.check()
-    success, buff = sim.run_simulation(report=True)
+    success, _ = sim.run_simulation(report=True)
     assert success
 
 
 @requires_exe("mf6")
 def test_load_and_run_sim_when_namefile_uses_abs_paths(
     function_tmpdir, example_data_path
 ):
-    ws = function_tmpdir / "ws"
-    ml_name = "freyberg"
-    nam_name = "mfsim.nam"
-    nam_path = ws / nam_name
-    copytree(example_data_path / f"mf6-{ml_name}", ws)
-
-    with set_dir(ws):
+    # copy model input files to temp workspace
+    model_name = "freyberg"
+    workspace = function_tmpdir / "ws"
+    copytree(example_data_path / f"mf6-{model_name}", workspace)
+
+    # sub abs paths into namefile
+    with set_dir(workspace):
+        nam_path = workspace / "mfsim.nam"
         lines = open(nam_path).readlines()
         with open(nam_path, "w") as f:
             for l in lines:
-                pattern = f"{ml_name}."
+                pattern = f"{model_name}."
                 if pattern in l:
                     l = l.replace(
-                        pattern, str(ws.absolute()) + os.sep + pattern
+                        pattern, str(workspace.absolute()) + os.sep + pattern
                     )
                 f.write(l)
 
-    sim = MFSimulation.load(nam_name, sim_ws=ws)
+    # load, check and run simulation
+    sim = MFSimulation.load(sim_ws=workspace)
     sim.check()
-    success, buff = sim.run_simulation(report=True)
+    success, _ = sim.run_simulation(report=True)
     assert success
 
 
@@ -301,40 +303,53 @@ def test_load_and_run_sim_when_namefile_uses_abs_paths(
 def test_load_sim_when_namefile_uses_rel_paths(
     function_tmpdir, example_data_path, sep
 ):
-    ws = function_tmpdir / "ws"
-    ml_name = "freyberg"
-    nam_name = "mfsim.nam"
-    nam_path = ws / nam_name
-    copytree(example_data_path / f"mf6-{ml_name}", ws)
-
-    with set_dir(ws):
+    # copy model input files to temp workspace
+    model_name = "freyberg"
+    workspace = function_tmpdir / "ws"
+    copytree(example_data_path / f"mf6-{model_name}", workspace)
+
+    # sub rel paths into namefile
+    with set_dir(workspace):
+        nam_path = workspace / "mfsim.nam"
         lines = open(nam_path).readlines()
         with open(nam_path, "w") as f:
             for l in lines:
-                pattern = f"{ml_name}."
+                pattern = f"{model_name}."
                 if pattern in l:
                     if sep == "win":
                         l = to_win_sep(
                             l.replace(
-                                pattern, "../" + ws.name + "/" + ml_name + "."
+                                pattern,
+                                "../"
+                                + workspace.name
+                                + "/"
+                                + model_name
+                                + ".",
                             )
                         )
                     else:
                         l = to_posix_sep(
                             l.replace(
-                                pattern, "../" + ws.name + "/" + ml_name + "."
+                                pattern,
+                                "../"
+                                + workspace.name
+                                + "/"
+                                + model_name
+                                + ".",
                             )
                         )
                 f.write(l)
 
-    sim = MFSimulation.load(nam_name, sim_ws=ws)
+    # load and check simulation
+    sim = MFSimulation.load(sim_ws=workspace)
     sim.check()
 
     # don't run simulation with Windows sep on Linux or Mac
     if sep == "win" and platform.system() != "Windows":
         return
 
-    success, buff = sim.run_simulation(report=True)
+    # run simulation
+    success, _ = sim.run_simulation(report=True)
     assert success
 
 
@@ -343,36 +358,49 @@ def test_load_sim_when_namefile_uses_rel_paths(
 def test_write_simulation_always_writes_posix_path_separators(
     function_tmpdir, example_data_path, sep
 ):
-    ws = function_tmpdir / "ws"
-    ml_name = "freyberg"
-    nam_name = "mfsim.nam"
-    nam_path = ws / nam_name
-    copytree(example_data_path / f"mf6-{ml_name}", ws)
-
-    with set_dir(ws):
+    # copy model input files to temp workspace
+    model_name = "freyberg"
+    workspace = function_tmpdir / "ws"
+    copytree(example_data_path / f"mf6-{model_name}", workspace)
+
+    # use OS-specific path separators
+    with set_dir(workspace):
+        nam_path = workspace / "mfsim.nam"
         lines = open(nam_path).readlines()
         with open(nam_path, "w") as f:
             for l in lines:
-                pattern = f"{ml_name}."
+                pattern = f"{model_name}."
                 if pattern in l:
                     if sep == "win":
                         l = to_win_sep(
                             l.replace(
-                                pattern, "../" + ws.name + "/" + ml_name + "."
+                                pattern,
+                                "../"
+                                + workspace.name
+                                + "/"
+                                + model_name
+                                + ".",
                             )
                         )
                     else:
                         l = to_posix_sep(
                             l.replace(
-                                pattern, "../" + ws.name + "/" + ml_name + "."
+                                pattern,
+                                "../"
+                                + workspace.name
+                                + "/"
+                                + model_name
+                                + ".",
                             )
                         )
                 f.write(l)
 
-    sim = MFSimulation.load(nam_name, sim_ws=ws)
+    # load and write simulation
+    sim = MFSimulation.load(sim_ws=workspace)
     sim.write_simulation()
 
-    lines = open(ws / "mfsim.nam").readlines()
+    # make sure posix separators were written
+    lines = open(workspace / "mfsim.nam").readlines()
     assert all("\\" not in l for l in lines)
 
 

flopy/utils/gridgen.py

@@ -406,9 +406,9 @@ def add_refinement_features(self, features, featuretype, level, layers):
         """
         Parameters
         ----------
-        features : str, list, or collection object
+        features : str or array-like
             features can be
-                a string containing the name of a shapefile
+                a str, the name of a shapefile in the workspace
                 a list of points, lines, or polygons
                 flopy.utils.geometry.Collection object
                 a list of flopy.utils.geometry objects
@@ -435,15 +435,19 @@ def add_refinement_features(self, features, featuretype, level, layers):
 
         # Create shapefile or set shapefile to feature
         rfname = f"rf{len(self._rfdict)}"
-        if isinstance(features, list):
+        if isinstance(features, str):
+            shapefile = features
+        elif isinstance(features, (list, tuple, np.ndarray)):
             rfname_w_path = os.path.join(self.model_ws, rfname)
             features_to_shapefile(features, featuretype, rfname_w_path)
-            shapefile = rfname
+            shapefile = f"{rfname}.shp"
         else:
-            shapefile = features
+            raise ValueError(
+                f"Features must be str (shapefile name) or array-like (of geometries)"
+            )
 
         self._rfdict[rfname] = [shapefile, featuretype, level]
-        sn = os.path.join(self.model_ws, f"{shapefile}.shp")
+        sn = os.path.join(self.model_ws, shapefile)
         assert os.path.isfile(sn), f"Shapefile does not exist: {sn}"
 
         for k in layers: