diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py index 7024c333..ed4ed748 100644 --- a/nlmod/gwf/output.py +++ b/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 diff --git a/nlmod/mfoutput/binaryfile.py b/nlmod/mfoutput/binaryfile.py index b10ae78a..a3197f0c 100644 --- a/nlmod/mfoutput/binaryfile.py +++ b/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: diff --git a/nlmod/mfoutput/mfoutput.py b/nlmod/mfoutput/mfoutput.py index 4f0ef27e..8d889941 100644 --- a/nlmod/mfoutput/mfoutput.py +++ b/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 diff --git a/tests/test_015_gwf_output.py b/tests/test_015_gwf_output.py index 7fa41eaf..8083bee5 100644 --- a/tests/test_015_gwf_output.py +++ b/tests/test_015_gwf_output.py @@ -46,7 +46,7 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): nlmod.gwf.dis(ds, gwf) # create node property flow - nlmod.gwf.npf(ds, gwf, save_flows=True) + nlmod.gwf.npf(ds, gwf, save_flows=True, save_specific_discharge=True) # Create the initial conditions package nlmod.gwf.ic(ds, gwf, starting_head=1.0) @@ -76,6 +76,9 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): da = get_budget_da("CHD", ds=None, gwf=gwf, fname=None) # gwf fname_cbc = os.path.join(ds.model_ws, ds.model_name + ".cbc") get_budget_da("CHD", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile) # fname + get_budget_da( + "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile + ) # fname # unstructured ds_unstr = refine( @@ -103,7 +106,7 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): nlmod.gwf.dis(ds_unstr, gwf_unstr) # create node property flow - nlmod.gwf.npf(ds_unstr, gwf_unstr) + nlmod.gwf.npf(ds_unstr, gwf_unstr, save_flows=True, save_specific_discharge=True) # Create the initial conditions package nlmod.gwf.ic(ds_unstr, gwf_unstr, starting_head=1.0) @@ -134,6 +137,9 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): da = get_budget_da( "CHD", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile ) # fname + _ = get_budget_da( + "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile + ) # fname def test_get_heads_da_from_file_structured_no_grb():