From 9909a601a39b685d10bc2fe88798475422537ead Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Mon, 29 Jun 2026 09:14:40 -0700 Subject: [PATCH 1/2] fix(cellbudgetfile): support aux vars with full3D in get_data() --- autotest/test_cellbudgetfile.py | 79 ++++++++++++++++++++++++++++++ flopy/utils/binaryfile/__init__.py | 44 ++++++++++++----- 2 files changed, 111 insertions(+), 12 deletions(-) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index 40a35308d..fc3e9d90a 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -1158,6 +1158,85 @@ def test_cellbudgetfile_get_ts_backwards_compatible_idx_format( ) +@pytest.mark.requires_exe("mf6") +def test_cellbudgetfile_full3D_aux_var(function_tmpdir): + """ + Reproduce GitHub issue #2774: get_data(full3D=True) should return + correct values for auxiliary variables (e.g. sat), not just for q. + """ + from flopy.mf6 import ( + MFSimulation, + ModflowGwf, + ModflowGwfchd, + ModflowGwfdis, + ModflowGwfic, + ModflowGwfnpf, + ModflowGwfoc, + ModflowIms, + ModflowTdis, + ) + + sim_name = "test_full3d_sat" + sim = MFSimulation(sim_name=sim_name, sim_ws=function_tmpdir, exe_name="mf6") + ModflowTdis(sim, nper=2, perioddata=[(1.0, 1, 1.0), (1.0, 1, 1.0)]) + ModflowIms(sim) + gwf = ModflowGwf(sim, modelname=sim_name, save_flows=True) + nlay, nrow, ncol = 1, 5, 5 + ModflowGwfdis( + gwf, nlay=nlay, nrow=nrow, ncol=ncol, delr=10.0, delc=10.0, top=10.0, botm=[0.0] + ) + ModflowGwfic(gwf, strt=5.0) + ModflowGwfnpf(gwf, k=1.0, icelltype=1, save_saturation=True) + ModflowGwfchd(gwf, stress_period_data=[[(0, 0, 0), 9.0], [(0, 4, 4), 1.0]]) + ModflowGwfoc( + gwf, + budget_filerecord=f"{sim_name}.cbc", + head_filerecord=f"{sim_name}.hds", + saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")], + ) + + sim.write_simulation() + success, _ = sim.run_simulation(silent=True) + assert success + + cbc = gwf.output.budget() + + # non-full3D + sat_rec = cbc.get_data(text="DATA-SAT") + assert len(sat_rec) > 0 + assert "sat" in sat_rec[0].dtype.names + assert "q" in sat_rec[0].dtype.names + assert np.allclose(sat_rec[0]["q"], 0.0), "q should be zero in DATA-SAT records" + assert not np.allclose(sat_rec[0]["sat"], 0.0), "sat values should be non-zero" + assert not np.allclose(sat_rec[0]["sat"], sat_rec[0]["sat"][0]), ( + "sat values should vary across cells" + ) + + # full3D with default variable="q" + q_3d = cbc.get_data(text="DATA-SAT", full3D=True) + assert q_3d[0].shape == (nlay, nrow, ncol) + assert np.allclose(np.ma.filled(q_3d[0], 0.0), 0.0) + + # full3D with variable="sat" + sat_3d = cbc.get_data(text="DATA-SAT", full3D=True, variable="sat") + assert sat_3d[0].shape == (nlay, nrow, ncol) + assert not np.allclose(np.ma.filled(sat_3d[0], 0.0), 0.0) + + # check recarray and full 3D array match + rec = sat_rec[0] + arr = sat_3d[0] + for node, sat_val in zip(rec["node"], rec["sat"]): + k = (node - 1) // (nrow * ncol) + rc = (node - 1) % (nrow * ncol) + r = rc // ncol + c = rc % ncol + np.testing.assert_allclose( + arr[k, r, c], + sat_val, + err_msg=f"3D SAT value mismatch at node {node}", + ) + + @pytest.mark.requires_exe("mf6") def test_cellbudgetfile_write_preserves_aux_vars(dis_sim, function_tmpdir): """Test that write() method preserves auxiliary variables in imeth=6 records.""" diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index c7838021c..f8be84b65 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -2894,6 +2894,7 @@ def get_data( paknam=None, paknam2=None, full3D=False, + variable="q", ) -> Union[list, np.ndarray]: """ Get data from the binary budget file. @@ -2922,6 +2923,10 @@ def get_data( If true, then return the record as a three dimensional numpy array, even for those list-style records written as part of a 'COMPACT BUDGET' MODFLOW budget file. (Default is False.) + variable : str + The variable name to extract when full3D is True and the record + contains auxiliary variables (e.g. 'sat' for 'DATA-SAT' records, + or 'qx'/'qy'/'qz' for 'DATA-SPDIS' records). Default is 'q'. Returns ------- @@ -3000,7 +3005,7 @@ def get_data( "'idx', or 'text'" ) return [ - self.get_record(idx, full3D=full3D) + self.get_record(idx, full3D=full3D, variable=variable) for idx, t in enumerate(select_indices) if t ] @@ -3298,7 +3303,7 @@ def _cellid_to_node(self, cellids) -> list[int]: nodes_0based = self.modelgrid.get_node(cellids) return (np.array(nodes_0based) + 1).tolist() - def get_record(self, idx, full3D=False): + def get_record(self, idx, full3D=False, variable="q"): """ Get a single data record from the budget file. @@ -3310,6 +3315,10 @@ def get_record(self, idx, full3D=False): If true, then return the record as a three dimensional numpy array, even for those list-style records written as part of a 'COMPACT BUDGET' MODFLOW budget file. (Default is False.) + variable : str + The variable name to extract when full3D is True and the record + contains auxiliary variables (e.g. 'sat' for 'DATA-SAT' records, + or 'qx'/'qy'/'qz' for 'DATA-SPDIS' records). Default is 'q'. Returns ------- @@ -3425,7 +3434,7 @@ def get_record(self, idx, full3D=False): if self.verbose: s += f"a list array of shape ({nlay}, {nrow}, {ncol})" print(s) - return self.__create3D(data) + return self.__create3D(data, variable=variable) else: if self.verbose: s += f"a numpy recarray of size ({nlist}, {2 + naux})" @@ -3451,7 +3460,7 @@ def get_record(self, idx, full3D=False): s += f"a numpy recarray of size ({nlist}, 2)" print(s) if full3D: - data = self.__create3D(data) + data = self.__create3D(data, variable=variable) if self.modelgrid is not None: return np.reshape(data, self.shape) else: @@ -3464,28 +3473,39 @@ def get_record(self, idx, full3D=False): # should not reach this point return - def __create3D(self, data): + def __create3D(self, data, variable="q"): """ - Convert a dictionary of {node: q, ...} into a numpy masked array. + Convert list budget data into a numpy masked array. Used to create full grid arrays when the full3D keyword is set to True in get_data. Parameters ---------- - data : dictionary - Dictionary with node keywords and flows (q) items. + data : numpy recarray + Record array with at least 'node' and the specified variable field. + variable : str + The field name to map into the 3D array. Default is 'q'. Returns ------- out : numpy masked array - List contains unique simulation times (totim) in binary file. + Masked array of shape self.shape with values from the specified + variable mapped to their grid positions. """ - out = np.ma.zeros(self.nnodes, dtype=data["q"].dtype) + if variable not in data.dtype.names: + raise ValueError( + f"variable '{variable}' not found in record. " + f"Available variables: {list(data.dtype.names)}" + ) + out = np.ma.zeros(self.nnodes, dtype=data[variable].dtype) out.mask = True - for [node, q] in zip(data["node"], data["q"]): + for node, val in zip(data["node"], data[variable]): idx = node - 1 - out.data[idx] += q + if variable == "q": + out.data[idx] += val + else: + out.data[idx] = val out.mask[idx] = False return np.ma.reshape(out, self.shape) From e5414d15ca983cfa290b106efbd61741b8c6cb6e Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Mon, 29 Jun 2026 09:20:51 -0700 Subject: [PATCH 2/2] only allow variable param with full3D --- autotest/test_cellbudgetfile.py | 9 +++++++++ flopy/utils/binaryfile/__init__.py | 6 ++++++ 2 files changed, 15 insertions(+) diff --git a/autotest/test_cellbudgetfile.py b/autotest/test_cellbudgetfile.py index fc3e9d90a..f99382be2 100644 --- a/autotest/test_cellbudgetfile.py +++ b/autotest/test_cellbudgetfile.py @@ -1237,6 +1237,15 @@ def test_cellbudgetfile_full3D_aux_var(function_tmpdir): ) +def test_cellbudgetfile_get_data_variable_without_full3D(example_data_path): + """variable= is only meaningful with full3D=True; without it, raise ValueError.""" + mf2005_model_path = example_data_path / "mf2005_test" + cbc = CellBudgetFile(mf2005_model_path / "test1tr.gitcbc") + with pytest.raises(ValueError, match="only used when full3D=True"): + cbc.get_data(text="WELLS", variable="IFACE") + cbc.close() + + @pytest.mark.requires_exe("mf6") def test_cellbudgetfile_write_preserves_aux_vars(dis_sim, function_tmpdir): """Test that write() method preserves auxiliary variables in imeth=6 records.""" diff --git a/flopy/utils/binaryfile/__init__.py b/flopy/utils/binaryfile/__init__.py index f8be84b65..fc349eeb5 100644 --- a/flopy/utils/binaryfile/__init__.py +++ b/flopy/utils/binaryfile/__init__.py @@ -3004,6 +3004,12 @@ def get_data( "get_data() missing 1 required argument: 'kstpkper', 'totim', " "'idx', or 'text'" ) + if variable != "q" and not full3D: + raise ValueError( + "'variable' is only used when full3D=True. " + "To access a specific field without full3D, index the " + "returned recarray directly, e.g. get_data(...)[0]['sat']." + ) return [ self.get_record(idx, full3D=full3D, variable=variable) for idx, t in enumerate(select_indices)