refactor(gridutil): improve arg handling in get_disu_kwargs (#2480)

Author: wpbonelli

autotest/test_gridutil.py

@@ -72,16 +72,16 @@ def test_get_lni_infers_layer_count_when_int_ncpl(ncpl, nodes, expected):
             61,
             np.array(61 * [50]),
             np.array(61 * [50]),
-            np.array([-10]),
-            np.array([-30.0, -50.0]),
+            -10,
+            np.array([-30.0]),
         ),
         (
             2,
             61,
             61,
             np.array(61 * [50]),
             np.array(61 * [50]),
-            np.array([-10]),
+            -10,
             np.array([-30.0, -50.0]),
         ),
         (
@@ -90,8 +90,8 @@ def test_get_lni_infers_layer_count_when_int_ncpl(ncpl, nodes, expected):
             4,  # ncol
             np.array(4 * [4.0]),  # delr
             np.array(3 * [3.0]),  # delc
-            np.array([-10]),  # top
-            np.array([-30.0]),  # botm
+            -10,  # top
+            -30.0,  # botm
         ),
     ],
 )
@@ -107,10 +107,6 @@ def test_get_disu_kwargs(nlay, nrow, ncol, delr, delc, tp, botm):
         return_vertices=True,
     )
 
-    from pprint import pprint
-
-    pprint(kwargs["area"])
-
     assert kwargs["nodes"] == nlay * nrow * ncol
     assert kwargs["nvert"] == (nrow + 1) * (ncol + 1)
 

flopy/utils/gridutil.py

@@ -98,12 +98,45 @@ def get_disu_kwargs(
     def get_nn(k, i, j):
         return k * nrow * ncol + i * ncol + j
 
-    if not isinstance(delr, np.ndarray):
-        delr = np.array(delr)
-    if not isinstance(delc, np.ndarray):
-        delc = np.array(delc)
-    assert delr.shape == (ncol,)
-    assert delc.shape == (nrow,)
+    # delr check
+    if np.isscalar(delr):
+        delr = delr * np.ones(ncol, dtype=float)
+    else:
+        assert np.asanyarray(delr).shape == (ncol,), (
+            "delr must be array with shape (ncol,), got {}".format(delr.shape)
+        )
+
+    # delc check
+    if np.isscalar(delc):
+        delc = delc * np.ones(nrow, dtype=float)
+    else:
+        assert np.asanyarray(delc).shape == (nrow,), (
+            "delc must be array with shape (nrow,), got {}".format(delc.shape)
+        )
+
+    # tp check
+    if np.isscalar(tp):
+        tp = tp * np.ones((nrow, ncol), dtype=float)
+    else:
+        assert np.asanyarray(tp).shape == (
+            nrow,
+            ncol,
+        ), "tp must be scalar or array with shape (nrow, ncol), got {}".format(tp.shape)
+
+    # botm check
+    if np.isscalar(botm):
+        botm = botm * np.ones((nlay, nrow, ncol), dtype=float)
+    elif np.asanyarray(botm).shape == (nlay,):
+        b = np.empty((nlay, nrow, ncol), dtype=float)
+        for k in range(nlay):
+            b[k] = botm[k]
+        botm = b
+    else:
+        assert np.asanyarray(botm).shape == (
+            nlay,
+            nrow,
+            ncol,
+        ), "botm must be array with shape (nlay, nrow, ncol), got {}".format(botm.shape)
 
     nodes = nlay * nrow * ncol
     iac = np.zeros((nodes), dtype=int)
@@ -126,16 +159,16 @@ def get_nn(k, i, j):
                 cl12.append(n + 1)
                 hwva.append(n + 1)
                 if k == 0:
-                    top[n] = tp.item() if isinstance(tp, np.ndarray) else tp
+                    top[n] = tp[i, j]
                 else:
-                    top[n] = botm[k - 1]
-                bot[n] = botm[k]
+                    top[n] = botm[k - 1, i, j]
+                bot[n] = botm[k, i, j]
                 # up
                 if k > 0:
                     ja.append(get_nn(k - 1, i, j))
                     iac[n] += 1
                     ihc.append(0)
-                    dz = botm[k - 1] - botm[k]
+                    dz = botm[k - 1, i, j] - botm[k, i, j]
                     cl12.append(0.5 * dz)
                     hwva.append(delr[j] * delc[i])
                 # back
@@ -172,9 +205,9 @@ def get_nn(k, i, j):
                     iac[n] += 1
                     ihc.append(0)
                     if k == 0:
-                        dz = tp - botm[k]
+                        dz = tp[i, j] - botm[k, i, j]
                     else:
-                        dz = botm[k - 1] - botm[k]
+                        dz = botm[k - 1, i, j] - botm[k, i, j]
                     cl12.append(0.5 * dz)
                     hwva.append(delr[j] * delc[i])
     ja = np.array(ja, dtype=int)
@@ -272,28 +305,31 @@ def get_disv_kwargs(
     if np.isscalar(delr):
         delr = delr * np.ones(ncol, dtype=float)
     else:
-        assert delr.shape == (ncol,), "delr must be array with shape (ncol,)"
+        assert np.asanyarray(delr).shape == (ncol,), (
+            "delr must be array with shape (ncol,), got {}".format(delr.shape)
+        )
 
     # delc check
     if np.isscalar(delc):
         delc = delc * np.ones(nrow, dtype=float)
     else:
-        assert delc.shape == (nrow,), "delc must be array with shape (nrow,)"
+        assert np.asanyarray(delc).shape == (nrow,), (
+            "delc must be array with shape (nrow,), got {}".format(delc.shape)
+        )
 
     # tp check
     if np.isscalar(tp):
         tp = tp * np.ones((nrow, ncol), dtype=float)
     else:
-        assert tp.shape == (
+        assert np.asanyarray(tp).shape == (
             nrow,
             ncol,
-        ), "tp must be scalar or array with shape (nrow, ncol)"
+        ), "tp must be scalar or array with shape (nrow, ncol), got {}".format(tp.shape)
 
     # botm check
     if np.isscalar(botm):
         botm = botm * np.ones((nlay, nrow, ncol), dtype=float)
-    elif isinstance(botm, list):
-        assert len(botm) == nlay, "if botm provided as a list it must have length nlay"
+    elif np.asanyarray(botm).shape == (nlay,):
         b = np.empty((nlay, nrow, ncol), dtype=float)
         for k in range(nlay):
             b[k] = botm[k]
@@ -303,7 +339,7 @@ def get_disv_kwargs(
             nlay,
             nrow,
             ncol,
-        ), "botm must be array with shape (nlay, nrow, ncol)"
+        ), "botm must be array with shape (nlay, nrow, ncol), got {}".format(botm.shape)
 
     # build vertices
     xv = np.cumsum(delr)