fix(HeadUFile): repro & fix modflowpy#1503, add tests

autotest/test_grid.py

@@ -1320,3 +1320,52 @@ def test_ncb_thick():
         raise AssertionError(
             "saturated_thickness is not properly indexing confining beds"
         )
+
+
+def test_usg_iverts():
+    iverts = [
+        [4, 3, 2, 1, 0, None],
+        [7, 0, 1, 6, 5, None],
+        [11, 10, 9, 8, 2, 3],
+        [1, 6, 13, 12, 8, 2],
+        [15, 14, 13, 6, 5, None],
+        [10, 9, 18, 17, 16, None],
+        [8, 12, 20, 19, 18, 9],
+        [22, 14, 13, 12, 20, 21],
+        [24, 17, 18, 19, 23, None],
+        [21, 20, 19, 23, 25, None],
+    ]
+    verts = [
+        [0.0, 22.5],
+        [5.1072, 22.5],
+        [7.5, 24.0324],
+        [7.5, 30.0],
+        [0.0, 30.0],
+        [0.0, 7.5],
+        [4.684, 7.5],
+        [0.0, 15.0],
+        [14.6582, 21.588],
+        [22.5, 24.3766],
+        [22.5, 30.0],
+        [15.0, 30.0],
+        [15.3597, 8.4135],
+        [7.5, 5.6289],
+        [7.5, 0.0],
+        [0.0, 0.0],
+        [30.0, 30.0],
+        [30.0, 22.5],
+        [25.3285, 22.5],
+        [24.8977, 7.5],
+        [22.5, 5.9676],
+        [22.5, 0.0],
+        [15.0, 0.0],
+        [30.0, 7.5],
+        [30.0, 15.0],
+        [30.0, 0.0],
+    ]
+
+    grid = UnstructuredGrid(verts, iverts, ncpl=[len(iverts)])
+
+    iverts = grid.iverts
+    if any(None in l for l in iverts):
+        raise ValueError("None type should not be returned in iverts list")

autotest/test_headufile.py

@@ -1,8 +1,4 @@
-"""
-HeadUFile get_ts tests using t505_test.py
-"""
-
-import os
+from pathlib import Path
 
 import pytest
 from autotest.conftest import requires_exe, requires_pkg
@@ -17,12 +13,12 @@
     ModflowOc,
 )
 from flopy.utils import HeadUFile
+from flopy.utils.binaryfile import get_lni
 from flopy.utils.gridgen import Gridgen
 
 
-@requires_exe("mfusg", "gridgen")
-@requires_pkg("shapely", "shapefile")
-def test_mfusg(tmpdir):
+@pytest.fixture(scope="module")
+def mfusg_model(module_tmpdir):
     from shapely.geometry import Polygon
 
     name = "dummy"
@@ -36,7 +32,7 @@ def test_mfusg(tmpdir):
     botm = [top - k * dz for k in range(1, nlay + 1)]
 
     # create dummy model and dis package for gridgen
-    m = Modflow(modelname=name, model_ws=str(tmpdir))
+    m = Modflow(modelname=name, model_ws=str(module_tmpdir))
     dis = ModflowDis(
         m,
         nlay=nlay,
@@ -49,7 +45,7 @@ def test_mfusg(tmpdir):
     )
 
     # Create and build the gridgen model with a refined area in the middle
-    g = Gridgen(dis, model_ws=str(tmpdir))
+    g = Gridgen(dis, model_ws=str(module_tmpdir))
 
     polys = [Polygon([(4, 4), (6, 4), (6, 6), (4, 6)])]
     g.add_refinement_features(polys, "polygon", 3, layers=[0])
@@ -68,7 +64,7 @@ def test_mfusg(tmpdir):
     name = "mymodel"
     m = MfUsg(
         modelname=name,
-        model_ws=str(tmpdir),
+        model_ws=str(module_tmpdir),
         exe_name="mfusg",
         structured=False,
     )
@@ -89,72 +85,63 @@ def test_mfusg(tmpdir):
 
     m.run_model()
 
-    # head is returned as a list of head arrays for each layer
-    head_file = os.path.join(str(tmpdir), f"{name}.hds")
-    head = HeadUFile(head_file).get_data()
+    # head contains a list of head arrays for each layer
+    head_file_path = Path(module_tmpdir / f"{name}.hds")
+    return m, HeadUFile(str(head_file_path))
+
+
+@requires_exe("mfusg", "gridgen")
+@requires_pkg("shapely", "shapefile")
+def test_get_ts_single_node(mfusg_model):
+    model, head_file = mfusg_model
+    head = head_file.get_data()
 
     # test if single node idx works
-    one_hds = HeadUFile(head_file).get_ts(idx=300)
-    if one_hds[0, 1] != head[0][300]:
-        raise AssertionError(
-            "Error head from 'get_ts' != head from 'get_data'"
-        )
+    one_hds = head_file.get_ts(idx=300)
+    assert one_hds[0, 1] == head[0][300], "head from 'get_ts' != head from 'get_data'"
+
+
+@requires_exe("mfusg", "gridgen")
+@requires_pkg("shapely", "shapefile")
+def test_get_ts_multiple_nodes(mfusg_model):
+    model, head_file = mfusg_model
+    grid = model.modelgrid
+    head = head_file.get_data()
+
+    # test if list of nodes for idx works
+    nodes = [500, 300, 182, 65]
+    multi_hds = head_file.get_ts(idx=nodes)
+    for i, node in enumerate(nodes):
+        li, ni = get_lni(grid.ncpl, node)
+        assert multi_hds[0, i + 1] == head[li][ni], "head from 'get_ts' != head from 'get_data'"
+
+
+@requires_exe("mfusg", "gridgen")
+@requires_pkg("shapely", "shapefile")
+def test_get_ts_all_nodes(mfusg_model):
+    model, head_file = mfusg_model
+    grid = model.modelgrid
+    head = head_file.get_data()
 
     # test if list of nodes for idx works
-    nodes = [300, 182, 65]
+    nodes = list(range(0, grid.nnodes))
+    multi_hds = head_file.get_ts(idx=nodes)
+    for i, node in enumerate(nodes):
+        li, ni = get_lni(grid.ncpl, node)
+        assert multi_hds[0, i + 1] == head[li][ni], "head from 'get_ts' != head from 'get_data'"
+
+
+@pytest.mark.xfail(reason="https://github.com/modflowpy/flopy/issues/1503")
+@requires_exe("mfusg", "gridgen")
+@requires_pkg("shapely", "shapefile")
+def test_get_ts_all_nodes_old(mfusg_model):
+    model, head_file = mfusg_model
+    grid = model.modelgrid
+    head = head_file.get_data()
 
-    multi_hds = HeadUFile(head_file).get_ts(idx=nodes)
+    # test if list of nodes for idx works
+    nodes = list(range(0, grid.nnodes))
+    multi_hds = head_file.__get_ts(idx=nodes)
     for i, node in enumerate(nodes):
-        if multi_hds[0, i + 1] != head[0][node]:
-            raise AssertionError(
-                "Error head from 'get_ts' != head from 'get_data'"
-            )
-
-
-def test_usg_iverts():
-    iverts = [
-        [4, 3, 2, 1, 0, None],
-        [7, 0, 1, 6, 5, None],
-        [11, 10, 9, 8, 2, 3],
-        [1, 6, 13, 12, 8, 2],
-        [15, 14, 13, 6, 5, None],
-        [10, 9, 18, 17, 16, None],
-        [8, 12, 20, 19, 18, 9],
-        [22, 14, 13, 12, 20, 21],
-        [24, 17, 18, 19, 23, None],
-        [21, 20, 19, 23, 25, None],
-    ]
-    verts = [
-        [0.0, 22.5],
-        [5.1072, 22.5],
-        [7.5, 24.0324],
-        [7.5, 30.0],
-        [0.0, 30.0],
-        [0.0, 7.5],
-        [4.684, 7.5],
-        [0.0, 15.0],
-        [14.6582, 21.588],
-        [22.5, 24.3766],
-        [22.5, 30.0],
-        [15.0, 30.0],
-        [15.3597, 8.4135],
-        [7.5, 5.6289],
-        [7.5, 0.0],
-        [0.0, 0.0],
-        [30.0, 30.0],
-        [30.0, 22.5],
-        [25.3285, 22.5],
-        [24.8977, 7.5],
-        [22.5, 5.9676],
-        [22.5, 0.0],
-        [15.0, 0.0],
-        [30.0, 7.5],
-        [30.0, 15.0],
-        [30.0, 0.0],
-    ]
-
-    grid = UnstructuredGrid(verts, iverts, ncpl=[len(iverts)])
-
-    iverts = grid.iverts
-    if any(None in l for l in iverts):
-        raise ValueError("None type should not be returned in iverts list")
+        li, ni = get_lni(grid.ncpl, node)
+        assert multi_hds[0, i + 1] == head[li][ni], "head from 'get_ts' != head from 'get_data'"

flopy/utils/binaryfile.py

@@ -1820,6 +1820,36 @@ def close(self):
         return
 
 
+def get_lni(ncpl, nn):
+    """
+    Get the layer and node index (within the layer), given the node number and node count per layer.
+
+    Parameters
+    ----------
+    ncpl : the node count per layer (int or array of ints)
+    nn : the grid node number
+
+    Returns
+    -------
+        A tuple (layer index, node index)
+    """
+
+    # nothing to do if ncpl doesn't vary by layer
+    if isinstance(ncpl, int):
+        return 0, nn
+
+    if not isinstance(ncpl, (list, tuple, np.ndarray)):
+        raise ValueError(f"ncpl must be int or collection of ints")
+
+    csum = np.cumsum([0] + list(ncpl))
+    layer = max(0, np.searchsorted(csum, nn) - 1)
+    nidx = nn - sum([ncpl[l] for l in range(0, layer)])
+    cs = list(ncpl)[layer]
+    if layer + 1 < len(csum) and nidx == cs:
+        return layer + 1, 0
+    return layer, nidx
+
+
 class HeadUFile(BinaryLayerFile):
     """
     Unstructured MODFLOW-USG HeadUFile Class.
@@ -1892,7 +1922,6 @@ def __init__(
             bintype="Head", precision=precision
         )
         super().__init__(filename, precision, verbose, kwargs)
-        return
 
     def _get_data_array(self, totim=0.0):
         """
@@ -1963,6 +1992,47 @@ def get_ts(self, idx):
 
         """
         times = self.get_times()
+        data = self.get_data(totim=times[0])
+        layers = len(data)
+        ncpl = [len(data[l]) for l in range(layers)]
+        result = []
+
+        if isinstance(idx, int):
+            li, ni = get_lni(ncpl, idx)
+            for i, time in enumerate(times):
+                data = self.get_data(totim=time)
+                value = data[li][ni]
+                result.append([time, value])
+        elif isinstance(idx, list) and all(isinstance(x, int) for x in idx):
+            for i, time in enumerate(times):
+                data = self.get_data(totim=time)
+                row = [time]
+                for nn in idx:
+                    li, ni = get_lni(ncpl, nn)
+                    value = data[li][ni]
+                    row += [value]
+                result.append(row)
+        else:
+            raise ValueError("idx must be an integer or a list of integers")
+
+        return np.array(result)
+
+    def __get_ts(self, idx):
+        """
+        Get a time series from the binary HeadUFile
+        Parameters
+        ----------
+        idx : int or list of ints
+            idx can be nodenumber or it can be a list in the form
+            [nodenumber, nodenumber, ...].  The nodenumber,
+            values must be zero based.
+        Returns
+        ----------
+        out : numpy array
+            Array has size (ntimes, ncells + 1).  The first column in the
+            data array will contain time (totim).
+        """
+        times = self.get_times()
 
         # find node number in layer that node is in
         data = self.get_data(totim=times[0])