diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index a8a660c89..dd61a4e72 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -656,3 +656,105 @@ def test_read_mf6_budgetfile(example_data_path): assert len(rch_zone_1) == 120 * 3 + 1 assert len(rch_zone_2) == 120 * 3 + 1 assert len(rch_zone_3) == 120 * 3 + 1 + + +def test_cellbudgetfile_get_ts_auxiliary_variables(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) + + with CellBudgetFile(cbc_fname, modelgrid=modelgrid) as v: + node = 54 - 1 + k = node // (15 * 10) + i = (node % (15 * 10)) // 10 + j = node % 10 + + 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_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]) + + +@pytest.mark.requires_exe("mf6") +def test_cellbudgetfile_get_ts_spdis_auxiliary_variables(function_tmpdir): + from flopy.mf6 import ( + MFSimulation, + ModflowGwf, + ModflowGwfchd, + ModflowGwfdis, + ModflowGwfic, + ModflowGwfnpf, + ModflowGwfoc, + ModflowIms, + ModflowTdis, + ) + + sim_name = "spdis_test" + 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) + gwf = ModflowGwf(sim, modelname=sim_name, save_flows=True) + nrow, ncol, nlay = 5, 5, 1 + dis = ModflowGwfdis( + gwf, + nrow=nrow, + ncol=ncol, + nlay=nlay, + delr=10.0, + delc=10.0, + top=10.0, + botm=[0.0], + ) + ic = ModflowGwfic(gwf, strt=5.0) + npf = ModflowGwfnpf(gwf, k=1.0, save_specific_discharge=True) + chd_spd = [[(0, 0, 0), 10.0], [(0, 4, 4), 0.0]] + chd = ModflowGwfchd(gwf, stress_period_data=chd_spd) + budget_file = f"{sim_name}.cbc" + head_file = f"{sim_name}.hds" + oc = ModflowGwfoc( + gwf, + budget_filerecord=budget_file, + head_filerecord=head_file, + saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")], + ) + + sim.write_simulation() + success, _ = sim.run_simulation(silent=True) + + 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"] + + for field in ["qx", "qy", "qz"]: + assert field in available_fields, ( + f"Expected field '{field}' not found in {available_fields}" + ) + + 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 + + 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]) diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index 5c9842b74..90fce3652 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -1604,7 +1604,7 @@ def get_data( if t ] - def get_ts(self, idx, text=None, times=None): + def get_ts(self, idx, text=None, times=None, variable="q"): """ Get a time series from the binary budget file. @@ -1619,6 +1619,11 @@ def get_ts(self, idx, text=None, times=None): 'RIVER LEAKAGE', 'STORAGE', 'FLOW RIGHT FACE', etc. times : iterable of floats List of times to from which to get time series. + variable : str, optional + Variable name to extract from the budget record. Default is 'q'. + For records with auxiliary variables (e.g., 'DATA-SPDIS', 'DATA-SAT'), + this can be used to access auxiliary data. Examples include 'qx', + 'qy', 'qz' for specific discharge components. Returns ------- @@ -1638,6 +1643,9 @@ def get_ts(self, idx, text=None, times=None): Examples -------- + >>> # Get specific discharge x-component from DATA-SPDIS record + >>> ts = cbb.get_ts(idx=(0, 5, 5), text='DATA-SPDIS', variable='qx') + """ # issue exception if text not provided if text is None: @@ -1669,45 +1677,61 @@ def get_ts(self, idx, text=None, times=None): for idx, t in enumerate(timesint): result[idx, 0] = t + full3d = True for itim, k in enumerate(kstpkper): - try: - v = self.get_data(kstpkper=k, text=text, full3D=True) - # skip missing data - required for storage - if len(v) > 0: + if full3d: + try: + v = self.get_data(kstpkper=k, text=text, full3D=True) + + # skip missing data - required for storage + if len(v) == 0: + continue + v = v[0] istat = 1 for k, i, j in kijlist: result[itim, istat] = v[k, i, j].copy() istat += 1 - except ValueError: + continue + except ValueError: + full3d = False + if not full3d: v = self.get_data(kstpkper=k, text=text) + # skip missing data - required for storage - if len(v) > 0: - if self.modelgrid is None: - s = ( - "A modelgrid instance must be provided during " - "instantiation to get IMETH=6 timeseries data" + if len(v) == 0: + continue + + if self.modelgrid is None: + s = ( + "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 + ] + + for vv in v: + available = vv.dtype.names + if variable not in available: + raise ValueError( + f"Variable '{variable}' not found in budget record. " + f"Available variables: {list(available)}" ) - raise AssertionError(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 - ] - - for vv in v: - field = vv.dtype.names[2] - dix = np.asarray(np.isin(vv["node"], ndx)).nonzero()[0] - if len(dix) > 0: - result[itim, 1:] = vv[field][dix] + + dix = np.asarray(np.isin(vv["node"], ndx)).nonzero()[0] + if len(dix) > 0: + result[itim, 1:] = vv[variable][dix] return result