added support to get_ts for HeadUFile

autotest/t420_test.py

@@ -0,0 +1,140 @@
+"""
+HeadUFile get_ts tests using t505_test.py
+"""
+
+import os
+import sys
+import platform
+import shutil
+import numpy as np
+
+try:
+    from shapely.geometry import Polygon
+except ImportWarning as e:
+    print("Shapely not installed, tests cannot be run.")
+    Polygon = None
+
+
+import flopy
+from flopy.utils.gridgen import Gridgen
+
+try:
+    import matplotlib
+    import matplotlib.pyplot as plt
+    from matplotlib.collections import QuadMesh, PathCollection, LineCollection
+except:
+    print("Matplotlib not installed, tests cannot be run.")
+    matplotlib = None
+    plt = None
+
+# Set gridgen executable
+gridgen_exe = "gridgen"
+if platform.system() in "Windows":
+    gridgen_exe += ".exe"
+gridgen_exe = flopy.which(gridgen_exe)
+
+# set mfusg executable
+mfusg_exe = "mfusg"
+if platform.system() in "Windows":
+    mfusg_exe += ".exe"
+mfusg_exe = flopy.which(mfusg_exe)
+
+# set up the example folder
+tpth = os.path.join("temp", "t420")
+if not os.path.isdir(tpth):
+    os.makedirs(tpth)
+
+# set up a gridgen workspace
+gridgen_ws = os.path.join(tpth, "gridgen_t420")
+if not os.path.exists(gridgen_ws):
+    os.makedirs(gridgen_ws)
+
+def test_mfusg():
+
+    name = "dummy"
+    nlay = 3
+    nrow = 10
+    ncol = 10
+    delr = delc = 1.0
+    top = 1
+    bot = 0
+    dz = (top - bot) / nlay
+    botm = [top - k * dz for k in range(1, nlay + 1)]
+
+    # create dummy model and dis package for gridgen
+    m = flopy.modflow.Modflow(modelname=name, model_ws=gridgen_ws)
+    dis = flopy.modflow.ModflowDis(
+        m,
+        nlay=nlay,
+        nrow=nrow,
+        ncol=ncol,
+        delr=delr,
+        delc=delc,
+        top=top,
+        botm=botm,
+    )
+
+    # Create and build the gridgen model with a refined area in the middle
+    g = Gridgen(dis, model_ws=gridgen_ws)
+    polys = [Polygon([(4, 4), (6, 4), (6, 6), (4, 6)])]
+    g.add_refinement_features(polys, "polygon", 3, layers=[0])
+    g.build()
+
+    chdspd = []
+    for x, y, head in [(0, 10, 1.0), (10, 0, 0.0)]:
+        ra = g.intersect([(x, y)], "point", 0)
+        ic = ra["nodenumber"][0]
+        chdspd.append([ic, head, head])
+
+    # gridprops = g.get_gridprops()
+    gridprops = g.get_gridprops_disu5()
+
+    # create the mfusg modoel
+    ws = os.path.join(tpth, "gridgen_mfusg")
+    name = "mymodel"
+    m = flopy.modflow.Modflow(
+        modelname=name,
+        model_ws=ws,
+        version="mfusg",
+        exe_name=mfusg_exe,
+        structured=False,
+    )
+    disu = flopy.modflow.ModflowDisU(m, **gridprops)
+    bas = flopy.modflow.ModflowBas(m)
+    lpf = flopy.modflow.ModflowLpf(m)
+    chd = flopy.modflow.ModflowChd(m, stress_period_data=chdspd)
+    sms = flopy.modflow.ModflowSms(m)
+    oc = flopy.modflow.ModflowOc(m, stress_period_data={(0, 0): ["save head"]})
+    m.write_input()
+
+    # MODFLOW-USG does not have vertices, so we need to create
+    # and unstructured grid and then assign it to the model. This
+    # will allow plotting and other features to work properly.
+    gridprops_ug = g.get_gridprops_unstructuredgrid()
+    ugrid = flopy.discretization.UnstructuredGrid(**gridprops_ug, angrot=-15)
+    m.modelgrid = ugrid
+
+    if mfusg_exe is not None:
+        m.run_model()
+
+        # head is returned as a list of head arrays for each layer
+        head_file = os.path.join(ws, f"{name}.hds")
+        head = flopy.utils.HeadUFile(head_file).get_data()
+
+        # test if single node idx works
+        one_hds = flopy.utils.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'")
+
+        # test if list of nodes for idx works
+        nodes = [300,182,65]
+
+        multi_hds = flopy.utils.HeadUFile(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'")
+
+    return
+
+if __name__ == '__main__':
+    test_mfusg()

flopy/utils/binaryfile.py

@@ -1939,30 +1939,58 @@ def get_databytes(self, header):
 
     def get_ts(self, idx):
         """
-        Get a time series from the binary HeadUFile (not implemented).
+        Get a time series from the binary HeadUFile
 
         Parameters
         ----------
-        idx : tuple of ints, or a list of a tuple of ints
-            idx can be (layer, row, column) or it can be a list in the form
-            [(layer, row, column), (layer, row, column), ...].  The layer,
-            row, and column values must be zero based.
+        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).
 
-        See Also
-        --------
+        """
+        times = self.get_times()
 
-        Notes
-        -----
+        # find node number in layer that node is in
+        data = self.get_data(totim=times[0])
+        nodelay = [len(data[lay]) for lay in range(len(data))]
+        nodelay_cumsum = np.cumsum([0] + nodelay)
 
-        Examples
-        --------
+        if isinstance(idx, int):
+            layer = np.searchsorted(nodelay_cumsum, idx)
+            nnode = idx - nodelay_cumsum[layer - 1]
 
-        """
-        msg = "HeadUFile: get_ts() is not implemented"
-        raise NotImplementedError(msg)
+            result = []
+            for i, time in enumerate(times):
+                data = self.get_data(totim=time)
+                result.append([time, data[layer - 1][nnode]])
+
+        elif isinstance(idx, list):
+
+            result = []
+            for i, time in enumerate(times):
+                data = self.get_data(totim=time)
+                row = [time]
+
+                for node in idx:
+                    if isinstance(node, int):
+                        layer = np.searchsorted(nodelay_cumsum, node)
+                        nnode = node - nodelay_cumsum[layer - 1]
+                        row += [data[layer - 1][nnode]]
+                    else:
+                        errmsg = "idx must be an integer or a list of integers"
+                        raise Exception(errmsg)
+
+                result.append(row)
+
+        else:
+            errmsg = "idx must be an integer or a list of integers"
+            raise Exception(errmsg)
+
+        return np.array(result)