refactor(flopy.utils.postprocessing.get_water_table) (#1560)

* refactor(flopy.utils.postprocessing): vectorize get_water_table; remove per_idx argument. Instead, if a 3D head array is passed (nlay x nrow x ncol), a 2D water table is computed; if a 4D head array is passed (with multiple stress periods); a 3D water table is returned (nper x nrow x ncol). 
* refactor(flopy.utils.postprocessing::get_water_table): add hdry, hnoflo and masked_values arguments

autotest/test_postprocessing.py

@@ -243,26 +243,32 @@ def test_get_transmissivities(tmpdir):
 
 
 def test_get_water_table():
-    nodata = -9999.0
-    hds = np.ones((3, 3, 3), dtype=float) * nodata
+    hdry = -1e30
+    hds = np.ones((3, 3, 3), dtype=float) * hdry
+    hds[-1, :, :] = 2.0
+    hds[1, 1, 1] = 1.0
+    hds[0, -1, -1] = 1e30
+    wt = get_water_table(hds)
+    assert wt.shape == (3, 3)
+    assert wt[1, 1] == 1.0
+    assert np.sum(wt) == 17.0
+
+    hdry = -9999
+    hds = np.ones((3, 3, 3), dtype=float) * hdry
     hds[-1, :, :] = 2.0
     hds[1, 1, 1] = 1.0
-    wt = get_water_table(hds, nodata=nodata)
+    hds[0, -1, -1] = 9999
+    wt = get_water_table(hds, hdry=-9999, hnoflo=9999)
     assert wt.shape == (3, 3)
     assert wt[1, 1] == 1.0
     assert np.sum(wt) == 17.0
 
     hds2 = np.array([hds, hds])
-    wt = get_water_table(hds2, nodata=nodata)
+    wt = get_water_table(hds2, hdry=-9999, hnoflo=9999)
     assert wt.shape == (2, 3, 3)
     assert np.sum(wt[:, 1, 1]) == 2.0
     assert np.sum(wt) == 34.0
 
-    wt = get_water_table(hds2, nodata=nodata, per_idx=0)
-    assert wt.shape == (3, 3)
-    assert wt[1, 1] == 1.0
-    assert np.sum(wt) == 17.0
-
 
 def test_get_sat_thickness_gradients(tmpdir):
     nodata = -9999.0

flopy/utils/postprocessing.py

@@ -123,7 +123,7 @@ def get_transmissivities(
     return T
 
 
-def get_water_table(heads, nodata, per_idx=None):
+def get_water_table(heads, hdry=-1e30, hnoflo=1e30, masked_values=None):
     """
     Get a 2D array representing the water table elevation for each
     stress period in heads array.
@@ -132,11 +132,16 @@ def get_water_table(heads, nodata, per_idx=None):
     ----------
     heads : 3 or 4-D np.ndarray
         Heads array.
-    nodata : real
-        HDRY value indicating dry cells.
-    per_idx : int or sequence of ints
-        stress periods to return. If None,
-        returns all stress periods (default is None).
+    hdry : real
+        The head that is assigned to cells that are converted to dry
+        during a simulation. By default, -1e30.
+    hnoflo : real
+        The value of head assigned to all inactive (no flow) cells
+        throughout the simulation, including vertical pass-through cells in
+        MODFLOW 6. By default, 1e30.
+    masked_values : list
+        List of any values (in addition to hdry and hnoflo) that should
+        be masked in the water table calculation.
 
     Returns
     -------
@@ -145,25 +150,20 @@ def get_water_table(heads, nodata, per_idx=None):
 
     """
     heads = np.array(heads, ndmin=4)
-    nper, nlay, nrow, ncol = heads.shape
-    if per_idx is None:
-        per_idx = list(range(nper))
-    elif np.isscalar(per_idx):
-        per_idx = [per_idx]
-    wt = []
-    for per in per_idx:
-        wt_per = []
-        for i in range(nrow):
-            for j in range(ncol):
-                for k in range(nlay):
-                    if heads[per, k, i, j] != nodata:
-                        wt_per.append(heads[per, k, i, j])
-                        break
-                    elif k == nlay - 1:
-                        wt_per.append(nodata)
-        assert len(wt_per) == nrow * ncol
-        wt.append(np.reshape(wt_per, (nrow, ncol)))
-    return np.squeeze(wt)
+    mask = (heads == hdry) | (heads == hnoflo)
+    if masked_values is not None:
+        for val in masked_values:
+            mask = mask | (heads == val)
+    k = (~mask).argmax(axis=1)
+    per, i, j = np.indices(k.shape)
+    wt = heads[per.ravel(), k.ravel(), i.ravel(), j.ravel()].reshape(k.shape)
+    wt = np.squeeze(wt)
+    mask = (wt == hdry) | (wt == hnoflo)
+    if masked_values is not None:
+        for val in masked_values:
+            mask = mask | (wt == val)
+    wt = np.ma.masked_array(wt, mask)
+    return wt
 
 
 def get_gradients(heads, m, nodata, per_idx=None):