From 1fa44d59dec885a5c695a13be1150122675ff92d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ruben=20Calj=C3=A9?= Date: Mon, 30 Oct 2023 17:09:49 +0100 Subject: [PATCH] replace three methods by _get_flopy_data_object --- nlmod/gwf/output.py | 73 +++++++------------------------------ nlmod/gwt/output.py | 47 +++++++++++------------- nlmod/mfoutput/mfoutput.py | 74 +++++++++++++++++++++++++++++++++++++- 3 files changed, 107 insertions(+), 87 deletions(-) diff --git a/nlmod/gwf/output.py b/nlmod/gwf/output.py index 5b271b03..3e73d86b 100644 --- a/nlmod/gwf/output.py +++ b/nlmod/gwf/output.py @@ -1,5 +1,4 @@ import logging -import os import warnings import flopy @@ -9,16 +8,20 @@ from shapely.geometry import Point from ..dims.grid import modelgrid_from_ds -from ..mfoutput.mfoutput import _get_budget_da, _get_heads_da, _get_time_index +from ..mfoutput.mfoutput import ( + _get_budget_da, + _get_heads_da, + _get_time_index, + _get_flopy_data_object, +) logger = logging.getLogger(__name__) def get_headfile(ds=None, gwf=None, fname=None, grbfile=None): - """Get modflow HeadFile object. + """Get flopy HeadFile object. - Provide one of ds, gwf or fname_hds. Not that it really matters but if - all are provided hierarchy is as follows: fname_hds > ds > gwf + Provide one of ds, gwf or fname. Parameters ---------- @@ -36,31 +39,7 @@ def get_headfile(ds=None, gwf=None, fname=None, grbfile=None): headobj : flopy.utils.HeadFile HeadFile object handle """ - msg = "Load the heads using either the ds, gwf or fname_hds" - assert ((ds is not None) + (gwf is not None) + (fname is not None)) >= 1, msg - - if fname is None: - if ds is None: - return gwf.output.head() - else: - fname = os.path.join(ds.model_ws, ds.model_name + ".hds") - # 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") - else: - grbfile = None - - if fname is not None: - if grbfile is not None: - mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid - else: - logger.warning(msg) - warnings.warn(msg) - mg = None - headobj = flopy.utils.HeadFile(fname, modelgrid=mg) - return headobj + return _get_flopy_data_object("head", ds, gwf, fname, grbfile) def get_heads_da( @@ -121,10 +100,9 @@ def get_heads_da( def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None): - """Get modflow CellBudgetFile object. + """Get flopy CellBudgetFile object. - Provide one of ds, gwf or fname_cbc. Not that it really matters but if - all are provided hierarchy is as follows: fname_cbc > ds > gwf + Provide one of ds, gwf or fname. Parameters ---------- @@ -140,35 +118,10 @@ def get_cellbudgetfile(ds=None, gwf=None, fname=None, grbfile=None): Returns ------- - cbc : flopy.utils.CellBudgetFile + cbc : flopy.utils.HeadFile, flopy.utils.CellBudgetFile or CellBudgetFile object """ - msg = "Load the budgets using either the ds or the gwf" - assert ((ds is not None) + (gwf is not None) + (fname is not None)) == 1, msg - - if fname is None: - if ds is None: - return gwf.output.budget() - else: - fname = os.path.join(ds.model_ws, ds.model_name + ".cbc") - # 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") - else: - grbfile = None - if fname is not None: - if grbfile is not None: - mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid - else: - logger.error("Cannot create budget data-array without grid information.") - raise ValueError( - "Please provide grid information by passing path to the " - "binary grid file with `grbfile=`." - ) - cbc = flopy.utils.CellBudgetFile(fname, modelgrid=mg) - return cbc + return _get_flopy_data_object("budget", ds, gwf, fname, grbfile) def get_budget_da( diff --git a/nlmod/gwt/output.py b/nlmod/gwt/output.py index 87cf580d..96c5f72b 100644 --- a/nlmod/gwt/output.py +++ b/nlmod/gwt/output.py @@ -1,41 +1,36 @@ import logging -import os -import warnings -import flopy import numpy as np import xarray as xr from ..dims.layers import calculate_thickness -from ..mfoutput.mfoutput import _get_heads_da, _get_time_index +from ..mfoutput.mfoutput import _get_heads_da, _get_time_index, _get_flopy_data_object logger = logging.getLogger(__name__) def get_concentration_obj(ds=None, gwt=None, fname=None, grbfile=None): - msg = "Load the concentration using either the ds, gwt or a fname_conc" - assert ((ds is not None) + (gwt is not None) + (fname is not None)) == 1, msg + """Get flopy HeadFile object connected to the file with the concetration of cells. - if fname is None: - if ds is None: - return gwt.output.concentration() - else: - fname = os.path.join(ds.model_ws, f"{ds.model_name}_gwt.ucn") - # 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") - else: - grbfile = None - if fname is not None: - if grbfile is not None: - mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid - else: - logger.warning(msg) - warnings.warn(msg) - mg = None - concobj = flopy.utils.HeadFile(fname, text="concentration", modelgrid=mg) + Provide one of ds, gwf or fname. + + Parameters + ---------- + ds : xarray.Dataset, optional + model dataset, by default None + gwf : flopy.mf6.ModflowGwf, optional + groundwater flow model, by default None + fname : str, optional + path to heads file, by default None + grbfile : str + path to file containing binary grid information + + Returns + ------- + headobj : flopy.utils.HeadFile + HeadFile object handle + """ + concobj = _get_flopy_data_object("concentration", ds, gwt, fname, grbfile) return concobj diff --git a/nlmod/mfoutput/mfoutput.py b/nlmod/mfoutput/mfoutput.py index 69bc5164..7a7676fe 100644 --- a/nlmod/mfoutput/mfoutput.py +++ b/nlmod/mfoutput/mfoutput.py @@ -1,9 +1,12 @@ +import os import logging import dask import xarray as xr -from ..dims.grid import get_dims_coords_from_modelgrid +import flopy + +from ..dims.grid import get_dims_coords_from_modelgrid, modelgrid_from_ds from ..dims.resample import get_affine_mod_to_world from ..dims.time import ds_time_idx from .binaryfile import _get_binary_budget_data, _get_binary_head_data @@ -217,3 +220,72 @@ def _get_budget_da( da = _create_da(stacked_arr, modelgrid, cbcobj.get_times()) return da + + +def _get_flopy_data_object(var, ds=None, gwml=None, fname=None, grbfile=None): + """Get modflow HeadFile or CellBudgetFile object, containg heads, budgets or + concentrations + + Provide one of ds, gwf or fname. + + Parameters + ---------- + var : str + The name of the variable. Can be 'head', 'budget' or 'concentration'. + ds : xarray.Dataset, optional + model dataset, by default 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 + fname_cbc is passed as only argument. + + Returns + ------- + flopy.utils.HeadFile or flopy.utils.CellBudgetFile + """ + if var == "head": + ml_name = "gwf" + extension = ".hds" + flopy_class = flopy.utils.HeadFile + elif var == "budget": + ml_name = "gwf" + extension = ".cbc" + flopy_class = flopy.utils.CellBudgetFile + elif var == "concentration": + ml_name = "gwt" + extension = "_gwt.ucn" + flopy_class = flopy.utils.HeadFile + else: + raise (ValueError(f"Unknown variable {var}")) + msg = f"Load the {var}s using either ds, {ml_name} or fname" + assert ((ds is not None) + (gwml is not None) + (fname is not None)) == 1, msg + if fname is None: + if ds is 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: + # 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") + else: + grbfile = None + if grbfile is not None and os.path.exists(grbfile): + mg = flopy.mf6.utils.MfGrdFile(grbfile).modelgrid + elif ds is not None: + mg = modelgrid_from_ds(ds) + else: + logger.error(f"Cannot create {var} data-array without grid information.") + raise ValueError( + "Please provide grid information by passing path to the " + "binary grid file with `grbfile=`." + ) + if isinstance(flopy_class, flopy.utils.HeadFile): + return flopy_class(fname, modelgrid=mg, text=var) + else: + return flopy_class(fname, modelgrid=mg)