refactor(contour_array): add layer param, update docstrings, expand t…

…ests (#1975)

* add layer param to flopy.export.utils.contour_array()
* pop/pass layer kwarg to contour_array() in export_array_contours()
* test export_array_contours() disu grids with constant/variable ncpl
* test PlotMapView.contour_array() array formats and with/without layer
* update PlotMapView.contour_array() docstring (dimension expectations)
* refactor some disu test grids & models into fixtures

autotest/test_export.py

@@ -139,6 +139,44 @@ def disu_sim(name, tmpdir, missing_arrays=False):
     return sim
 
 
+@pytest.fixture
+def unstructured_grid(example_data_path):
+    ws = example_data_path / "unstructured"
+
+    # load vertices
+    verts = load_verts(ws / "ugrid_verts.dat")
+
+    # load the index list into iverts, xc, and yc
+    iverts, xc, yc = load_iverts(ws / "ugrid_iverts.dat", closed=True)
+
+    # create a 3 layer model grid
+    ncpl = np.array(3 * [len(iverts)])
+    nnodes = np.sum(ncpl)
+
+    top = np.ones(nnodes)
+    botm = np.ones(nnodes)
+
+    # set top and botm elevations
+    i0 = 0
+    i1 = ncpl[0]
+    elevs = [100, 0, -100, -200]
+    for ix, cpl in enumerate(ncpl):
+        top[i0:i1] *= elevs[ix]
+        botm[i0:i1] *= elevs[ix + 1]
+        i0 += cpl
+        i1 += cpl
+
+    return UnstructuredGrid(
+        vertices=verts,
+        iverts=iverts,
+        xcenters=xc,
+        ycenters=yc,
+        top=top,
+        botm=botm,
+        ncpl=ncpl,
+    )
+
+
 @requires_pkg("shapefile")
 @pytest.mark.parametrize("pathlike", (True, False))
 def test_output_helper_shapefile_export(
@@ -613,7 +651,7 @@ def test_export_array2(function_tmpdir):
 
 
 @requires_pkg("shapefile", "shapely")
-def test_export_array_contours(function_tmpdir):
+def test_export_array_contours_structured(function_tmpdir):
     nrow = 7
     ncol = 11
     crs = 4431
@@ -648,6 +686,61 @@ def test_export_array_contours(function_tmpdir):
     assert os.path.isfile(filename), "did not create contour shapefile"
 
 
+@requires_pkg("shapefile", "shapely")
+def test_export_array_contours_unstructured(
+    function_tmpdir, unstructured_grid
+):
+    from shapefile import Reader
+
+    grid = unstructured_grid
+    fname = function_tmpdir / "myarraycontours1.shp"
+    export_array_contours(grid, fname, np.arange(grid.nnodes))
+    assert fname.is_file(), "did not create contour shapefile"
+
+    # visual debugging
+    grid.plot(alpha=0.2)
+    with Reader(fname) as r:
+        shapes = r.shapes()
+        for s in shapes:
+            x = [i[0] for i in s.points[:]]
+            y = [i[1] for i in s.points[:]]
+            plt.plot(x, y)
+
+    # plt.show()
+
+
+from autotest.test_gridgen import sim_disu_diff_layers
+
+
+@requires_pkg("shapefile", "shapely")
+def test_export_array_contours_unstructured_diff_layers(
+    function_tmpdir, sim_disu_diff_layers
+):
+    from shapefile import Reader
+
+    gwf = sim_disu_diff_layers.get_model()
+    grid = gwf.modelgrid
+    a = np.arange(grid.nnodes)
+    for layer in range(3):
+        fname = function_tmpdir / f"contours.{layer}.shp"
+        export_array_contours(grid, fname, a, layer=layer)
+        assert fname.is_file(), "did not create contour shapefile"
+
+    # visual debugging
+    fig, axes = plt.subplots(1, 3, subplot_kw={"aspect": "equal"})
+    for layer, ax in enumerate(axes):
+        fname = function_tmpdir / f"contours.{layer}.shp"
+        with Reader(fname) as r:
+            shapes = r.shapes()
+            for s in shapes:
+                x = [i[0] for i in s.points[:]]
+                y = [i[1] for i in s.points[:]]
+                ax.plot(x, y)
+            grid.plot(ax=ax, alpha=0.2, layer=layer)
+
+    # plt.show()
+
+
 @requires_pkg("shapefile", "shapely")
 def test_export_contourf(function_tmpdir, example_data_path):
     from shapefile import Reader
@@ -1331,52 +1424,18 @@ def test_vtk_vector(function_tmpdir, example_data_path):
 
 
 @requires_pkg("vtk")
-def test_vtk_unstructured(function_tmpdir, example_data_path):
+def test_vtk_unstructured(function_tmpdir, unstructured_grid):
     from vtkmodules.util.numpy_support import vtk_to_numpy
     from vtkmodules.vtkIOLegacy import vtkUnstructuredGridReader
 
-    u_data_ws = example_data_path / "unstructured"
-
-    # load vertices
-    verts = load_verts(u_data_ws / "ugrid_verts.dat")
-
-    # load the index list into iverts, xc, and yc
-    iverts, xc, yc = load_iverts(u_data_ws / "ugrid_iverts.dat", closed=True)
-
-    # create a 3 layer model grid
-    ncpl = np.array(3 * [len(iverts)])
-    nnodes = np.sum(ncpl)
-
-    top = np.ones(nnodes)
-    botm = np.ones(nnodes)
-
-    # set top and botm elevations
-    i0 = 0
-    i1 = ncpl[0]
-    elevs = [100, 0, -100, -200]
-    for ix, cpl in enumerate(ncpl):
-        top[i0:i1] *= elevs[ix]
-        botm[i0:i1] *= elevs[ix + 1]
-        i0 += cpl
-        i1 += cpl
-
-    # create the modelgrid
-    modelgrid = UnstructuredGrid(
-        vertices=verts,
-        iverts=iverts,
-        xcenters=xc,
-        ycenters=yc,
-        top=top,
-        botm=botm,
-        ncpl=ncpl,
-    )
+    grid = unstructured_grid
 
     outfile = function_tmpdir / "disu_grid.vtu"
     vtkobj = Vtk(
-        modelgrid=modelgrid, vertical_exageration=2, binary=True, smooth=False
+        modelgrid=grid, vertical_exageration=2, binary=True, smooth=False
     )
-    vtkobj.add_array(modelgrid.top, "top")
-    vtkobj.add_array(modelgrid.botm, "botm")
+    vtkobj.add_array(grid.top, "top")
+    vtkobj.add_array(grid.botm, "botm")
     vtkobj.write(outfile)
 
     assert is_binary_file(outfile)
@@ -1390,7 +1449,9 @@ def test_vtk_unstructured(function_tmpdir, example_data_path):
 
     top2 = vtk_to_numpy(data.GetCellData().GetArray("top"))
 
-    assert np.allclose(np.ravel(top), top2), "Field data not properly written"
+    assert np.allclose(
+        np.ravel(grid.top), top2
+    ), "Field data not properly written"
 
 
 @requires_pkg("vtk", "pyvista")

autotest/test_gridgen.py

@@ -164,13 +164,11 @@ def test_mf6disv(function_tmpdir):
     plt.close("all")
 
 
-@pytest.mark.slow
-@requires_exe("mf6", "gridgen")
-@requires_pkg("shapely", "shapefile")
-def test_mf6disu(function_tmpdir):
+@pytest.fixture
+def sim_disu_diff_layers(function_tmpdir):
     from shapely.geometry import Polygon
 
-    name = "dummy"
+    name = "disu_diff_layers"
     nlay = 3
     nrow = 10
     ncol = 10
@@ -232,6 +230,17 @@ def test_mf6disu(function_tmpdir):
         head_filerecord=head_file,
         saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
     )
+
+    return sim
+
+
+@pytest.mark.slow
+@requires_exe("mf6", "gridgen")
+@requires_pkg("shapely", "shapefile")
+def test_mf6disu(sim_disu_diff_layers):
+    sim = sim_disu_diff_layers
+    ws = sim.sim_path
+    gwf = sim.get_model()
     sim.write_simulation()
 
     gwf.modelgrid.set_coord_info(angrot=15)
@@ -243,9 +252,9 @@ def test_mf6disu(function_tmpdir):
     assert np.allclose(gwf.modelgrid.ncpl, np.array([436, 184, 112]))
 
     # write grid and model shapefiles
-    fname = function_tmpdir / "grid.shp"
+    fname = ws / "grid.shp"
     gwf.modelgrid.write_shapefile(fname)
-    fname = function_tmpdir / "model.shp"
+    fname = ws / "model.shp"
     gwf.export(fname)
 
     sim.run_simulation(silent=True)
@@ -272,7 +281,7 @@ def test_mf6disu(function_tmpdir):
         ax.set_title(f"Layer {ilay + 1}")
         pmv.plot_vector(spdis["qx"], spdis["qy"], color="white")
     fname = "results.png"
-    fname = function_tmpdir / fname
+    fname = ws / fname
     plt.savefig(fname)
     plt.close("all")
 
@@ -307,7 +316,7 @@ def test_mf6disu(function_tmpdir):
         sim_name=name,
         version="mf6",
         exe_name="mf6",
-        sim_ws=function_tmpdir,
+        sim_ws=ws,
     )
     gwf = sim.get_model(name)
 

autotest/test_plot_map_view.py

@@ -31,6 +31,12 @@
 from flopy.utils import CellBudgetFile, EndpointFile, HeadFile, PathlineFile
 
 
+@pytest.fixture
+def rng():
+    # set seed so parametrized plot tests are comparable
+    return np.random.default_rng(0)
+
+
 @requires_pkg("shapely")
 def test_map_view():
     m = flopy.modflow.Modflow(rotation=20.0)
@@ -162,21 +168,19 @@ def test_map_view_bc_UZF_3lay(example_data_path):
         ), f"Unexpected collection type: {type(col)}"
 
 
-def test_map_view_contour(function_tmpdir):
-    arr = np.random.rand(10, 10) * 100
-    arr[-1, :] = np.nan
-    delc = np.array([10] * 10, dtype=float)
-    delr = np.array([8] * 10, dtype=float)
-    top = np.ones((10, 10), dtype=float)
-    botm = np.ones((3, 10, 10), dtype=float)
+@pytest.mark.parametrize("ndim", [1, 2, 3])
+def test_map_view_contour_array_structured(function_tmpdir, ndim, rng):
+    nlay, nrow, ncol = 3, 10, 10
+    ncpl = nrow * ncol
+    delc = np.array([10] * nrow, dtype=float)
+    delr = np.array([8] * ncol, dtype=float)
+    top = np.ones((nrow, ncol), dtype=float)
+    botm = np.ones((nlay, nrow, ncol), dtype=float)
     botm[0] = 0.75
     botm[1] = 0.5
     botm[2] = 0.25
-    idomain = np.ones((3, 10, 10))
+    idomain = np.ones((nlay, nrow, ncol))
     idomain[0, 0, :] = 0
-    vmin = np.nanmin(arr)
-    vmax = np.nanmax(arr)
-    levels = np.linspace(vmin, vmax, 7)
 
     grid = StructuredGrid(
         delc=delc,
@@ -185,16 +189,49 @@ def test_map_view_contour(function_tmpdir):
         botm=botm,
         idomain=idomain,
         lenuni=1,
-        nlay=3,
-        nrow=10,
-        ncol=10,
+        nlay=nlay,
+        nrow=nrow,
+        ncol=ncol,
     )
 
-    pmv = PlotMapView(modelgrid=grid, layer=0)
-    contours = pmv.contour_array(a=arr)
-    plt.savefig(function_tmpdir / "map_view_contour.png")
+    # define full grid 1D array to contour
+    arr = rng.random(nlay * nrow * ncol) * 100
+
+    for l in range(nlay):
+        if ndim == 1:
+            # full grid 1D array
+            pmv = PlotMapView(modelgrid=grid, layer=l)
+            contours = pmv.contour_array(a=arr)
+            fname = f"map_view_contour_{ndim}d_l{l}_full.png"
+            plt.savefig(function_tmpdir / fname)
+            plt.clf()
+
+            # 1 layer slice
+            pmv = PlotMapView(modelgrid=grid, layer=l)
+            contours = pmv.contour_array(a=arr[(l * ncpl) : ((l + 1) * ncpl)])
+            fname = f"map_view_contour_{ndim}d_l{l}_1lay.png"
+            plt.savefig(function_tmpdir / fname)
+            plt.clf()
+        elif ndim == 2:
+            # 1 layer as 2D
+            # arr[-1, :] = np.nan  # add nan to test nan handling
+            pmv = PlotMapView(modelgrid=grid, layer=l)
+            contours = pmv.contour_array(
+                a=arr.reshape(nlay, nrow, ncol)[l, :, :]
+            )
+            plt.savefig(function_tmpdir / f"map_view_contour_{ndim}d_l{l}.png")
+            plt.clf()
+        elif ndim == 3:
+            # full grid as 3D
+            pmv = PlotMapView(modelgrid=grid, layer=l)
+            contours = pmv.contour_array(a=arr.reshape(nlay, nrow, ncol))
+            plt.savefig(function_tmpdir / f"map_view_contour_{ndim}d_l{l}.png")
+            plt.clf()
 
     # if we ever revert from standard contours to tricontours, restore this nan check
+    # vmin = np.nanmin(arr)
+    # vmax = np.nanmax(arr)
+    # levels = np.linspace(vmin, vmax, 7)
     # for ix, lev in enumerate(contours.levels):
     #     if not np.allclose(lev, levels[ix]):
     #         raise AssertionError("TriContour NaN catch Failed")

flopy/discretization/structuredgrid.py

@@ -1637,7 +1637,7 @@ def get_plottable_layer_array(self, a, layer):
                 plotarray = plotarray.reshape(self.shape)
                 plotarray = plotarray[layer, :, :]
         else:
-            raise Exception("Array to plot must be of dimension 1, 2, or 3")
+            raise ValueError("Array to plot must be of dimension 1, 2, or 3")
         msg = f"{plotarray.shape} /= {required_shape}"
         assert plotarray.shape == required_shape, msg
         return plotarray

flopy/discretization/vertexgrid.py

@@ -528,7 +528,7 @@ def get_plottable_layer_array(self, a, layer):
                 a = np.squeeze(a, axis=1)
                 plotarray = a[layer, :]
             else:
-                raise Exception(
+                raise ValueError(
                     "Array has 3 dimensions so one of them must be of size 1 "
                     "for a VertexGrid."
                 )
@@ -540,7 +540,7 @@ def get_plottable_layer_array(self, a, layer):
                 plotarray = plotarray.reshape(self.nlay, self.ncpl)
                 plotarray = plotarray[layer, :]
         else:
-            raise Exception("Array to plot must be of dimension 1 or 2")
+            raise ValueError("Array to plot must be of dimension 1 or 2")
         msg = f"{plotarray.shape[0]} /= {required_shape}"
         assert plotarray.shape == required_shape, msg
         return plotarray