diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py index ed4ed748..4b4e4ea2 100644 --- a/nlmod/gwf/output.py +++ b/nlmod/gwf/output.py @@ -8,17 +8,19 @@ from shapely.geometry import Point from ..dims.grid import modelgrid_from_ds +from ..dims.resample import get_affine_world_to_mod from ..mfoutput.mfoutput import ( _get_budget_da, _get_heads_da, _get_time_index, _get_flopy_data_object, + _get_grb_file, ) logger = logging.getLogger(__name__) -def get_headfile(ds=None, gwf=None, fname=None, grbfile=None): +def get_headfile(ds=None, gwf=None, fname=None, grb_file=None): """Get flopy HeadFile object. Provide one of ds, gwf or fname. @@ -31,7 +33,7 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None): groundwater flow model, by default None fname : str, optional path to heads file, by default None - grbfile : str + grb_file : str path to file containing binary grid information Returns @@ -39,14 +41,14 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None): flopy.utils.HeadFile HeadFile object handle """ - return _get_flopy_data_object("head", ds, gwf, fname, grbfile) + return _get_flopy_data_object("head", ds, gwf, fname, grb_file) def get_heads_da( ds=None, gwf=None, fname=None, - grbfile=None, + grb_file=None, delayed=False, chunked=False, **kwargs, @@ -62,7 +64,7 @@ def get_heads_da( Flopy groundwaterflow object. fname : path, optional path to a binary heads file - grbfile : str, optional + grb_file : str, optional path to file containing binary grid information, only needed if reading output from file using fname delayed : bool, optional @@ -75,7 +77,7 @@ def get_heads_da( da : xarray.DataArray heads data array. """ - hobj = get_headfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile) + hobj = get_headfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file) # gwf.output.head() defaults to a structured grid if gwf is not None and ds is None and fname is None: kwargs["modelgrid"] = gwf.modelgrid @@ -99,7 +101,7 @@ def get_heads_da( return da -def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None): +def get_cellbudgetfile(ds=None, gwf=None, fname=None, grb_file=None): """Get flopy CellBudgetFile object. Provide one of ds, gwf or fname. @@ -111,9 +113,9 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None): gwf : flopy.mf6.ModflowGwf, optional groundwater flow model, by default None fname_cbc : str, optional - path to cell budget file, by default None\ - grbfile : str, optional - path to file containing binary grid information, only needed if + path to cell budget file, by default None + grb_file : str, optional + path to file containing binary grid information, only needed if fname_cbc is passed as only argument. Returns @@ -121,7 +123,7 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None): flopy.utils.CellBudgetFile CellBudgetFile object handle """ - return _get_flopy_data_object("budget", ds, gwf, fname, grbfile) + return _get_flopy_data_object("budget", ds, gwf, fname, grb_file) def get_budget_da( @@ -129,7 +131,7 @@ def get_budget_da( ds=None, gwf=None, fname=None, - grbfile=None, + grb_file=None, column="q", delayed=False, chunked=False, @@ -148,7 +150,7 @@ def get_budget_da( fname : path, optional specify the budget file to load, if not provided budget file will be obtained from ds or gwf. - grbfile : str + grb_file : str path to file containing binary grid information, only needed if reading output from file using fname column : str @@ -164,7 +166,7 @@ def get_budget_da( da : xarray.DataArray budget data array. """ - cbcobj = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grbfile=grbfile) + cbcobj = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file) da = _get_budget_da(cbcobj, text, column=column, **kwargs) da.attrs["units"] = "m3/d" @@ -232,7 +234,157 @@ def get_gwl_from_wet_cells(head, layer="layer", botm=None): return gwl -def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True): +def get_flow_residuals(ds, gwf=None, fname=None, grb_file=None, kstpkper=None): + """ + Get the flow residuals of a MODFLOW 6 simulation. + + Parameters + ---------- + ds : xarray.Dataset + Xarray dataset with model data. + gwf : flopy ModflowGwf, optional + Flopy groundwaterflow object. One of ds or gwf must be provided. + fname : path, optional + specify the budget file to load, if not provided budget file will + be obtained from ds or gwf. + grb_file : str + The location of the grb-file. grb_file is determied from ds when None. The + default is None. + kstpkper : tuple of 2 ints, optional + The index of the timestep and the stress period to include in the result. Include + all data in the budget-file when None. The default is None. + + Returns + ------- + da : xr.DataArray + The flow residual in each cell, in m3/d. + + """ + if grb_file is None: + grb_file = _get_grb_file(ds) + grb = flopy.mf6.utils.MfGrdFile(grb_file) + cbf = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file) + dims = ds["botm"].dims + coords = ds["botm"].coords + flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper) + mask_active = np.diff(grb.ia) > 0 + flowja_index = grb.ia[:-1][mask_active] + if kstpkper is None: + # loop over all timesteps/stress-periods + residuals = [] + for iflowja in flowja: + # residuals.append(flopy.mf6.utils.get_residuals(iflowja, grb_file)) + # use our own faster method instead of a for loop: + residual = np.full(grb.shape, np.NaN) + residual.ravel()[mask_active] = iflowja.flatten()[flowja_index] + residuals.append(residual) + dims = ("time",) + dims + coords = dict(coords) | {"time": _get_time_index(cbf, ds)} + else: + # residuals = flopy.mf6.utils.get_residuals(flowja[0], grb_file) + # use our own faster method instead of a for loop: + residuals = np.full(grb.shape, np.NaN) + residuals.ravel()[mask_active] = flowja[0].flatten()[flowja_index] + da = xr.DataArray(residuals, dims=dims, coords=coords) + return da + + +def get_flow_lower_face( + ds, gwf=None, fname=None, grb_file=None, kstpkper=None, lays=None +): + """ + Get the flow over the lower face of all model cells + + The flow Lower Face (flf) used to be written to the budget file in previous versions + of MODFLOW. In MODFLOW 6 we determine these flows from the flow-ja-face-records. + + Parameters + ---------- + ds : xarray.Dataset + Xarray dataset with model data. + gwf : flopy ModflowGwf, optional + Flopy groundwaterflow object. One of ds or gwf must be provided. + fname : path, optional + specify the budget file to load, if not provided budget file will + be obtained from ds or gwf. + grb_file : str, optional + The location of the grb-file. grb_file is determied from ds when None. The + default is None. + kstpkper : tuple of 2 ints, optional + The index of the timestep and the stress period to include in the result. Include + all data in the budget-file when None. The default is None. + lays : int or list of ints, optional + The layers to include in the result. When lays is None, all layers are included. + The default is None. + + Returns + ------- + da : xr.DataArray + The flow over the lower face of each cell, in m3/d. + + """ + if grb_file is None: + grb_file = _get_grb_file(ds) + cbf = get_cellbudgetfile(ds=ds, gwf=gwf, fname=fname, grb_file=grb_file) + flowja = cbf.get_data(text="FLOW-JA-FACE", kstpkper=kstpkper) + + if ds.gridtype == "vertex": + # determine flf_index first + grb = flopy.mf6.utils.MfGrdFile(grb_file) + + if lays is None: + lays = range(grb.nlay) + if isinstance(lays, int): + lays = [lays] + shape = (len(lays), len(ds.icell2d)) + + flf_index = np.full(shape, -1) + # get these properties outside of the for loop to increase speed + grb_ia = grb.ia + grb_ja = grb.ja + for ilay, lay in enumerate(lays): + ja_start_next_layer = (lay + 1) * grb.ncpl + for icell2d in range(grb.ncpl): + node = lay * grb.ncpl + icell2d + ia = np.arange(grb_ia[node], grb_ia[node + 1]) + mask = grb_ja[ia] >= ja_start_next_layer + if mask.any(): + # assert mask.sum() == 1 + flf_index[ilay, icell2d] = int(ia[mask]) + coords = ds["botm"][lays].coords + else: + coords = ds["botm"].coords + dims = ds["botm"].dims + + if kstpkper is None: + # loop over all tiesteps/stress-periods + flfs = [] + for iflowja in flowja: + if ds.gridtype == "vertex": + flf = np.full(shape, np.NaN) + mask = flf_index >= 0 + flf[mask] = iflowja[0, 0, flf_index[mask]] + else: + _, _, flf = flopy.mf6.utils.get_structured_faceflows(iflowja, grb_file) + flfs.append(flf) + dims = ("time",) + dims + coords = dict(coords) | {"time": _get_time_index(cbf, ds)} + else: + if ds.gridtype == "vertex": + flfs = np.full(shape, np.NaN) + mask = flf_index >= 0 + flfs[mask] = flowja[0][0, 0, flf_index[mask]] + else: + _, _, flfs = flopy.mf6.utils.get_structured_faceflows(flowja[0], grb_file) + da = xr.DataArray(flfs, dims=dims, coords=coords) + if ds.gridtype != "vertex" and lays is not None: + da = da.isel(layer=lays) + return da + + +def get_head_at_point( + head, x, y, ds=None, gi=None, drop_nan_layers=True, rotated=False +): """Get the head at a certain point from a head DataArray for all cells. Parameters @@ -253,12 +405,20 @@ def get_head_at_point(head, x, y, ds=None, gi=None, drop_nan_layers=True): None. drop_nan_layers : bool, optional Drop layers that are NaN at all timesteps. The default is True. + rotated : bool, optional + If the model grid has a rotation, and rotated is False, x and y are in model + coordinates. Otherwise x and y are in real world coordinates. The defaults is + False. Returns ------- head_point : xarray.DataArray A DataArray with dimensions (time, layer). """ + if rotated and "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0: + # calculate model coordinates from the specified real-world coordinates + x, y = get_affine_world_to_mod(ds) * (x, y) + if "icell2d" in head.dims: if gi is None: if ds is None: diff --git a/nlmod/mfoutput/binaryfile.py b/nlmod/mfoutput/binaryfile.py index a3197f0c..9d228a1b 100644 --- a/nlmod/mfoutput/binaryfile.py +++ b/nlmod/mfoutput/binaryfile.py @@ -138,13 +138,24 @@ def _get_binary_budget_data(kstpkper, fobj, text, column="q"): idx = np.array([idx]) header = fobj.recordarray[idx] - ipos = fobj.iposarray[idx].item() - imeth = header["imeth"][0] t = header["text"][0] if isinstance(t, bytes): t = t.decode("utf-8") + data = [] + for ipos in fobj.iposarray[idx]: + data.append(_get_binary_budget_record(fobj, ipos, header, column)) + + if len(data) == 1: + return data[0] + else: + return np.ma.sum(data, axis=0) + + +def _get_binary_budget_record(fobj, ipos, header, column): + """Get a single data record from the budget file.""" + imeth = header["imeth"][0] nlay = abs(header["nlay"][0]) nrow = header["nrow"][0] ncol = header["ncol"][0] diff --git a/nlmod/mfoutput/mfoutput.py b/nlmod/mfoutput/mfoutput.py index 64158e3d..d26fc73c 100644 --- a/nlmod/mfoutput/mfoutput.py +++ b/nlmod/mfoutput/mfoutput.py @@ -239,7 +239,7 @@ def _get_budget_da( return da -def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None): +def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grb_file=None): """Get modflow HeadFile or CellBudgetFile object, containg heads, budgets or concentrations @@ -255,7 +255,7 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None): groundwater flow or transport model, by default None fname : str, optional path to Head- or CellBudgetFile, by default None - grbfile : str, optional + grb_file : str, optional path to file containing binary grid information, if None modelgrid information is obtained from ds. By default None @@ -283,14 +283,11 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None): # return gwf.output.head(), gwf.output.budget() or gwt.output.concentration() return getattr(gwml.output, var)() fname = os.path.join(ds.model_ws, ds.model_name + extension) - if grbfile is None and ds is not None: + if grb_file is None and ds is not None: # get grb file - if ds.gridtype == "vertex": - grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb") - elif ds.gridtype == "structured": - grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb") - if grbfile is not None and os.path.exists(grbfile): - modelgrid = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid + grb_file = _get_grb_file(ds) + if grb_file is not None and os.path.exists(grb_file): + modelgrid = flopy.mf6.utils.MfGrdFile(grb_file).modelgrid elif ds is not None: modelgrid = modelgrid_from_ds(ds) else: @@ -307,3 +304,11 @@ def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None): logger.warning(msg) warnings.warn(msg) return flopy.utils.HeadFile(fname, text=var, modelgrid=modelgrid) + + +def _get_grb_file(ds): + if ds.gridtype == "vertex": + grb_file = os.path.join(ds.model_ws, ds.model_name + ".disv.grb") + elif ds.gridtype == "structured": + grb_file = os.path.join(ds.model_ws, ds.model_name + ".dis.grb") + return grb_file diff --git a/tests/test_015_gwf_output.py b/tests/test_015_gwf_output.py index 074e920b..fc3608ee 100644 --- a/tests/test_015_gwf_output.py +++ b/tests/test_015_gwf_output.py @@ -67,17 +67,17 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): assert np.array_equal(da.values, heads_correct, equal_nan=True) fname_hds = os.path.join(ds.model_ws, ds.model_name + ".hds") - grbfile = os.path.join(ds.model_ws, ds.model_name + ".dis.grb") - da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grbfile=grbfile) # fname + grb_file = os.path.join(ds.model_ws, ds.model_name + ".dis.grb") + da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grb_file=grb_file) # fname assert np.array_equal(da.values, heads_correct, equal_nan=True) # budget da = get_budget_da("CHD", ds=ds, gwf=None, fname=None) # ds 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("CHD", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file) # fname get_budget_da( - "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile + "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file ) # fname # unstructured @@ -127,18 +127,18 @@ def test_create_small_model_grid_only(tmpdir, model_name="test"): assert np.array_equal(da.values, heads_correct, equal_nan=True) fname_hds = os.path.join(ds.model_ws, ds.model_name + ".hds") - grbfile = os.path.join(ds.model_ws, ds.model_name + ".disv.grb") - da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grbfile=grbfile) # fname + grb_file = os.path.join(ds.model_ws, ds.model_name + ".disv.grb") + da = get_heads_da(ds=None, gwf=None, fname=fname_hds, grb_file=grb_file) # fname assert np.array_equal(da.values, heads_correct, equal_nan=True) # budget da = get_budget_da("CHD", ds=ds_unstr, gwf=None, fname=None) # ds da = get_budget_da("CHD", ds=None, gwf=gwf_unstr, fname=None) # gwf da = get_budget_da( - "CHD", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile + "CHD", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file ) # fname _ = get_budget_da( - "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grbfile=grbfile + "DATA-SPDIS", column="qz", ds=None, gwf=None, fname=fname_cbc, grb_file=grb_file ) # fname @@ -150,8 +150,8 @@ def test_get_heads_da_from_file_structured_no_grb(): def test_get_heads_da_from_file_structured_with_grb(): fname_hds = "./tests/data/mf6output/structured/test.hds" - grbfile = "./tests/data/mf6output/structured/test.dis.grb" - nlmod.gwf.output.get_heads_da(fname=fname_hds, grbfile=grbfile) + grb_file = "./tests/data/mf6output/structured/test.dis.grb" + nlmod.gwf.output.get_heads_da(fname=fname_hds, grb_file=grb_file) def test_get_budget_da_from_file_structured_no_grb(): @@ -162,8 +162,8 @@ def test_get_budget_da_from_file_structured_no_grb(): def test_get_budget_da_from_file_structured_with_grb(): fname_cbc = "./tests/data/mf6output/structured/test.cbc" - grbfile = "./tests/data/mf6output/structured/test.dis.grb" - nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grbfile=grbfile) + grb_file = "./tests/data/mf6output/structured/test.dis.grb" + nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grb_file=grb_file) def test_get_heads_da_from_file_vertex_no_grb(): @@ -174,8 +174,8 @@ def test_get_heads_da_from_file_vertex_no_grb(): def test_get_heads_da_from_file_vertex_with_grb(): fname_hds = "./tests/data/mf6output/vertex/test.hds" - grbfile = "./tests/data/mf6output/vertex/test.disv.grb" - nlmod.gwf.output.get_heads_da(fname=fname_hds, grbfile=grbfile) + grb_file = "./tests/data/mf6output/vertex/test.disv.grb" + nlmod.gwf.output.get_heads_da(fname=fname_hds, grb_file=grb_file) def test_get_budget_da_from_file_vertex_no_grb(): @@ -186,8 +186,8 @@ def test_get_budget_da_from_file_vertex_no_grb(): def test_get_budget_da_from_file_vertex_with_grb(): fname_cbc = "./tests/data/mf6output/vertex/test.cbc" - grbfile = "./tests/data/mf6output/vertex/test.disv.grb" - nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grbfile=grbfile) + grb_file = "./tests/data/mf6output/vertex/test.disv.grb" + nlmod.gwf.output.get_budget_da("CHD", fname=fname_cbc, grb_file=grb_file) def test_postprocess_head(): @@ -199,3 +199,19 @@ def test_postprocess_head(): nlmod.gwf.get_gwl_from_wet_cells(head, botm=ds["botm"]) nlmod.gwf.get_head_at_point(head, float(ds.x.mean()), float(ds.y.mean()), ds=ds) + + +def test_get_flow_residuals(): + ds = test_001_model.get_ds_from_cache("basic_sea_model") + da = nlmod.gwf.output.get_flow_residuals(ds) + assert "time" in da.dims + da = nlmod.gwf.output.get_flow_residuals(ds, kstpkper=(0, 0)) + assert "time" not in da.dims + + +def test_get_flow_lower_face(): + ds = test_001_model.get_ds_from_cache("basic_sea_model") + da = nlmod.gwf.output.get_flow_lower_face(ds) + assert "time" in da.dims + da = nlmod.gwf.output.get_flow_lower_face(ds, kstpkper=(0, 0)) + assert "time" not in da.dims