add compare check for example outputs

.pre-commit-config.yaml

@@ -5,7 +5,9 @@ repos:
   rev: v3.2.0
   hooks:
     - id: trailing-whitespace
+      exclude: ^test/__compare__
     - id: end-of-file-fixer
+      exclude: ^test/__compare__
     - id: check-yaml
     - id: check-added-large-files
     - id: check-toml

docs/examples/twri.py

@@ -173,7 +173,7 @@ def plot_head(head, workspace):
 tdis = flopy4.mf6.simulation.Tdis.from_time(time)
 
 # Create workspace
-workspace = Path(__file__).parent / "twri"
+workspace = Path(__file__).parent / "twri" / "list_stresspkg"
 workspace.mkdir(parents=True, exist_ok=True)
 
 # Create simulation
@@ -260,7 +260,7 @@ def plot_head(head, workspace):
 gwf.rcha = [rcha]
 
 # create new workspace
-workspace = Path(__file__).parent / "twri2"
+workspace = Path(__file__).parent / "twri" / "array_stresspkg"
 workspace.mkdir(parents=True, exist_ok=True)
 sim.workspace = workspace
 

test/test_examples.py

@@ -2,6 +2,7 @@
 
 import sys
 
+import numpy as np
 import pytest
 from conftest import EXAMPLES_PATH
 from modflow_devtools.misc import run_cmd
@@ -19,8 +20,73 @@ def pytest_generate_tests(metafunc):
         metafunc.parametrize("example_script", scripts.values(), ids=scripts.keys())
 
 
+def compare_hds(compare_fpth, check_path):
+    from flopy.utils import HeadFile
+
+    if compare_fpth.suffix == ".hds":
+        hds_compare = HeadFile(compare_fpth, precision="double").get_data()
+    elif compare_fpth.suffix == ".npy":
+        hds_compare = np.load(compare_fpth)
+
+    # check *.hds files
+    for f in check_path.rglob("*.hds"):
+        hds = HeadFile(f, precision="double")
+        assert np.allclose(hds_compare, hds.get_data())
+
+
+def compare_grb(compare_fpth, check_path):
+    from flopy.mf6.utils import MfGrdFile
+
+    grb_compare = MfGrdFile(compare_fpth)
+
+    # check *.grb files
+    for f in check_path.rglob("*.grb"):
+        grb = MfGrdFile(f)
+        np.testing.assert_equal(grb_compare._datadict, grb._datadict)
+
+
+def compare_bud(compare_fpth, check_path):
+    from flopy.utils.binaryfile import CellBudgetFile
+
+    bud_compare = CellBudgetFile(compare_fpth)
+
+    # check *.bud and *.cbc files
+    pattern = ["*.bud", "*.cbc"]
+    files = [f for p in pattern for f in check_path.rglob(p)]
+    for f in files:
+        bud = CellBudgetFile(f)
+        names = bud_compare.get_unique_record_names()
+        for n in names:
+            bud_cmp = bud_compare.get_data(text=n)
+            bud_chk = bud.get_data(text=n)
+            if isinstance(bud_cmp[0], np.rec.recarray):
+                assert len(bud_cmp) == len(bud_chk)
+                for ra_cmp, ra in zip(bud_cmp, bud_chk):
+                    assert len(ra_cmp) == len(ra)
+                    for i, rec in enumerate(ra_cmp):
+                        assert all(np.allclose(x, y) for x, y in zip(rec, ra[i]))
+            else:
+                assert np.allclose(bud_compare.get_data(text=n), bud.get_data(text=n))
+
+
+def compare(example_script):
+    from pathlib import Path
+
+    check_path = Path(f"{example_script.parent}/{example_script.stem}")
+    compare_path = Path(f"{example_script.parent.parent.parent}/test/__compare__/test_examples")
+
+    for f in compare_path.glob(f"{example_script.stem}.*"):
+        if f.suffix == ".bud" or f.suffix == ".cbc":
+            compare_bud(f, check_path)
+        if f.suffix == ".hds" or f.match("*hds.npy"):
+            compare_hds(f, check_path)
+        if f.suffix == ".grb":
+            compare_grb(f, check_path)
+
+
 @pytest.mark.slow
 def test_scripts(example_script):
     args = [sys.executable, example_script]
     stdout, stderr, retcode = run_cmd(*args, verbose=True)
     assert not retcode, stdout + stderr
+    compare(example_script)

test/test_mf6_codec.py

@@ -302,6 +302,55 @@ def test_dumps_chdg():
     pprint(loaded)
 
 
+def test_dumps_rcha():
+    from flopy4.mf6.gwf import Dis, Gwf, Rcha
+
+    nlay = 3
+    nrow = 10
+    ncol = 10
+
+    dis = Dis(nlay=nlay, nrow=nrow, ncol=ncol)
+    gwf = Gwf(dis=dis)
+
+    recharge = np.full((nrow, ncol), FILL_DNODATA, dtype=float)
+    irch = np.full((nrow, ncol), 1, dtype=int)
+    irch[0, 0] = 2
+    irch[9, 9] = 2
+    recharge[0, 0] = 1.0
+    recharge[9, 9] = 0.0
+    rch = Rcha(
+        parent=gwf,
+        irch=np.expand_dims(irch.ravel(), axis=0),
+        recharge=np.expand_dims(recharge.ravel(), axis=0),
+        save_flows=True,
+        print_input=True,
+        dims={"nper": 1},
+    )
+
+    dumped = dumps(COMPONENT_CONVERTER.unstructure(rch))
+    print("RCH dump:")
+    print(dumped)
+
+    assert "BEGIN PERIOD 1" in dumped
+    assert "END PERIOD 1" in dumped
+
+    period_section = dumped.split("BEGIN PERIOD 1")[1].split("END PERIOD 1")[0].strip()
+    lines = [line.strip() for line in period_section.split("\n") if line.strip()]
+
+    assert len(lines) == 6
+    assert "READASARRAYS" in dumped
+    dump_irch = [int(x) for x in lines[2].split()]
+    dump_lidx = np.array(dump_irch)
+    assert np.allclose(irch, dump_lidx.reshape(nrow, ncol))
+    dump_rch = [float(x) for x in lines[5].split()]
+    dump_recharge = np.array(dump_rch)
+    assert np.allclose(recharge, dump_recharge.reshape(nrow, ncol))
+
+    loaded = loads(dumped)
+    print("RCHA load:")
+    pprint(loaded)
+
+
 def test_dumps_wel():
     from flopy4.mf6.gwf import Dis, Gwf, Wel
 

test/test_mf6_component.py

@@ -435,15 +435,15 @@ def test_quickstart_grid(function_tmpdir):
     )
     npf = Npf(parent=gwf, icelltype=0, k=1.0)
 
-    # chd grid based input, step 1 data head array
-    head = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float)
-    head[0, 0, 0] = 1.0
-    head[0, 9, 9] = 0.0
-    # TODO: support dict style input keyed on SP with separate grid arrays
-    chd = Chdg(parent=gwf, head=np.expand_dims(head.ravel(), axis=0))
-    # headnone = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float)
-    # ts_head = np.stack((head.ravel(), headnone.ravel()), axis=0)
-    # chd = Chdg(parent=gwf, head=ts_head)
+    # Chdg
+    GRID_NODATA = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float)
+    head = np.repeat(np.expand_dims(GRID_NODATA, axis=0), repeats=1, axis=0)
+    head[0, 0, 0, 0] = 1.0
+    head[0, 0, 9, 9] = 0.0
+    chd = Chdg(
+        parent=gwf,
+        head=head.reshape(1, -1),
+    )
 
     sim.write()
     sim.run()