diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index 3e64c1b..d001872 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -45,7 +45,7 @@ assert chd.data["head"][0, 0] == 1.0 assert chd.data.head.sel(per=0)[99] == 0.0 -assert np.allclose(chd.data.head[:, 1:99], np.full(98, 1e30)) +assert np.allclose(chd.data.head[:, 1:99], np.full(98, 3e30)) assert gwf.dis.data.botm.sel(lay=0, col=0, row=0) == 0.0 diff --git a/flopy4/mf6/binding.py b/flopy4/mf6/binding.py index a3c05fa..f2074c0 100644 --- a/flopy4/mf6/binding.py +++ b/flopy4/mf6/binding.py @@ -33,6 +33,8 @@ def _get_binding_type(component: Component) -> str: elif isinstance(component, Solution): return f"{component.slntype}6" else: + if len(cls_name) == 4 and (cls_name[3] == "g" or cls_name[3] == "a"): + return f"{cls_name[0:3]}6" return f"{cls_name.upper()}6" def _get_binding_terms(component: Component) -> tuple[str, ...] | None: diff --git a/flopy4/mf6/constants.py b/flopy4/mf6/constants.py index a1b3b89..e053bfe 100644 --- a/flopy4/mf6/constants.py +++ b/flopy4/mf6/constants.py @@ -2,5 +2,5 @@ MF6 = "mf6" FILL_DEFAULT = np.nan -FILL_DNODATA = 1e30 +FILL_DNODATA = 3e30 LENBOUNDNAME = 40 diff --git a/flopy4/mf6/converter/unstructure.py b/flopy4/mf6/converter/unstructure.py index 76a3f95..147bdc7 100644 --- a/flopy4/mf6/converter/unstructure.py +++ b/flopy4/mf6/converter/unstructure.py @@ -7,9 +7,11 @@ import xarray as xr import xattree from modflow_devtools.dfn.schema.block import block_sort_key +from xattree import XatSpec from flopy4.mf6.binding import Binding from flopy4.mf6.component import Component +from flopy4.mf6.constants import FILL_DNODATA from flopy4.mf6.context import Context from flopy4.mf6.spec import FileInOut @@ -144,9 +146,89 @@ def oc_setting_data(rec): return data +def _unstructure_block_param( + block_name: str, + field_name: str, + xatspec: XatSpec, + value: Component, + data: dict[str, Any], + blocks: dict, + period_data: dict, +) -> None: + # Skip child components that have been processed as bindings + if isinstance(value, Context) and field_name in xatspec.children: + child_spec = xatspec.children[field_name] + if hasattr(child_spec, "metadata") and "block" in child_spec.metadata: # type: ignore + if child_spec.metadata["block"] == block_name: # type: ignore + return + + # filter out empty values and false keywords, and convert: + # - paths to records + # - datetimes to ISO format + # - filter out false keywords + # - 'auxiliary' fields to tuples + # - xarray DataArrays with 'nper' dim to dict of kper-sliced datasets + # - other values to their original form + # TODO: use cattrs converters for field unstructuring? + match field_value := data[field_name]: + case None: + return + case bool(): + if field_value: + blocks[block_name][field_name] = field_value + case Path(): + field_spec = xatspec.attrs[field_name] + field_meta = getattr(field_spec, "metadata", {}) + t = _path_to_tuple(field_name, field_value, inout=field_meta.get("inout", "fileout")) + # name may have changed e.g dropping '_file' suffix + blocks[block_name][t[0]] = t + case datetime(): + blocks[block_name][field_name] = field_value.isoformat() + case t if ( + field_name == "auxiliary" and hasattr(field_value, "values") and field_value is not None + ): + blocks[block_name][field_name] = tuple(field_value.values.tolist()) + case xr.DataArray() if "nper" in field_value.dims: + has_spatial_dims = any( + dim in field_value.dims for dim in ["nlay", "nrow", "ncol", "nodes"] + ) + if has_spatial_dims: + field_value = _hack_structured_grid_dims( + field_value, + structured_grid_dims=value.parent.data.dims, # type: ignore + ) + if block_name == "period": + if not np.issubdtype(field_value.dtype, np.number): + dat = _hack_period_non_numeric(field_name, field_value) + for n, v in dat.items(): + period_data[n] = v + else: + period_data[field_name] = { + kper: field_value.isel(nper=kper) # type: ignore + for kper in range(field_value.sizes["nper"]) + } + else: + blocks[block_name][field_name] = field_value + + case _: + blocks[block_name][field_name] = field_value + + def unstructure_component(value: Component) -> dict[str, Any]: - from flopy4.mf6.constants import FILL_DNODATA + xatspec = xattree.get_xatspec(type(value)) + if "readarraygrid" in xatspec.attrs: + return _unstructure_grid_component(value) + elif "readasarrays" in xatspec.attrs: + return _unstructure_layer_component(value) + else: + return _unstructure_component(value) + + +def _unstructure_layer_component(value: Component) -> dict[str, Any]: + return {} + +def _unstructure_grid_component(value: Component) -> dict[str, Any]: blockspec = dict(sorted(value.dfn.blocks.items(), key=block_sort_key)) # type: ignore blocks: dict[str, dict[str, Any]] = {} xatspec = xattree.get_xatspec(type(value)) @@ -167,67 +249,53 @@ def unstructure_component(value: Component) -> dict[str, Any]: blocks[block_name] = {} for field_name in block.keys(): - # Skip child components that have been processed as bindings - if isinstance(value, Context) and field_name in xatspec.children: - child_spec = xatspec.children[field_name] - if hasattr(child_spec, "metadata") and "block" in child_spec.metadata: # type: ignore - if child_spec.metadata["block"] == block_name: # type: ignore - continue - - # filter out empty values and false keywords, and convert: - # - paths to records - # - datetimes to ISO format - # - filter out false keywords - # - 'auxiliary' fields to tuples - # - xarray DataArrays with 'nper' dim to dict of kper-sliced datasets - # - other values to their original form - # TODO: use cattrs converters for field unstructuring? - match field_value := data[field_name]: - case None: - continue - case bool(): - if field_value: - blocks[block_name][field_name] = field_value - case Path(): - field_spec = xatspec.attrs[field_name] - field_meta = getattr(field_spec, "metadata", {}) - t = _path_to_tuple( - field_name, field_value, inout=field_meta.get("inout", "fileout") - ) - # name may have changed e.g dropping '_file' suffix - blocks[block_name][t[0]] = t - case datetime(): - blocks[block_name][field_name] = field_value.isoformat() - case t if ( - field_name == "auxiliary" - and hasattr(field_value, "values") - and field_value is not None - ): - blocks[block_name][field_name] = tuple(field_value.values.tolist()) - case xr.DataArray() if "nper" in field_value.dims: - has_spatial_dims = any( - dim in field_value.dims for dim in ["nlay", "nrow", "ncol", "nodes"] - ) - if has_spatial_dims: - field_value = _hack_structured_grid_dims( - field_value, - structured_grid_dims=value.parent.data.dims, # type: ignore - ) - if block_name == "period": - if not np.issubdtype(field_value.dtype, np.number): - dat = _hack_period_non_numeric(field_name, field_value) - for n, v in dat.items(): - period_data[n] = v - else: - period_data[field_name] = { - kper: field_value.isel(nper=kper) # type: ignore - for kper in range(field_value.sizes["nper"]) - } - else: - blocks[block_name][field_name] = field_value - - case _: - blocks[block_name][field_name] = field_value + _unstructure_block_param( + block_name, field_name, xatspec, value, data, blocks, period_data + ) + + # invert key order, (arr_name, kper) -> (kper, arr_name) + for arr_name, periods in period_data.items(): + for kper, arr in periods.items(): + if kper not in period_blocks: + period_blocks[kper] = {} + period_blocks[kper][arr_name] = arr + + # setup indexed period blocks, combine arrays into datasets + for kper, block in period_blocks.items(): + key = f"period {kper + 1}" + for arr_name, val in block.items(): + if not np.all(val == FILL_DNODATA): + if key not in blocks: + blocks[key] = {} + blocks[key][arr_name] = val + + return {name: block for name, block in blocks.items() if name != "period"} + + +def _unstructure_component(value: Component) -> dict[str, Any]: + blockspec = dict(sorted(value.dfn.blocks.items(), key=block_sort_key)) # type: ignore + blocks: dict[str, dict[str, Any]] = {} + xatspec = xattree.get_xatspec(type(value)) + data = xattree.asdict(value) + + # create child component binding blocks + blocks.update(_make_binding_blocks(value)) + + # process blocks in order, unstructuring fields as needed, + # then slice period data into separate kper-indexed blocks + # each of which contains a dataset indexed for that period. + for block_name, block in blockspec.items(): + period_data = {} # type: ignore + period_blocks = {} # type: ignore + period_block_name = None + + if block_name not in blocks: + blocks[block_name] = {} + + for field_name in block.keys(): + _unstructure_block_param( + block_name, field_name, xatspec, value, data, blocks, period_data + ) # invert key order, (arr_name, kper) -> (kper, arr_name) for arr_name, periods in period_data.items(): diff --git a/flopy4/mf6/gwf/__init__.py b/flopy4/mf6/gwf/__init__.py index 8c69b9a..b020010 100644 --- a/flopy4/mf6/gwf/__init__.py +++ b/flopy4/mf6/gwf/__init__.py @@ -8,8 +8,10 @@ from xattree import xattree from flopy4.mf6.gwf.chd import Chd +from flopy4.mf6.gwf.chdg import Chdg from flopy4.mf6.gwf.dis import Dis from flopy4.mf6.gwf.drn import Drn +from flopy4.mf6.gwf.drng import Drng from flopy4.mf6.gwf.ic import Ic from flopy4.mf6.gwf.npf import Npf from flopy4.mf6.gwf.oc import Oc @@ -21,7 +23,7 @@ from flopy4.mf6.utils import open_cbc, open_hds from flopy4.utils import to_path -__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Sto", "Wel", "Rch"] +__all__ = ["Gwf", "Chd", "Chdg", "Dis", "Drn", "Drng", "Ic", "Npf", "Oc", "Rch", "Sto", "Wel"] def convert_grid(value): @@ -79,9 +81,11 @@ def budget(self): npf: Npf | None = field(block="packages", default=None) sto: Sto | None = field(block="packages", default=None) chd: list[Chd] = field(block="packages") - wel: list[Wel] = field(block="packages") + chdg: list[Chdg] = field(block="packages") drn: list[Drn] = field(block="packages") + drng: list[Drng] = field(block="packages") rch: list[Rch] = field(block="packages") + wel: list[Wel] = field(block="packages") output: Output = attrs.field( default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True) ) diff --git a/flopy4/mf6/gwf/chdg.py b/flopy4/mf6/gwf/chdg.py new file mode 100644 index 0000000..e3b5da8 --- /dev/null +++ b/flopy4/mf6/gwf/chdg.py @@ -0,0 +1,51 @@ +from pathlib import Path +from typing import ClassVar, Optional + +import numpy as np + +# from attrs import Converter +from numpy.typing import NDArray +from xattree import xattree + +# from flopy4.mf6.converter import structure_array +from flopy4.mf6.package import Package +from flopy4.mf6.spec import array, field, path +from flopy4.mf6.utils.grid import update_maxbound +from flopy4.utils import to_path + + +@xattree +class Chdg(Package): + multi_package: ClassVar[bool] = True + auxiliary: Optional[list[str]] = array(block="options", default=None) + auxmultname: Optional[str] = field(block="options", default=None) + print_input: bool = field(block="options", default=False) + print_flows: bool = field(block="options", default=False) + readarraygrid: bool = field(block="options", default=True) + save_flows: bool = field(block="options", default=False) + obs_filerecord: Optional[Path] = path( + block="options", default=None, converter=to_path, inout="fileout" + ) + export_array_netcdf: bool = field(block="options", default=False) + dev_no_newton: bool = field(default=False, block="options") + maxbound: Optional[int] = field(block="dimensions", default=None, init=False) + head: Optional[NDArray[np.float64]] = array( + block="period", + dims=( + "nper", + "nodes", + ), + default=None, + # converter=Converter(structure_array, takes_self=True, takes_field=True), + on_setattr=update_maxbound, + ) + aux: Optional[NDArray[np.float64]] = array( + block="period", + dims=( + "nper", + "nodes", + ), + default=None, + # converter=Converter(structure_array, takes_self=True, takes_field=True), + on_setattr=update_maxbound, + ) diff --git a/flopy4/mf6/gwf/drng.py b/flopy4/mf6/gwf/drng.py new file mode 100644 index 0000000..b8ea38f --- /dev/null +++ b/flopy4/mf6/gwf/drng.py @@ -0,0 +1,62 @@ +from pathlib import Path +from typing import ClassVar, Optional + +import numpy as np + +# from attrs import Converter +from numpy.typing import NDArray +from xattree import xattree + +# from flopy4.mf6.converter import structure_array +from flopy4.mf6.package import Package +from flopy4.mf6.spec import array, field, path +from flopy4.mf6.utils.grid import update_maxbound +from flopy4.utils import to_path + + +@xattree +class Drng(Package): + multi_package: ClassVar[bool] = True + auxiliary: Optional[list[str]] = array(block="options", default=None) + auxmultname: Optional[str] = field(block="options", default=None) + print_input: bool = field(block="options", default=False) + print_flows: bool = field(block="options", default=False) + readarraygrid: bool = field(block="options", default=True) + save_flows: bool = field(block="options", default=False) + obs_filerecord: Optional[Path] = path( + block="options", default=None, converter=to_path, inout="fileout" + ) + export_array_netcdf: bool = field(block="options", default=False) + mover: bool = field(block="options", default=False) + dev_cubic_scaling: bool = field(default=False, block="options") + maxbound: Optional[int] = field(block="dimensions", default=None, init=False) + elev: Optional[NDArray[np.float64]] = array( + block="period", + dims=( + "nper", + "nodes", + ), + default=None, + # converter=Converter(structure_array, takes_self=True, takes_field=True), + on_setattr=update_maxbound, + ) + cond: Optional[NDArray[np.float64]] = array( + block="period", + dims=( + "nper", + "nodes", + ), + default=None, + # converter=Converter(structure_array, takes_self=True, takes_field=True), + on_setattr=update_maxbound, + ) + aux: Optional[NDArray[np.float64]] = array( + block="period", + dims=( + "nper", + "nodes", + ), + default=None, + # converter=Converter(structure_array, takes_self=True, takes_field=True), + on_setattr=update_maxbound, + ) diff --git a/test/test_mf6_codec.py b/test/test_mf6_codec.py index 84db7e6..f0cde5c 100644 --- a/test/test_mf6_codec.py +++ b/test/test_mf6_codec.py @@ -2,9 +2,11 @@ from pprint import pprint +import numpy as np import pytest from flopy4.mf6.codec import dumps, loads +from flopy4.mf6.constants import FILL_DNODATA from flopy4.mf6.converter import COMPONENT_CONVERTER @@ -249,14 +251,57 @@ def test_dumps_chd(): assert len(lines) == 2 assert "1 1 1 10.0" in dumped # First CHD cell - node 1 assert "1 10 10 20.0" in dumped # Second CHD cell - node 100 - assert "1e+30" not in dumped - assert "1.0e+30" not in dumped + assert "3e+30" not in dumped + assert "3.0e+30" not in dumped loaded = loads(dumped) print("CHD load:") pprint(loaded) +def test_dumps_chdg(): + from flopy4.mf6.gwf import Chdg, Dis, Gwf + + nlay = 1 + nrow = 10 + ncol = 10 + + dis = Dis(nlay=nlay, nrow=nrow, ncol=ncol) + gwf = Gwf(dis=dis) + + head = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float) + head[0, 0, 0] = 1.0 + head[0, 9, 9] = 0.0 + chd = Chdg( + parent=gwf, + head=np.expand_dims(head.ravel(), axis=0), + save_flows=True, + print_input=True, + dims={"nper": 1}, + ) + + dumped = dumps(COMPONENT_CONVERTER.unstructure(chd)) + print("CHD dump:") + print(dumped) + + assert "BEGIN PERIOD 1" in dumped + assert "END PERIOD 1" in dumped + + period_section = dumped.split("BEGIN PERIOD 1")[1].split("END PERIOD 1")[0].strip() + lines = [line.strip() for line in period_section.split("\n") if line.strip()] + + assert len(lines) == 12 + assert "READARRAYGRID" in dumped + assert "MAXBOUND 2" in dumped + dump_data = [[float(x) for x in line.split()] for line in lines[2:12]] + dump_head = np.array(dump_data) + assert np.allclose(head, dump_head) + + loaded = loads(dumped) + print("CHDG load:") + pprint(loaded) + + def test_dumps_wel(): from flopy4.mf6.gwf import Dis, Gwf, Wel @@ -291,8 +336,8 @@ def test_dumps_wel(): assert "1 3 4 -100.0" in dumped # (0,2,3) -> node 24 assert "2 6 8 -50.0" in dumped # (1,5,7) -> node 158 assert "3 9 2 25.0" in dumped # (2,8,1) -> node 282 - assert "1e+30" not in dumped - assert "1.0e+30" not in dumped + assert "3e+30" not in dumped + assert "3.0e+30" not in dumped loaded = loads(dumped) print("WEL load:") @@ -356,8 +401,8 @@ def test_dumps_drn(): assert "1 2 2 12.0 1.5" in dumped # Period 2: (0,1,1) assert "1 3 4 9.0 0.8" in dumped # Period 2: (0,2,3) assert "2 4 3 7.0 2.2" in dumped # Period 2: (1,3,2) - assert "1e+30" not in dumped - assert "1.0e+30" not in dumped + assert "3e+30" not in dumped + assert "3.0e+30" not in dumped loaded = loads(dumped) print("DRN load:") @@ -369,9 +414,9 @@ def test_dumps_npf(): dis = Dis(nlay=2, nrow=5, ncol=5) gwf = Gwf(dis=dis) - drn = Npf(parent=gwf, cvoptions=Npf.CvOptions(dewatered=True), k=1.0) + npf = Npf(parent=gwf, cvoptions=Npf.CvOptions(dewatered=True), k=1.0) - dumped = dumps(COMPONENT_CONVERTER.unstructure(drn)) + dumped = dumps(COMPONENT_CONVERTER.unstructure(npf)) print("NPF dump:") print(dumped) @@ -407,8 +452,8 @@ def test_dumps_chd_2(): assert "100.0" in dumped # Left boundary assert "95.0" in dumped # Right boundary assert "98.0" in dumped # Bottom boundary - assert "1e+30" not in dumped - assert "1.0e+30" not in dumped + assert "3e+30" not in dumped + assert "3.0e+30" not in dumped loaded = loads(dumped) print("CHD load:") @@ -450,8 +495,8 @@ def test_dumps_wel_with_aux(): # node q aux_value assert "1 2 3 -75.0 1.0" in dumped # (0,1,2) -> node 8, q=-75.0, aux=1.0 assert "2 4 5 -25.0 2.0" in dumped # (1,3,4) -> node 45, q=-25.0, aux=2.0 - assert "1e+30" not in dumped - assert "1.0e+30" not in dumped + assert "3e+30" not in dumped + assert "3.0e+30" not in dumped loaded = loads(dumped) print("WEL+aux load:") diff --git a/test/test_mf6_component.py b/test/test_mf6_component.py index 03cd9ff..28c62e5 100644 --- a/test/test_mf6_component.py +++ b/test/test_mf6_component.py @@ -11,7 +11,7 @@ from flopy4.mf6.component import COMPONENTS from flopy4.mf6.constants import FILL_DNODATA, LENBOUNDNAME -from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc +from flopy4.mf6.gwf import Chd, Chdg, Dis, Gwf, Ic, Npf, Oc from flopy4.mf6.ims import Ims from flopy4.mf6.simulation import Simulation from flopy4.mf6.tdis import Tdis @@ -398,6 +398,57 @@ def test_quickstart(function_tmpdir): sim.run() +def test_quickstart_grid(function_tmpdir): + sim_name = "quickstart" + gwf_name = "mymodel" + + # dimensions + nlay = 1 + nrow = 10 + ncol = 10 + nstp = 1 + + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) + # time = Time(perlen=[1.0, 1.0], nstp=[1, 1], tsmult=[1.0, 1.0]) + ims = Ims(models=[gwf_name]) + dis = Dis( + nlay=nlay, + nrow=nrow, + ncol=ncol, + top=1.0, + botm=0.0, + ) + sim = Simulation( + tdis=time, + workspace=function_tmpdir, + name=sim_name, + solutions={"ims": ims}, + ) + gwf = Gwf(parent=sim, dis=dis, name=gwf_name) + ic = Ic(parent=gwf) + oc = Oc( + parent=gwf, + budget_file=f"{gwf_name}.bud", + head_file=f"{gwf_name}.hds", + save_head=["all"], + save_budget=["all"], + ) + npf = Npf(parent=gwf, icelltype=0, k=1.0) + + # chd grid based input, step 1 data head array + head = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float) + head[0, 0, 0] = 1.0 + head[0, 9, 9] = 0.0 + # TODO: support dict style input keyed on SP with separate grid arrays + chd = Chdg(parent=gwf, head=np.expand_dims(head.ravel(), axis=0)) + # headnone = np.full((nlay, nrow, ncol), FILL_DNODATA, dtype=float) + # ts_head = np.stack((head.ravel(), headnone.ravel()), axis=0) + # chd = Chdg(parent=gwf, head=ts_head) + + sim.write() + sim.run() + + def test_write_ascii(function_tmpdir): sim_name = "sim" gwf_name = "gwf"