feat(CellBudget): add support for full3D keyword

autotest/t004_test_utilarray.py

@@ -17,15 +17,15 @@
 
 def test_load_txt_free():
     a = np.ones((10,), dtype=np.float32) * 250.0
-    fp = StringIO(u"10*250.0")
+    fp = StringIO("10*250.0")
     fa = Util2d.load_txt(a.shape, fp, a.dtype, "(FREE)")
     np.testing.assert_equal(fa, a)
     assert fa.dtype == a.dtype
 
     a = np.arange(10, dtype=np.int32).reshape((2, 5))
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         0 1,2,3, 4
         5 6, 7,  8 9
     """
@@ -39,7 +39,7 @@ def test_load_txt_free():
     a[1, 0] = 2.2
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         5*1.0
         2.2 2*1.0, +1E-00 1.0
     """
@@ -54,7 +54,7 @@ def test_load_txt_fixed():
     a = np.arange(10, dtype=np.int32).reshape((2, 5))
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         01234X
         56789
     """
@@ -66,7 +66,7 @@ def test_load_txt_fixed():
 
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         0123X
         4
         5678
@@ -81,7 +81,7 @@ def test_load_txt_fixed():
     a = np.array([[-1, 1, -2, 2, -3], [3, -4, 4, -5, 5]], np.int32)
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         -1 1-2 2-3
         3 -44 -55
     """
@@ -96,7 +96,7 @@ def test_load_block():
     a = np.ones((2, 5), dtype=np.int32) * 4
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         1
         1 2 1 5 4
     """
@@ -111,7 +111,7 @@ def test_load_block():
     a[0, 2:4] = 6.0
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         3
         1 2 1 5 4.0
         1 2 2 2 9.0
@@ -127,7 +127,7 @@ def test_load_block():
     a[0, 2:4] = 8
     fp = StringIO(
         dedent(
-            u"""\
+            """\
         1
         1 1 3 4 8
     """

autotest/t050_test.py

@@ -46,8 +46,9 @@ def test_vtk_export_array2d():
     assert nlines == 17615
 
     # with smoothing
-    m.dis.top.export(output_dir, fmt="vtk", name="top_smooth",
-                     binary=False, smooth=True)
+    m.dis.top.export(
+        output_dir, fmt="vtk", name="top_smooth", binary=False, smooth=True
+    )
     filetocheck = os.path.join(output_dir, "top_smooth.vtk")
     nlines1 = count_lines_in_file(filetocheck)
     assert nlines1 == 17615
@@ -79,7 +80,11 @@ def test_vtk_export_array3d():
 
     # with point scalars
     m.upw.hk.export(
-        output_dir, fmt="vtk", name="hk_points", point_scalars=True, binary=False
+        output_dir,
+        fmt="vtk",
+        name="hk_points",
+        point_scalars=True,
+        binary=False,
     )
     filetocheck = os.path.join(output_dir, "hk_points.vtk")
     nlines1 = count_lines_in_file(filetocheck)
@@ -117,7 +122,9 @@ def test_vtk_transient_array_2d():
     kpers = [0, 1, 1096]
 
     # export and check
-    m.rch.rech.export(output_dir, fmt="vtk", kpers=kpers, binary=False, xml=True)
+    m.rch.rech.export(
+        output_dir, fmt="vtk", kpers=kpers, binary=False, xml=True
+    )
     filetocheck = os.path.join(output_dir, "rech_000001.vtk")
     nlines = count_lines_in_file(filetocheck)
     assert nlines == 26837
@@ -179,7 +186,9 @@ def test_vtk_export_packages():
     # transient package drain
     kpers = [0, 1, 1096]
     output_dir = os.path.join(cpth, "DRN")
-    m.drn.export(output_dir, fmt="vtk", binary=False, xml=True, kpers=kpers, pvd=True)
+    m.drn.export(
+        output_dir, fmt="vtk", binary=False, xml=True, kpers=kpers, pvd=True
+    )
     filetocheck = os.path.join(output_dir, "DRN_000001.vtu")
     nlines3 = count_lines_in_file(filetocheck)
     assert nlines3 == 27239
@@ -498,8 +507,9 @@ def test_vtk_vertex():
         return
 
     # disv test
-    workspace = os.path.join("..", "examples", "data", "mf6",
-                             "test003_gwfs_disv")
+    workspace = os.path.join(
+        "..", "examples", "data", "mf6", "test003_gwfs_disv"
+    )
     # outfile = os.path.join("vtk_transient_test", "vtk_pacakages")
     outfile = os.path.join("temp", "t050", "vtk_disv", "disv.vtk")
     sim = flopy.mf6.MFSimulation.load(sim_ws=workspace)
@@ -538,8 +548,9 @@ def test_vtk_pathline():
     ws = os.path.join("..", "examples", "data", "freyberg")
     modelpth = os.path.join("temp", "t050")
     outfile = os.path.join(modelpth, "pathline_test", "pathline.vtk")
-    ml = flopy.modflow.Modflow.load("freyberg.nam", model_ws=ws,
-                                    exe_name="mf2005")
+    ml = flopy.modflow.Modflow.load(
+        "freyberg.nam", model_ws=ws, exe_name="mf2005"
+    )
     ml.change_model_ws(new_pth=modelpth)
     ml.write_input()
     ml.run_model()
@@ -587,8 +598,8 @@ def test_vtk_pathline():
 
     maxtime = 0
     for p in plines:
-        if np.max(p['time']) > maxtime:
-            maxtime = np.max(p['time'])
+        if np.max(p["time"]) > maxtime:
+            maxtime = np.max(p["time"])
 
     if not len(totim) == 12054:
         raise AssertionError("Array size is incorrect for modpath VTK")
@@ -613,4 +624,4 @@ def test_vtk_pathline():
     test_vtk_vector()
     test_vtk_unstructured()
     test_vtk_vertex()
-    test_vtk_pathline()
+    test_vtk_pathline()

autotest/t079_test_cbc_full3D.py

@@ -0,0 +1,205 @@
+import os
+import sys
+import shutil
+
+import numpy as np
+import flopy
+
+ex_pths = (
+    os.path.join("..", "examples", "data", "freyberg"),
+    os.path.join("..", "examples", "data", "mf6-freyberg"),
+    os.path.join("..", "examples", "data", "mf6", "test003_gwfs_disv"),
+    os.path.join("..", "examples", "data", "mf6", "test003_gwftri_disv"),
+)
+ismf6_lst = ["mf6" in pth for pth in ex_pths]
+names = [os.path.basename(pth) for pth in ex_pths]
+
+tpth = os.path.join("temp", "t079")
+# delete the directory if it exists
+if os.path.isdir(tpth):
+    shutil.rmtree(tpth)
+# make the directory
+os.makedirs(tpth)
+
+mf6_exe = "mf6"
+mf2005_exe = "mf2005"
+if sys.platform == "win32":
+    mf6_exe += ".exe"
+    mf2005_exe += ".exe"
+
+
+def load_mf2005(name, ws_in):
+    name_file = f"{name}.nam"
+    ml = flopy.modflow.Modflow.load(
+        name_file,
+        model_ws=ws_in,
+        exe_name=mf2005_exe,
+        check=False,
+    )
+
+    # change work space
+    ws_out = os.path.join(tpth, name)
+    ml.change_model_ws(ws_out)
+
+    # save all budget data to a cell-by cell file
+    oc = ml.get_package("OC")
+    oc.reset_budgetunit()
+    oc.stress_period_data = {(0, 0): ["save budget"]}
+
+    return ml
+
+
+def load_mf6(name, ws_in):
+    sim = flopy.mf6.MFSimulation.load(
+        sim_name=name,
+        exe_name=mf6_exe,
+        sim_ws=ws_in,
+    )
+
+    # change work space
+    ws_out = os.path.join(tpth, name)
+    sim.set_sim_path(ws_out)
+
+    # get the groundwater flow model(s) and redefine the output control
+    # file to save cell by cell output
+    models = sim.model_names
+    for model in models:
+        gwf = sim.get_model(model)
+        gwf.name_file.save_flows = True
+
+        gwf.remove_package("oc")
+        budget_filerecord = f"{model}.cbc"
+        oc = flopy.mf6.ModflowGwfoc(
+            gwf,
+            budget_filerecord=budget_filerecord,
+            saverecord=[("BUDGET", "ALL")],
+        )
+
+    return sim
+
+
+def cbc_eval_size(cbcobj, nnodes, shape3d):
+    cbc_pth = cbcobj.filename
+
+    assert cbcobj.nnodes == nnodes, (
+        f"{cbc_pth} nnodes ({cbcobj.nnodes}) " f"does not equal {nnodes}"
+    )
+    a = np.squeeze(np.ones(cbcobj.shape, dtype=float))
+    b = np.squeeze(np.ones(shape3d, dtype=float))
+    assert a.shape == b.shape, (
+        f"{cbc_pth} shape {cbcobj.shape} " f"does not conform to {shape3d}"
+    )
+
+
+def cbc_eval_data(cbcobj, shape3d):
+    cbc_pth = cbcobj.filename
+    print(f"{cbc_pth}:\n")
+    cbcobj.list_unique_records()
+
+    names = cbcobj.get_unique_record_names(decode=True)
+    times = cbcobj.get_times()
+    for name in names:
+        text = name.strip()
+        arr = np.squeeze(
+            cbcobj.get_data(text=text, totim=times[0], full3D=True)[0]
+        )
+        if text != "FLOW-JA-FACE":
+            b = np.squeeze(np.ones(shape3d, dtype=float))
+            assert arr.shape == b.shape, (
+                f"{cbc_pth} shape {arr.shape} for '{text}' budget item "
+                f"does not conform to {shape3d}"
+            )
+
+
+def cbc_eval(cbcobj, nnodes, shape3d, modelgrid):
+    cbc_pth = cbcobj.filename
+    cbc_eval_size(cbcobj, nnodes, shape3d)
+    cbc_eval_data(cbcobj, shape3d)
+    cbcobj.close()
+
+    cobj_mg = flopy.utils.CellBudgetFile(
+        cbc_pth,
+        modelgrid=modelgrid,
+        verbose=True,
+    )
+    cbc_eval_size(cobj_mg, nnodes, shape3d)
+    cbc_eval_data(cobj_mg, shape3d)
+    cobj_mg.close()
+
+    return
+
+
+def clean_run(name):
+    ws = os.path.join(tpth, name)
+    if os.path.isdir(ws):
+        shutil.rmtree(ws)
+
+
+def mf6_eval(name, ws_in):
+
+    sim = load_mf6(name, ws_in)
+
+    # write the simulation
+    sim.write_simulation()
+
+    # run the simulation
+    sim.run_simulation()
+
+    # get the groundwater model and determine the size of the model grid
+    gwf_name = list(sim.model_names)[0]
+    gwf = sim.get_model(gwf_name)
+    nnodes, shape3d = gwf.modelgrid.nnodes, gwf.modelgrid.shape
+
+    # get the cell by cell object
+    cbc = gwf.output.budget()
+
+    # evaluate the full3D option
+    cbc_eval(cbc, nnodes, shape3d, gwf.modelgrid)
+
+    # clean the run
+    clean_run(name)
+
+    return
+
+
+def mf2005_eval(name, ws_in):
+    ml = load_mf2005(name, ws_in)
+
+    # write the model
+    ml.write_input()
+
+    # run the model
+    ml.run_model()
+
+    # determine the size of the model grid
+    nnodes, shape3d = ml.modelgrid.nnodes, ml.modelgrid.shape
+
+    # get the cell by cell object
+    fpth = os.path.join(tpth, name, f"{name}.cbc")
+    cbc = flopy.utils.CellBudgetFile(fpth)
+
+    # evaluate the full3D option
+    cbc_eval(cbc, nnodes, shape3d, ml.modelgrid)
+
+    # clean the run
+    clean_run(name)
+
+
+def test_cbc_full3D():
+    for (name, ismf6, ws_in) in zip(names, ismf6_lst, ex_pths):
+        if ismf6:
+            yield mf6_eval, name, ws_in
+        else:
+            yield mf2005_eval, name, ws_in
+
+
+def main():
+    for (name, ismf6, ws_in) in zip(names, ismf6_lst, ex_pths):
+        if ismf6:
+            mf6_eval(name, ws_in)
+        else:
+            mf2005_eval(name, ws_in)
+
+
+if __name__ == "__main__":
+    main()