From 2bb0c64dc79dc05ea573dbb55e99fccdd88f4eaa Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Fri, 7 Nov 2025 07:38:24 -0500 Subject: [PATCH 1/4] fix(cellbudgetfile): fix get_ts support for aux vars --- autotest/test_cellbudgetfile.py | 264 ++++++++++++++++++++++++----- flopy/utils/binaryfile/__init__.py | 190 +++++++++++---------- 2 files changed, 327 insertions(+), 127 deletions(-) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index dd61a4e72c..b56c20b964 100644 --- a/autotest/test_cellbudgetfile.py +++ b/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" diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index f0781d56a3..dfa851a510 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -1610,10 +1610,14 @@ def get_ts(self, idx, text=None, times=None, variable="q"): 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, tuple of ints, or list of such + Acceptable values depend on grid type: + + - Structured grids (DIS): (layer, row, column) or list of such + - Vertex grids (DISV): (layer, cellid) or list of such + - Unstructured grids (DISU): node number, (node number,) or list of such + + All indices must be zero-based. text : str The text identifier for the record. Examples include 'RIVER LEAKAGE', 'STORAGE', 'FLOW RIGHT FACE', etc. @@ -1631,33 +1635,40 @@ def get_ts(self, idx, text=None, times=None, variable="q"): Array has size (ntimes, ncells + 1). The first column in the data array will contain time (totim). - See Also - -------- - Notes ----- - The layer, row, and column values must be zero-based, and must be - within the following ranges: 0 <= k < nlay; 0 <= i < nrow; 0 <= j < ncol + Index ranges (zero-based): + + - DIS: 0 <= layer < nlay, 0 <= row < nrow, 0 <= col < ncol + - DISV: 0 <= layer < nlay, 0 <= cellid < ncpl + - DISU: 0 <= node < nnodes Examples -------- - >>> # Get specific discharge x-component from DATA-SPDIS record + >>> # DIS grid: layer 0, row 5, column 5 >>> ts = cbb.get_ts(idx=(0, 5, 5), text='DATA-SPDIS', variable='qx') + >>> # DISV grid: layer 0, cell 12 + >>> ts = cbb.get_ts(idx=(0, 12), text='DATA-SPDIS', variable='qx') + + >>> # DISU grid: node 10 + >>> ts = cbb.get_ts(idx=10, text='FLOW-JA-FACE') + """ - # issue exception if text not provided if text is None: - raise Exception( + raise ValueError( "text keyword must be provided to CellBudgetFile get_ts() method." ) - kijlist = self._build_kijlist(idx) - nstation = self._get_nstation(idx, kijlist) + cellids = self._cellids(idx) + ncells = len(cellids) - # Initialize result array and put times in first column - result = self._init_result(nstation) + result = np.empty((len(self.kstpkper), ncells + 1), dtype=self.realtype) + result[:, :] = np.nan + if len(self.times) == result.shape[0]: + result[:, 0] = np.array(self.times) timesint = self.get_times() kstpkper = self.get_kstpkper() @@ -1674,53 +1685,36 @@ def get_ts(self, idx, text=None, times=None, variable="q"): f"number of time steps in cell budget file ({nsteps})" ) timesint = times - for idx, t in enumerate(timesint): - result[idx, 0] = t - - full3d = True - for itim, k in enumerate(kstpkper): - if full3d: - try: - v = self.get_data(kstpkper=k, text=text, full3D=True) + for i, t in enumerate(timesint): + result[i, 0] = t - # skip missing data - required for storage - if len(v) == 0: - continue + use_full3d = variable == "q" and ( + self.modelgrid is None or self.modelgrid.grid_type == "structured" + ) - v = v[0] - istat = 1 - for k, i, j in kijlist: - result[itim, istat] = v[k, i, j].copy() - istat += 1 + for itim, kstpkper_ in enumerate(kstpkper): + if use_full3d: + v = self.get_data(kstpkper=kstpkper_, text=text, full3D=True) + if len(v) == 0: continue - except ValueError: - full3d = False - if not full3d: - v = self.get_data(kstpkper=k, text=text) - # skip missing data - required for storage + v = v[0] + istat = 1 + for k, i, j in cellids: + result[itim, istat] = v[k, i, j].copy() + istat += 1 + else: + v = self.get_data(kstpkper=kstpkper_, text=text) if len(v) == 0: continue if self.modelgrid is None: - s = ( + raise ValueError( "A modelgrid instance must be provided during " "instantiation to get IMETH=6 timeseries data" ) - raise ValueError(s) - - if self.modelgrid.grid_type == "structured": - ndx = [ - lrc[0] * (self.modelgrid.nrow * self.modelgrid.ncol) - + lrc[1] * self.modelgrid.ncol - + (lrc[2] + 1) - for lrc in kijlist - ] - else: - ndx = [ - lrc[0] * self.modelgrid.ncpl + (lrc[-1] + 1) for lrc in kijlist - ] + nodes = self._cellid_to_node(cellids) for vv in v: available = vv.dtype.names if variable not in available: @@ -1729,51 +1723,77 @@ def get_ts(self, idx, text=None, times=None, variable="q"): f"Available variables: {list(available)}" ) - dix = np.asarray(np.isin(vv["node"], ndx)).nonzero()[0] + dix = np.asarray(np.isin(vv["node"], nodes)).nonzero()[0] if len(dix) > 0: result[itim, 1:] = vv[variable][dix] return result - def _build_kijlist(self, idx): - if isinstance(idx, list): - kijlist = idx + def _cellids(self, idx): + if isinstance(idx, int): + idx_list = [idx] elif isinstance(idx, tuple): - kijlist = [idx] + idx_list = [idx] + elif isinstance(idx, list): + idx_list = idx else: - raise Exception("Could not build kijlist from ", idx) + raise TypeError( + f"Invalid index type, expected int, tuple " + f"or list of such, got {type(idx)}" + ) - # Check to make sure that k, i, j are within range, otherwise - # the seek approach won't work. Can't use k = -1, for example. - for k, i, j in kijlist: - fail = False - if k < 0 or k > self.nlay - 1: - fail = True - if i < 0 or i > self.nrow - 1: - fail = True - if j < 0 or j > self.ncol - 1: - fail = True - if fail: - raise Exception( - "Invalid cell index. Cell {} not within model grid: {}".format( - (k, i, j), (self.nlay, self.nrow, self.ncol) + grid_type = "structured" if self.modelgrid is None else self.modelgrid.grid_type + + cellid = [] + for item in idx_list: + if grid_type == "structured": + if not isinstance(item, tuple) or len(item) != 3: + raise ValueError( + f"Expected DIS cell index (layer, row, col), got: {item}" ) - ) - return kijlist + k, i, j = item + if k < 0 or k >= self.nlay: + raise ValueError(f"Layer index {k} out of range [0, {self.nlay})") + if i < 0 or i >= self.nrow: + raise ValueError(f"Row index {i} out of range [0, {self.nrow})") + if j < 0 or j >= self.ncol: + raise ValueError(f"Column index {j} out of range [0, {self.ncol})") + cellid.append((k, i, j)) + elif grid_type == "vertex": + if not isinstance(item, tuple) or len(item) != 2: + raise ValueError( + f"Expected DISV cell index (layer, cellid), got: {item}" + ) + k, cell = item + if k < 0 or k >= self.nlay: + raise Exception(f"Layer index {k} out of range [0, {self.nlay})") + if cell < 0 or cell >= self.modelgrid.ncpl: + raise ValueError( + f"Cell index {cell} out of range [0, {self.modelgrid.ncpl})" + ) + cellid.append((k, cell)) + else: + if not isinstance(item, (int, np.integer)): + raise Exception(f"Expected DISU node number , got: {item}") + node = int(item) + if node < 0 or node >= self.modelgrid.nnodes: + raise Exception( + f"Node index {node} out of range [0, {self.modelgrid.nnodes})" + ) + cellid.append(node) - def _get_nstation(self, idx, kijlist): - if isinstance(idx, list): - return len(kijlist) - elif isinstance(idx, tuple): - return 1 + return cellid - def _init_result(self, nstation): - # Initialize result array and put times in first column - result = np.empty((len(self.kstpkper), nstation + 1), dtype=self.realtype) - result[:, :] = np.nan - if len(self.times) == result.shape[0]: - result[:, 0] = np.array(self.times) - return result + def _cellid_to_node(self, cellids) -> list[int]: + if self.modelgrid.grid_type == "structured": + return [ + k * (self.nrow * self.ncol) + i * self.ncol + j + 1 + for k, i, j in cellids + ] + elif self.modelgrid.grid_type == "vertex": + return [k * self.modelgrid.ncpl + cell + 1 for k, cell in cellids] + else: + return [node + 1 for node in cellids] def get_record(self, idx, full3D=False): """ From f98516a053e506cda19aeee1363062d528ffe93c Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Wed, 12 Nov 2025 07:25:48 -0500 Subject: [PATCH 2/4] testing wip --- autotest/test_cellbudgetfile.py | 120 +++++++++++++++++++++++--------- 1 file changed, 86 insertions(+), 34 deletions(-) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index b56c20b964..aacbab9b6b 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -687,6 +687,47 @@ def test_cellbudgetfile_get_ts_aux_vars_mf2005(example_data_path): assert np.array_equal(ts_default[:, 0], ts_iface[:, 0], equal_nan=True) +def test_cellbudgetfile_get_ts_aux_vars_mf6_readme_example(function_tmpdir): + from flopy.mf6 import ( + MFSimulation, + ModflowGwf, + ModflowGwfchd, + ModflowGwfdis, + ModflowGwfic, + ModflowGwfnpf, + ModflowGwfoc, + ModflowIms, + ModflowTdis, + ) + + name = 'mymodel' + sim = MFSimulation(sim_name=name, sim_ws=function_tmpdir, exe_name='mf6') + tdis = ModflowTdis(sim) + ims = ModflowIms(sim) + gwf = ModflowGwf(sim, modelname=name, save_flows=True) + dis = ModflowGwfdis(gwf, nrow=10, ncol=10) + ic = ModflowGwfic(gwf) + npf = ModflowGwfnpf(gwf, save_specific_discharge=True) + chd = ModflowGwfchd(gwf, stress_period_data=[[(0, 0, 0), 1.], + [(0, 9, 9), 0.]]) + budget_file = name + '.bud' + head_file = name + '.hds' + oc = ModflowGwfoc(gwf, + budget_filerecord=budget_file, + head_filerecord=head_file, + saverecord=[('HEAD', 'ALL'), ('BUDGET', 'ALL')]) + sim.write_simulation(silent=True) + sim.run_simulation(silent=True) + + hds = gwf.output.head().get_data() + cbc = gwf.output.budget() + + cellid = (0,5,5) + + head = hds.get_ts(idx=cellid) + spdis = cbc.get_ts(idx=cellid, text='DATA-SPDIS') + + @pytest.fixture def dis_sim(function_tmpdir): from flopy.mf6 import ( @@ -742,54 +783,61 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_dis(dis_sim): 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") + text = "DATA-SPDIS" + spdis = cbc.get_data(text=text) + for field in ["node", "q", "qx", "qy", "qz"]: + assert field in spdis[0].dtype.names - 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 + text = "CHD" + chd = cbc.get_data(text=text) + for field in ["node", "node2", "q"]: + assert field in chd[0].dtype.names + + spdis_q = cbc.get_ts(idx=cellid, text=text, variable="q") + spdis_qx = cbc.get_ts(idx=cellid, text=text, variable="qx") + spdis_qy = cbc.get_ts(idx=cellid, text=text, variable="qy") + spdis_qz = cbc.get_ts(idx=cellid, text=text, variable="qz") + + chd_q = cbc.get_ts(idx=cellid, ) + + assert spdis_q.shape == spdis_qx.shape == spdis_qy.shape == spdis_qz.shape + assert spdis_q.shape[1] == 2 # time + 1 data column + assert spdis_qx.shape[1] == 2 + assert spdis_qy.shape[1] == 2 + assert spdis_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]) + assert np.array_equal(spdis_q[:, 0], spdis_qx[:, 0]) + assert np.array_equal(spdis_q[:, 0], spdis_qy[:, 0]) + assert np.array_equal(spdis_q[:, 0], spdis_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], + spdis_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], + spdis_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], + spdis_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], + spdis_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" + assert not np.allclose(spdis_qx[:, 1], 0.0), "qx should have non-zero flow" + assert not np.allclose(spdis_qy[:, 1], 0.0), "qy should have non-zero flow" + assert np.allclose(spdis_q[:, 1], 0.0), "q should be zero for internal cells" + assert np.allclose(spdis_qz[:, 1], 0.0), "qz should be zero for single layer" @pytest.mark.requires_exe("mf6") @@ -822,12 +870,10 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_disv(dis_sim): 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 + for field in ["node", "q", "qx", "qy", "qz"]: + assert field in spdis[0].dtype.names - cellid = (0, 4) # cell in center of layer 0 + cellid = (0, 4) nn = 4 ts_q = cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="q") @@ -866,6 +912,10 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_disv(dis_sim): 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" + chd = cbc.get_data(text="CHD") + for field in ["node", "node2", "q"]: + assert field in chd[0].dtype.names + @pytest.mark.requires_exe("mf6") def test_cellbudgetfile_get_ts_aux_vars_mf6_disu(dis_sim): @@ -898,10 +948,8 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_disu(dis_sim): 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 + for field in ["node", "q", "qx", "qy", "qz"]: + assert field in spdis[0].dtype.names ts_q = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="q") ts_qx = cbc.get_ts(idx=nn, text="DATA-SPDIS", variable="qx") @@ -938,3 +986,7 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_disu(dis_sim): 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" + + chd = cbc.get_data(text="CHD") + for field in ["node", "node2", "q"]: + assert field in chd[0].dtype.names From ab1594ae5ae0324c859e667f2fb92483a0db7f2d Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Wed, 3 Dec 2025 13:39:01 -0500 Subject: [PATCH 3/4] support list or tuple cell ids --- autotest/test_cellbudgetfile.py | 29 ++++++++++++++++------------- flopy/pakbase.py | 2 +- flopy/utils/binaryfile/__init__.py | 4 ++-- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index aacbab9b6b..188eac1cc2 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -700,32 +700,33 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_readme_example(function_tmpdir): ModflowTdis, ) - name = 'mymodel' - sim = MFSimulation(sim_name=name, sim_ws=function_tmpdir, exe_name='mf6') + name = "mymodel" + sim = MFSimulation(sim_name=name, sim_ws=function_tmpdir, exe_name="mf6") tdis = ModflowTdis(sim) ims = ModflowIms(sim) gwf = ModflowGwf(sim, modelname=name, save_flows=True) dis = ModflowGwfdis(gwf, nrow=10, ncol=10) ic = ModflowGwfic(gwf) npf = ModflowGwfnpf(gwf, save_specific_discharge=True) - chd = ModflowGwfchd(gwf, stress_period_data=[[(0, 0, 0), 1.], - [(0, 9, 9), 0.]]) - budget_file = name + '.bud' - head_file = name + '.hds' - oc = ModflowGwfoc(gwf, - budget_filerecord=budget_file, - head_filerecord=head_file, - saverecord=[('HEAD', 'ALL'), ('BUDGET', 'ALL')]) + chd = ModflowGwfchd(gwf, stress_period_data=[[(0, 0, 0), 1.0], [(0, 9, 9), 0.0]]) + budget_file = name + ".bud" + head_file = name + ".hds" + oc = ModflowGwfoc( + gwf, + budget_filerecord=budget_file, + head_filerecord=head_file, + saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")], + ) sim.write_simulation(silent=True) sim.run_simulation(silent=True) hds = gwf.output.head().get_data() cbc = gwf.output.budget() - cellid = (0,5,5) + cellid = (0, 5, 5) head = hds.get_ts(idx=cellid) - spdis = cbc.get_ts(idx=cellid, text='DATA-SPDIS') + spdis = cbc.get_ts(idx=cellid, text="DATA-SPDIS") @pytest.fixture @@ -802,7 +803,9 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_dis(dis_sim): spdis_qy = cbc.get_ts(idx=cellid, text=text, variable="qy") spdis_qz = cbc.get_ts(idx=cellid, text=text, variable="qz") - chd_q = cbc.get_ts(idx=cellid, ) + chd_q = cbc.get_ts( + idx=cellid, + ) assert spdis_q.shape == spdis_qx.shape == spdis_qy.shape == spdis_qz.shape assert spdis_q.shape[1] == 2 # time + 1 data column diff --git a/flopy/pakbase.py b/flopy/pakbase.py index 70c6a23ddc..71a7f38e69 100644 --- a/flopy/pakbase.py +++ b/flopy/pakbase.py @@ -515,7 +515,7 @@ def __getitem__(self, item): # Oc88 but is not a MfList spd = getattr(self, "stress_period_data") if isinstance(item, MfList): - if not isinstance(item, list) and not isinstance(item, tuple): + if not isinstance(item, (list, tuple)): msg = f"package.__getitem__() kper {item} not in data.keys()" assert item in list(spd.data.keys()), msg return spd[item] diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index dfa851a510..0d53f8a018 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -1747,7 +1747,7 @@ def _cellids(self, idx): cellid = [] for item in idx_list: if grid_type == "structured": - if not isinstance(item, tuple) or len(item) != 3: + if not isinstance(item, (list, tuple)) or len(item) != 3: raise ValueError( f"Expected DIS cell index (layer, row, col), got: {item}" ) @@ -1760,7 +1760,7 @@ def _cellids(self, idx): raise ValueError(f"Column index {j} out of range [0, {self.ncol})") cellid.append((k, i, j)) elif grid_type == "vertex": - if not isinstance(item, tuple) or len(item) != 2: + if not isinstance(item, (list, tuple)) or len(item) != 2: raise ValueError( f"Expected DISV cell index (layer, cellid), got: {item}" ) From 937f86f0694a6a453a43f5fe04817fc29e96676b Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Wed, 3 Dec 2025 14:44:26 -0500 Subject: [PATCH 4/4] avoid rounding: use existing dtype in full 3D conversion --- autotest/test_cellbudgetfile.py | 55 +++++++++++++++++++----------- flopy/utils/binaryfile/__init__.py | 2 +- 2 files changed, 37 insertions(+), 20 deletions(-) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index 188eac1cc2..81c3bb1fb8 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -720,13 +720,21 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_readme_example(function_tmpdir): sim.write_simulation(silent=True) sim.run_simulation(silent=True) - hds = gwf.output.head().get_data() + hds = gwf.output.head() cbc = gwf.output.budget() cellid = (0, 5, 5) - - head = hds.get_ts(idx=cellid) - spdis = cbc.get_ts(idx=cellid, text="DATA-SPDIS") + assert ( + hds.get_ts(idx=cellid)[0][1] == hds.get_data()[cellid[0], cellid[1], cellid[2]] + ) + assert ( + cbc.get_ts(idx=cellid, text="DATA-SPDIS", variable="qx")[0][1] + == cbc.get_data(text="DATA-SPDIS")[0][gwf.modelgrid.get_node([cellid])[0]]["qx"] + ) + cellid = (0, 0, 0) + assert ( + cbc.get_ts(idx=cellid, text="CHD")[0][1] == cbc.get_data(text="CHD")[0][0]["q"] + ) @pytest.fixture @@ -785,28 +793,19 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_dis(dis_sim): gwf = sim.get_model() cbc = gwf.output.budget() - cellid = (0, 2, 2) - nn = gwf.modelgrid.get_node(cellid)[0] - text = "DATA-SPDIS" spdis = cbc.get_data(text=text) for field in ["node", "q", "qx", "qy", "qz"]: assert field in spdis[0].dtype.names - text = "CHD" - chd = cbc.get_data(text=text) - for field in ["node", "node2", "q"]: - assert field in chd[0].dtype.names + cellid = (0, 2, 2) + nn = gwf.modelgrid.get_node(cellid)[0] spdis_q = cbc.get_ts(idx=cellid, text=text, variable="q") spdis_qx = cbc.get_ts(idx=cellid, text=text, variable="qx") spdis_qy = cbc.get_ts(idx=cellid, text=text, variable="qy") spdis_qz = cbc.get_ts(idx=cellid, text=text, variable="qz") - chd_q = cbc.get_ts( - idx=cellid, - ) - assert spdis_q.shape == spdis_qx.shape == spdis_qy.shape == spdis_qz.shape assert spdis_q.shape[1] == 2 # time + 1 data column assert spdis_qx.shape[1] == 2 @@ -823,25 +822,43 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_dis(dis_sim): assert np.allclose( spdis_q[i, 1], rec["q"][mask][0], - ), f"get_ts() q value doesn't match get_data() at time {i}" + ), f"get_ts() SPDIS q value doesn't match get_data() at time {i}" assert np.allclose( spdis_qx[i, 1], rec["qx"][mask][0], - ), f"get_ts() qx value doesn't match get_data() at time {i}" + ), f"get_ts() SPDIS qx value doesn't match get_data() at time {i}" assert np.allclose( spdis_qy[i, 1], rec["qy"][mask][0], - ), f"get_ts() qy value doesn't match get_data() at time {i}" + ), f"get_ts() SPDIS qy value doesn't match get_data() at time {i}" assert np.allclose( spdis_qz[i, 1], rec["qz"][mask][0], - ), f"get_ts() qz value doesn't match get_data() at time {i}" + ), f"get_ts() SPDIS qz value doesn't match get_data() at time {i}" assert not np.allclose(spdis_qx[:, 1], 0.0), "qx should have non-zero flow" assert not np.allclose(spdis_qy[:, 1], 0.0), "qy should have non-zero flow" assert np.allclose(spdis_q[:, 1], 0.0), "q should be zero for internal cells" assert np.allclose(spdis_qz[:, 1], 0.0), "qz should be zero for single layer" + text = "CHD" + chd = cbc.get_data(text=text) + for field in ["node", "node2", "q"]: + assert field in chd[0].dtype.names + + cellid = (0, 0, 0) + nn = gwf.modelgrid.get_node(cellid)[0] + + chd_q = cbc.get_ts(idx=cellid, text=text) + + # check get_ts() values match get_data() for each time step + for i, rec in enumerate(chd): + mask = rec["node"] == nn + 1 # 1-based + assert np.allclose( + chd_q[i, 1], + rec["q"][mask][0], + ), f"get_ts() CHD q value doesn't match get_data() at time {i}" + @pytest.mark.requires_exe("mf6") def test_cellbudgetfile_get_ts_aux_vars_mf6_disv(dis_sim): diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 0d53f8a018..0f089a21f7 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -1978,7 +1978,7 @@ def __create3D(self, data): List contains unique simulation times (totim) in binary file. """ - out = np.ma.zeros(self.nnodes, dtype=np.float32) + out = np.ma.zeros(self.nnodes, dtype=data["q"].dtype) out.mask = True for [node, q] in zip(data["node"], data["q"]): idx = node - 1