feat(cellbudgetfile): support aux vars with full3D in get_data() (#2776)

Add a 'variable' parameter (default "q") to get_data() and get_record() to allow specifying the field to extract to a 3D grid-shaped array. Previously "q" was returned unconditionally.

'variable' is only supported when full3D=True and the method's default behavior is unchanged.

sat = cbc.get_data(text="DATA-SAT", full3D=True, variable="sat")[0]
This is consistent with the existing `variable` parameter in get_ts(), added in #2568

Addresses #2774.

Author: wpbonelli

autotest/test_cellbudgetfile.py

@@ -1158,6 +1158,94 @@ 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}",
+        )
+
+
+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."""

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
         -------
@@ -2999,8 +3004,14 @@ 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)
+            self.get_record(idx, full3D=full3D, variable=variable)
             for idx, t in enumerate(select_indices)
             if t
         ]
@@ -3298,7 +3309,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 +3321,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 +3440,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 +3466,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 +3479,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)