allow reading specific columns from budgetfile, e.g. for DATA-SPDIS

nlmod/gwf/output.py

@@ -130,6 +130,7 @@ def get_budget_da(
     gwf=None,
     fname=None,
     grbfile=None,
+    column="q",
     delayed=False,
     chunked=False,
     **kwargs,
@@ -150,6 +151,9 @@ def get_budget_da(
     grbfile : str
         path to file containing binary grid information, only needed if reading
         output from file using fname
+    column : str
+        name of column in rec-array to read, default is 'q' which contains the fluxes
+        for most budget datasets.
     delayed : bool, optional
         if delayed is True, do not load output data into memory, default is False.
     chunked : bool, optional
@@ -161,7 +165,7 @@ def get_budget_da(
         budget data array.
     """
     cbcobj = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile)
-    da = _get_budget_da(cbcobj, text, **kwargs)
+    da = _get_budget_da(cbcobj, text, column=column, **kwargs)
     da.attrs["units"] = "m3/d"
 
     # set time index if ds/gwt are provided

nlmod/mfoutput/binaryfile.py

@@ -48,14 +48,14 @@ def _get_binary_head_data(kstpkper, fobj):
     return arr
 
 
-def __create3D(data, fobj):
-    """Copy of CellBudgetFile.__create3D.
+def __create3D(data, fobj, column="q"):
+    """Adapted from CellBudgetFile.__create3D.
 
     See flopy.utils.binaryfile.CellBudgetFile.__create3D.
     """
     out = np.ma.zeros(fobj.nnodes, dtype=np.float32)
     out.mask = True
-    for [node, q] in zip(data["node"], data["q"]):
+    for [node, q] in zip(data["node"], data[column]):
         idx = node - 1
         out.data[idx] += q
         out.mask[idx] = False
@@ -105,7 +105,7 @@ def _select_data_indices_budget(fobj, text, kstpkper):
     return select_indices
 
 
-def _get_binary_budget_data(kstpkper, fobj, text):
+def _get_binary_budget_data(kstpkper, fobj, text, column="q"):
     """Get budget data from binary CellBudgetFile.
 
     Code copied from flopy.utils.binaryfile.CellBudgetFile and adapted to
@@ -121,6 +121,9 @@ def _get_binary_budget_data(kstpkper, fobj, text):
         file object
     text : str
         text indicating which dataset to read
+    column : str
+        name of column in rec-array to read, default is 'q' which contains the fluxes
+        for most budget datasets.
 
     Returns
     -------
@@ -211,7 +214,7 @@ def _get_binary_budget_data(kstpkper, fobj, text):
             dtype = np.dtype(dtype_list)
             nlist = binaryread(f, np.int32)[0]
             data = binaryread(f, dtype, shape=(nlist,))
-            data = __create3D(data, fobj)
+            data = __create3D(data, fobj, column=column)
             if fobj.modelgrid is not None:
                 return np.reshape(data, fobj.shape)
             else:

nlmod/mfoutput/mfoutput.py

@@ -177,6 +177,7 @@ def _get_budget_da(
     cbcobj,
     text,
     modelgrid=None,
+    column="q",
     **kwargs,
 ):
     """Get budget data array based on CellBudgetFile and text string.
@@ -193,6 +194,9 @@ def _get_budget_da(
     modelgrid : flopy.discretization.Grid, optional
         flopy modelgrid object, default is None, in which case the modelgrid
         is derived from `cbcobj.modelgrid`
+    column : str
+        name of column in rec-array to read, default is 'q' which contains the fluxes
+        for most budget datasets.
 
     Returns
     -------
@@ -215,6 +219,7 @@ def _get_budget_da(
         shape=shape,
         fobj=cbcobj,
         text=text,
+        column=column,
     )
 
     # create data array