fix(cellbudgetfile): fix get_ts support for aux vars

autotest/test_cellbudgetfile.py

@@ -7,6 +7,7 @@
 from autotest.conftest import get_example_data_path
 from flopy.mf6.modflow.mfsimulation import MFSimulation
 from flopy.utils.binaryfile import CellBudgetFile
+from flopy.utils.gridutil import get_disu_kwargs, get_disv_kwargs
 
 # test low-level CellBudgetFile._build_index() method
 
@@ -658,34 +659,36 @@ def test_read_mf6_budgetfile(example_data_path):
     assert len(rch_zone_3) == 120 * 3 + 1
 
 
-def test_cellbudgetfile_get_ts_auxiliary_variables(example_data_path):
+def test_cellbudgetfile_get_ts_aux_vars_mf2005(example_data_path):
     from flopy.discretization import StructuredGrid
 
     mf2005_model_path = example_data_path / "mf2005_test"
     cbc_fname = mf2005_model_path / "test1tr.gitcbc"
 
-    # modelgrid for test1tr (1 layer, 15 rows, 10 cols)
-    modelgrid = StructuredGrid(nrow=15, ncol=10, nlay=1)
+    nlay = 1
+    nrow = 15
+    ncol = 10
+    grid = StructuredGrid(nlay=nlay, nrow=nrow, ncol=ncol)
 
-    with CellBudgetFile(cbc_fname, modelgrid=modelgrid) as v:
-        node = 54 - 1
-        k = node // (15 * 10)
-        i = (node % (15 * 10)) // 10
-        j = node % 10
+    cbc = CellBudgetFile(cbc_fname, modelgrid=grid)
+    node = 53
+    cellid = grid.get_lrc(node)[0]
 
-        ts_default = v.get_ts(idx=(k, i, j), text="WELLS")
-        ts_q_explicit = v.get_ts(idx=(k, i, j), text="WELLS", variable="q")
-        assert ts_default.shape[1] == 2  # time + 1 data column
-        np.testing.assert_array_equal(ts_default, ts_q_explicit)
+    ts_default = cbc.get_ts(idx=cellid, text="WELLS")
+    ts_q_explicit = cbc.get_ts(idx=cellid, text="WELLS", variable="q")
+    assert ts_default.shape == ts_q_explicit.shape
+    assert ts_default.shape[1] == 2  # time + 1 data column
+    assert np.array_equal(ts_default, ts_q_explicit, equal_nan=True)
 
-        ts_iface = v.get_ts(idx=(k, i, j), text="WELLS", variable="IFACE")
-        assert ts_iface.shape[1] == 2  # time + 1 data column
-        assert not np.array_equal(ts_default[:, 1], ts_iface[:, 1])
-        np.testing.assert_array_equal(ts_default[:, 0], ts_iface[:, 0])
+    ts_iface = cbc.get_ts(idx=cellid, text="WELLS", variable="IFACE")
+    assert ts_iface.shape[1] == 2  # time + 1 data column
+    assert ts_default.shape == ts_iface.shape
+    assert not np.array_equal(ts_default[:, 1], ts_iface[:, 1])
+    assert np.array_equal(ts_default[:, 0], ts_iface[:, 0], equal_nan=True)
 
 
-@pytest.mark.requires_exe("mf6")
-def test_cellbudgetfile_get_ts_spdis_auxiliary_variables(function_tmpdir):
+@pytest.fixture
+def dis_sim(function_tmpdir):
     from flopy.mf6 import (
         MFSimulation,
         ModflowGwf,
@@ -698,7 +701,7 @@ def test_cellbudgetfile_get_ts_spdis_auxiliary_variables(function_tmpdir):
         ModflowTdis,
     )
 
-    sim_name = "spdis_test"
+    sim_name = "test_ts_aux_vars"
     sim = MFSimulation(sim_name=sim_name, sim_ws=function_tmpdir, exe_name="mf6")
     tdis = ModflowTdis(sim, nper=2, perioddata=[(1.0, 1, 1.0), (1.0, 1, 1.0)])
     ims = ModflowIms(sim)
@@ -727,34 +730,211 @@ def test_cellbudgetfile_get_ts_spdis_auxiliary_variables(function_tmpdir):
         saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
     )
 
+    return sim
+
+
+@pytest.mark.requires_exe("mf6")
+def test_cellbudgetfile_get_ts_aux_vars_mf6_dis(dis_sim):
+    sim = dis_sim
     sim.write_simulation()
     success, _ = sim.run_simulation(silent=True)
+    assert success
+
+    gwf = sim.get_model()
+    cbc = gwf.output.budget()
+    spdis = cbc.get_data(text="DATA-SPDIS")
+    available_fields = list(spdis[0].dtype.names)
+    expected_fields = ["node", "q", "qx", "qy", "qz"]
+    for field in expected_fields:
+        assert field in available_fields
+
+    cellid = (0, 2, 2)
+    nn = gwf.modelgrid.get_node(cellid)[0]
+
+    ts_q = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="q")
+    ts_qx = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qx")
+    ts_qy = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qy")
+    ts_qz = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qz")
+
+    assert ts_q.shape == ts_qx.shape == ts_qy.shape == ts_qz.shape
+    assert ts_q.shape[1] == 2  # time + 1 data column
+    assert ts_qx.shape[1] == 2
+    assert ts_qy.shape[1] == 2
+    assert ts_qz.shape[1] == 2
+
+    assert np.array_equal(ts_q[:, 0], ts_qx[:, 0])
+    assert np.array_equal(ts_q[:, 0], ts_qy[:, 0])
+    assert np.array_equal(ts_q[:, 0], ts_qz[:, 0])
+
+    # check get_ts() values match get_data() for each time step
+    for i, rec in enumerate(spdis):
+        mask = rec["node"] == nn + 1  # 1-based
+        assert np.allclose(
+            ts_q[i, 1],
+            rec["q"][mask][0],
+        ), f"get_ts() q value doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qx[i, 1],
+            rec["qx"][mask][0],
+        ), f"get_ts() qx value doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qy[i, 1],
+            rec["qy"][mask][0],
+        ), f"get_ts() qy value doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qz[i, 1],
+            rec["qz"][mask][0],
+        ), f"get_ts() qz value doesn't match get_data() at time {i}"
+
+    assert not np.allclose(ts_qx[:, 1], 0.0), "qx should have non-zero flow"
+    assert not np.allclose(ts_qy[:, 1], 0.0), "qy should have non-zero flow"
+    assert np.allclose(ts_q[:, 1], 0.0), "q should be zero for internal cells"
+    assert np.allclose(ts_qz[:, 1], 0.0), "qz should be zero for single layer"
 
-    cbc_file = function_tmpdir / budget_file
-    with CellBudgetFile(cbc_file, modelgrid=gwf.modelgrid) as cbc:
-        spdis_data = cbc.get_data(text="DATA-SPDIS")
-        if len(spdis_data) == 0:
-            pytest.skip("No DATA-SPDIS records found")
 
-        available_fields = list(spdis_data[0].dtype.names)
-        expected_fields = ["node", "q", "qx", "qy", "qz"]
+@pytest.mark.requires_exe("mf6")
+def test_cellbudgetfile_get_ts_aux_vars_mf6_disv(dis_sim):
+    from flopy.mf6 import ModflowGwfchd, ModflowGwfdisv
+
+    sim = dis_sim
+    gwf = sim.get_model()
+    dis_grid = gwf.modelgrid
+    disv_kwargs = get_disv_kwargs(
+        nlay=dis_grid.nlay,
+        nrow=dis_grid.nrow,
+        ncol=dis_grid.ncol,
+        delr=dis_grid.delr,
+        delc=dis_grid.delc,
+        tp=dis_grid.top,
+        botm=dis_grid.botm,
+        xoff=dis_grid.xoffset,
+        yoff=dis_grid.yoffset,
+    )
+    gwf.remove_package("dis")
+    gwf.remove_package("chd")
+    disv = ModflowGwfdisv(gwf, **disv_kwargs)
+    chd_spd = [[(0, 0), 10.0], [(0, 24), 0.0]]
+    chd = ModflowGwfchd(gwf, stress_period_data=chd_spd)
 
-        for field in ["qx", "qy", "qz"]:
-            assert field in available_fields, (
-                f"Expected field '{field}' not found in {available_fields}"
-            )
+    sim.write_simulation()
+    success, _ = sim.run_simulation(silent=False)
+    assert success
+
+    cbc = gwf.output.budget()
+    spdis = cbc.get_data(text="DATA-SPDIS")
+    available_fields = list(spdis[0].dtype.names)
+    expected_fields = ["node", "q", "qx", "qy", "qz"]
+    for field in expected_fields:
+        assert field in available_fields
+
+    cellid = (0, 4)  # cell in center of layer 0
+    nn = 4
+
+    ts_q = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="q")
+    ts_qx = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qx")
+    ts_qy = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qy")
+    ts_qz = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qz")
+
+    assert ts_q.shape == ts_qx.shape == ts_qy.shape == ts_qz.shape
+    assert ts_q.shape[1] == 2  # time + 1 data column
+    assert ts_qx.shape[1] == 2
+    assert ts_qy.shape[1] == 2
+    assert ts_qz.shape[1] == 2
+
+    # check get_ts() values match get_data() for each time step
+    for i, rec in enumerate(spdis):
+        mask = rec["node"] == nn + 1  # 1-based
+        assert np.allclose(
+            ts_q[i, 1],
+            rec["q"][mask][0],
+        ), f"DISV get_ts() q doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qx[i, 1],
+            rec["qx"][mask][0],
+        ), f"DISV get_ts() qx doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qy[i, 1],
+            rec["qy"][mask][0],
+        ), f"DISV get_ts() qy doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qz[i, 1],
+            rec["qz"][mask][0],
+        ), f"DISV get_ts() qz doesn't match get_data() at time {i}"
+
+    assert not np.allclose(ts_qx[:, 1], 0.0), "qx should have non-zero flow"
+    assert not np.allclose(ts_qy[:, 1], 0.0), "qy should have non-zero flow"
+    assert np.allclose(ts_q[:, 1], 0.0), "q should be zero for internal cells"
+    assert np.allclose(ts_qz[:, 1], 0.0), "qz should be zero for single layer"
 
-        test_cell = (0, 2, 2)
-        ts_q = cbc.get_ts(idx=test_cell, text="DATA-SPDIS", variable="q")
-        ts_qx = cbc.get_ts(idx=test_cell, text="DATA-SPDIS", variable="qx")
-        ts_qy = cbc.get_ts(idx=test_cell, text="DATA-SPDIS", variable="qy")
-        ts_qz = cbc.get_ts(idx=test_cell, text="DATA-SPDIS", variable="qz")
 
-        assert ts_q.shape[1] == 2  # time + 1 data column
-        assert ts_qx.shape[1] == 2
-        assert ts_qy.shape[1] == 2
-        assert ts_qz.shape[1] == 2
+@pytest.mark.requires_exe("mf6")
+def test_cellbudgetfile_get_ts_aux_vars_mf6_disu(dis_sim):
+    from flopy.mf6 import ModflowGwfchd, ModflowGwfdisu
+
+    sim = dis_sim
+    gwf = sim.get_model()
+    dis_grid = gwf.modelgrid
+    disu_kwargs = get_disu_kwargs(
+        nlay=dis_grid.nlay,
+        nrow=dis_grid.nrow,
+        ncol=dis_grid.ncol,
+        delr=dis_grid.delr,
+        delc=dis_grid.delc,
+        tp=dis_grid.top,
+        botm=dis_grid.botm,
+        return_vertices=True,
+    )
+    gwf.remove_package("dis")
+    gwf.remove_package("chd")
+    disu = ModflowGwfdisu(gwf, **disu_kwargs)
+    chd_spd = [[0, 10.0], [24, 0.0]]
+    chd = ModflowGwfchd(gwf, stress_period_data=chd_spd)
 
-        np.testing.assert_array_equal(ts_q[:, 0], ts_qx[:, 0])
-        np.testing.assert_array_equal(ts_q[:, 0], ts_qy[:, 0])
-        np.testing.assert_array_equal(ts_q[:, 0], ts_qz[:, 0])
+    sim.write_simulation()
+    success, _ = sim.run_simulation(silent=False)
+    assert success
+
+    nn = 4
+
+    cbc = gwf.output.budget()
+    spdis = cbc.get_data(text="DATA-SPDIS")
+    available_fields = list(spdis[0].dtype.names)
+    expected_fields = ["node", "q", "qx", "qy", "qz"]
+    for field in expected_fields:
+        assert field in available_fields
+
+    ts_q = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="q")
+    ts_qx = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="qx")
+    ts_qy = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="qy")
+    ts_qz = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="qz")
+
+    assert ts_q.shape == ts_qx.shape == ts_qy.shape == ts_qz.shape
+    assert ts_q.shape[1] == 2  # time + 1 data column
+    assert ts_qx.shape[1] == 2
+    assert ts_qy.shape[1] == 2
+    assert ts_qz.shape[1] == 2
+
+    # check values match get_data()
+    for i, rec in enumerate(spdis):
+        mask = rec["node"] == nn + 1  # 1-based
+        assert np.allclose(
+            ts_q[i, 1],
+            rec["q"][mask][0],
+        ), f"DISU get_ts() q doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qx[i, 1],
+            rec["qx"][mask][0],
+        ), f"DISU get_ts() qx doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qy[i, 1],
+            rec["qy"][mask][0],
+        ), f"DISU get_ts() qy doesn't match get_data() at time {i}"
+        assert np.allclose(
+            ts_qz[i, 1],
+            rec["qz"][mask][0],
+        ), f"DISU get_ts() qz doesn't match get_data() at time {i}"
+
+    assert not np.allclose(ts_qx[:, 1], 0.0), "qx should have non-zero flow"
+    assert not np.allclose(ts_qy[:, 1], 0.0), "qy should have non-zero flow"
+    assert np.allclose(ts_q[:, 1], 0.0), "q should be zero for internal cells"
+    assert np.allclose(ts_qz[:, 1], 0.0), "qz should be zero for single layer"