Skip to content

Commit

Permalink
Merge pull request #285 from gwmod/282-getting-specific-discharge-dat…
Browse files Browse the repository at this point in the history
…a-as-dataarrays-from-data-spdis

allow reading specific columns from budgetfile, e.g. for DATA-SPDIS
  • Loading branch information
dbrakenhoff authored Nov 17, 2023
2 parents 6a6cf75 + 6e3de76 commit d4139c8
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 8 deletions.
6 changes: 5 additions & 1 deletion nlmod/gwf/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ def get_budget_da(
gwf=None,
fname=None,
grbfile=None,
column="q",
delayed=False,
chunked=False,
**kwargs,
Expand All @@ -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
Expand All @@ -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
Expand Down
13 changes: 8 additions & 5 deletions nlmod/mfoutput/binaryfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
-------
Expand Down Expand Up @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions nlmod/mfoutput/mfoutput.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
-------
Expand All @@ -215,6 +219,7 @@ def _get_budget_da(
shape=shape,
fobj=cbcobj,
text=text,
column=column,
)

# create data array
Expand Down
10 changes: 8 additions & 2 deletions tests/test_015_gwf_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down

0 comments on commit d4139c8

Please sign in to comment.