From 089d1970b378f5e657585c7cc65e51d01c058295 Mon Sep 17 00:00:00 2001 From: w-bonelli Date: Thu, 23 Oct 2025 21:20:05 -0400 Subject: [PATCH] draft object dtype perioddata for oc --- flopy4/mf6/codec/writer/filters.py | 12 +- flopy4/mf6/converter/__init__.py | 51 ++++++++ flopy4/mf6/converter/structure.py | 83 +++++++++++++ .../unstructure.py} | 112 +----------------- flopy4/mf6/gwf/chd.py | 8 +- flopy4/mf6/gwf/dis.py | 12 +- flopy4/mf6/gwf/drn.py | 10 +- flopy4/mf6/gwf/ic.py | 4 +- flopy4/mf6/gwf/npf.py | 18 +-- flopy4/mf6/gwf/oc.py | 63 ++++++---- flopy4/mf6/gwf/rch.py | 8 +- flopy4/mf6/gwf/sto.py | 12 +- flopy4/mf6/gwf/wel.py | 8 +- flopy4/mf6/tdis.py | 8 +- test/test_codec.py | 15 ++- 15 files changed, 249 insertions(+), 175 deletions(-) create mode 100644 flopy4/mf6/converter/__init__.py create mode 100644 flopy4/mf6/converter/structure.py rename flopy4/mf6/{converter.py => converter/unstructure.py} (67%) diff --git a/flopy4/mf6/codec/writer/filters.py b/flopy4/mf6/codec/writer/filters.py index 2719379a..8e140cc9 100644 --- a/flopy4/mf6/codec/writer/filters.py +++ b/flopy4/mf6/codec/writer/filters.py @@ -2,6 +2,7 @@ from io import StringIO from typing import Any, Literal +import attrs import numpy as np import xarray as xr from modflow_devtools.dfn.schema.v2 import FieldType @@ -202,7 +203,8 @@ def dataset2list(value: xr.Dataset): return # special case OC for now. - is_oc = all( + # TODO remove after properly handling object dtype period data arrays + is_oc = any( str(v.name).startswith("save_") or str(v.name).startswith("print_") for v in value.data_vars.values() ) @@ -211,9 +213,17 @@ def dataset2list(value: xr.Dataset): if (first := next(iter(value.data_vars.values()))).ndim == 0: if is_oc: for name in value.data_vars.keys(): + if not (name.startswith("save_") or name.startswith("print_")): + # TODO: not working yet + if name == "perioddata": + val = value[name] + val = val.item() if val.shape == () else val + yield attrs.astuple(val, recurse=True) + continue val = value[name] val = val.item() if val.shape == () else val yield (*name.split("_"), val) + else: vals = [] for name in value.data_vars.keys(): diff --git a/flopy4/mf6/converter/__init__.py b/flopy4/mf6/converter/__init__.py new file mode 100644 index 00000000..7ee74a2b --- /dev/null +++ b/flopy4/mf6/converter/__init__.py @@ -0,0 +1,51 @@ +from pathlib import Path +from typing import Any + +import cattr +import xattree +from cattr import Converter +from cattrs.gen import make_hetero_tuple_unstructure_fn + +from flopy4.mf6.component import Component +from flopy4.mf6.context import Context +from flopy4.mf6.converter.structure import structure_array +from flopy4.mf6.converter.unstructure import ( + unstructure_component, +) +from flopy4.mf6.gwf.oc import Oc + +__all__ = [ + "structure", + "unstructure", + "structure_array", + "unstructure_array", + "COMPONENT_CONVERTER", +] + + +def _make_converter() -> Converter: + converter = Converter(unstruct_strat=cattr.UnstructureStrategy.AS_TUPLE) + converter.register_unstructure_hook_factory(xattree.has, lambda _: xattree.asdict) + converter.register_unstructure_hook(Component, unstructure_component) + converter.register_unstructure_hook( + Oc.PrintSaveSetting, make_hetero_tuple_unstructure_fn(Oc.PrintSaveSetting, converter) + ) + converter.register_unstructure_hook( + Oc.Steps, make_hetero_tuple_unstructure_fn(Oc.Steps, converter) + ) + return converter + + +COMPONENT_CONVERTER = _make_converter() + + +def structure(data: dict[str, Any], path: Path) -> Component: + component = COMPONENT_CONVERTER.structure(data, Component) + if isinstance(component, Context): + component.workspace = path.parent + component.filename = path.name + return component + + +def unstructure(component: Component) -> dict[str, Any]: + return COMPONENT_CONVERTER.unstructure(component) diff --git a/flopy4/mf6/converter/structure.py b/flopy4/mf6/converter/structure.py new file mode 100644 index 00000000..2a5cbc93 --- /dev/null +++ b/flopy4/mf6/converter/structure.py @@ -0,0 +1,83 @@ +from typing import Any + +import numpy as np +import sparse +from numpy.typing import NDArray +from xattree import get_xatspec + +from flopy4.adapters import get_nn +from flopy4.mf6.config import SPARSE_THRESHOLD +from flopy4.mf6.constants import FILL_DNODATA + + +def structure_array(value, self_, field) -> NDArray: + """ + Convert a sparse dictionary representation of an array to a + dense numpy array or a sparse COO array. + + TODO: generalize this not only to dictionaries but to any + form that can be converted to an array (e.g. nested list) + """ + + if not isinstance(value, dict): + # if not a dict, assume it's a numpy array + # and let xarray deal with it if it isn't + return value + + spec = get_xatspec(type(self_)).flat + field = spec[field.name] + if not field.dims: + raise ValueError(f"Field {field} missing dims") + + # resolve dims + explicit_dims = self_.__dict__.get("dims", {}) + inherited_dims = dict(self_.parent.data.dims) if self_.parent else {} + dims = inherited_dims | explicit_dims + shape = [dims.get(d, d) for d in field.dims] + unresolved = [d for d in shape if isinstance(d, str)] + if any(unresolved): + raise ValueError(f"Couldn't resolve dims: {unresolved}") + + if np.prod(shape) > SPARSE_THRESHOLD: + a: dict[tuple[Any, ...], Any] = dict() + + def set_(arr, val, *ind): + arr[tuple(ind)] = val + + def final(arr): + coords = np.array(list(map(list, zip(*arr.keys())))) + return sparse.COO( + coords, + list(arr.values()), + shape=shape, + fill_value=field.default or FILL_DNODATA, + ) + else: + a = np.full(shape, FILL_DNODATA, dtype=field.dtype) # type: ignore + + def set_(arr, val, *ind): + arr[ind] = val + + def final(arr): + arr[arr == FILL_DNODATA] = field.default or FILL_DNODATA + return arr + + if "nper" in dims: + for kper, period in value.items(): + if kper == "*": + kper = 0 + match len(shape): + case 1: + set_(a, period, kper) + case _: + for cellid, v in period.items(): + nn = get_nn(cellid, **dims) + set_(a, v, kper, nn) + if kper == "*": + break + else: + for cellid, v in value.items(): + nn = get_nn(cellid, **dims) + set_(a, v, nn) + + return final(a) diff --git a/flopy4/mf6/converter.py b/flopy4/mf6/converter/unstructure.py similarity index 67% rename from flopy4/mf6/converter.py rename to flopy4/mf6/converter/unstructure.py index ee16db5f..9b6e1a4c 100644 --- a/flopy4/mf6/converter.py +++ b/flopy4/mf6/converter/unstructure.py @@ -3,25 +3,17 @@ from pathlib import Path from typing import Any -import numpy as np -import sparse import xarray as xr import xattree -from cattrs import Converter from modflow_devtools.dfn.schema.block import block_sort_key -from numpy.typing import NDArray -from xattree import get_xatspec -from flopy4.adapters import get_nn from flopy4.mf6.binding import Binding from flopy4.mf6.component import Component -from flopy4.mf6.config import SPARSE_THRESHOLD -from flopy4.mf6.constants import FILL_DNODATA from flopy4.mf6.context import Context from flopy4.mf6.spec import FileInOut -def path_to_tuple(name: str, value: Path, inout: FileInOut) -> tuple[str, ...]: +def _path_to_tuple(name: str, value: Path, inout: FileInOut) -> tuple[str, ...]: t = [name.upper()] if name.endswith("_file"): t[0] = name.replace("_file", "").upper() @@ -31,7 +23,7 @@ def path_to_tuple(name: str, value: Path, inout: FileInOut) -> tuple[str, ...]: return tuple(t) -def make_binding_blocks(value: Component) -> dict[str, dict[str, list[tuple[str, ...]]]]: +def _make_binding_blocks(value: Component) -> dict[str, dict[str, list[tuple[str, ...]]]]: if not isinstance(value, Context): return {} @@ -104,7 +96,7 @@ def unstructure_component(value: Component) -> dict[str, Any]: data = xattree.asdict(value) # create child component binding blocks - blocks.update(make_binding_blocks(value)) + blocks.update(_make_binding_blocks(value)) # process blocks in order, unstructuring fields as needed, # then slice period data into separate kper-indexed blocks @@ -132,6 +124,7 @@ def unstructure_component(value: Component) -> dict[str, Any]: # - '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 @@ -141,7 +134,7 @@ def unstructure_component(value: Component) -> dict[str, Any]: case Path(): field_spec = xatspec.attrs[field_name] field_meta = getattr(field_spec, "metadata", {}) - t = path_to_tuple( + t = _path_to_tuple( field_name, field_value, inout=field_meta.get("inout", "fileout") ) # name may have changed e.g dropping '_file' suffix @@ -197,98 +190,3 @@ def unstructure_component(value: Component) -> dict[str, Any]: del blocks["solutiongroup"] return {name: block for name, block in blocks.items() if name != period_block_name} - - -def _make_converter() -> Converter: - converter = Converter() - converter.register_unstructure_hook_factory(xattree.has, lambda _: xattree.asdict) - converter.register_unstructure_hook(Component, unstructure_component) - return converter - - -COMPONENT_CONVERTER = _make_converter() - - -def dict_to_array(value, self_, field) -> NDArray: - """ - Convert a sparse dictionary representation of an array to a - dense numpy array or a sparse COO array. - - TODO: generalize this not only to dictionaries but to any - form that can be converted to an array (e.g. nested list) - """ - - if not isinstance(value, dict): - # if not a dict, assume it's a numpy array - # and let xarray deal with it if it isn't - return value - - spec = get_xatspec(type(self_)).flat - field = spec[field.name] - if not field.dims: - raise ValueError(f"Field {field} missing dims") - - # resolve dims - explicit_dims = self_.__dict__.get("dims", {}) - inherited_dims = dict(self_.parent.data.dims) if self_.parent else {} - dims = inherited_dims | explicit_dims - shape = [dims.get(d, d) for d in field.dims] - unresolved = [d for d in shape if isinstance(d, str)] - if any(unresolved): - raise ValueError(f"Couldn't resolve dims: {unresolved}") - - if np.prod(shape) > SPARSE_THRESHOLD: - a: dict[tuple[Any, ...], Any] = dict() - - def set_(arr, val, *ind): - arr[tuple(ind)] = val - - def final(arr): - coords = np.array(list(map(list, zip(*arr.keys())))) - return sparse.COO( - coords, - list(arr.values()), - shape=shape, - fill_value=field.default or FILL_DNODATA, - ) - else: - a = np.full(shape, FILL_DNODATA, dtype=field.dtype) # type: ignore - - def set_(arr, val, *ind): - arr[ind] = val - - def final(arr): - arr[arr == FILL_DNODATA] = field.default or FILL_DNODATA - return arr - - if "nper" in dims: - for kper, period in value.items(): - if kper == "*": - kper = 0 - match len(shape): - case 1: - set_(a, period, kper) - case _: - for cellid, v in period.items(): - nn = get_nn(cellid, **dims) - set_(a, v, kper, nn) - if kper == "*": - break - else: - for cellid, v in value.items(): - nn = get_nn(cellid, **dims) - set_(a, v, nn) - - return final(a) - - -def structure(data: dict[str, Any], path: Path) -> Component: - component = COMPONENT_CONVERTER.structure(data, Component) - if isinstance(component, Context): - component.workspace = path.parent - component.filename = path.name - return component - - -def unstructure(component: Component) -> dict[str, Any]: - return COMPONENT_CONVERTER.unstructure(component) diff --git a/flopy4/mf6/gwf/chd.py b/flopy4/mf6/gwf/chd.py index 3896b066..365163f0 100644 --- a/flopy4/mf6/gwf/chd.py +++ b/flopy4/mf6/gwf/chd.py @@ -7,7 +7,7 @@ from xattree import xattree from flopy4.mf6.constants import LENBOUNDNAME -from flopy4.mf6.converter import dict_to_array +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_utils import update_maxbound @@ -38,7 +38,7 @@ class Chd(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) aux: Optional[NDArray[np.float64]] = array( @@ -48,7 +48,7 @@ class Chd(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) boundname: Optional[NDArray[np.str_]] = array( @@ -59,6 +59,6 @@ class Chd(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) diff --git a/flopy4/mf6/gwf/dis.py b/flopy4/mf6/gwf/dis.py index 9ed93ea2..cf8c4416 100644 --- a/flopy4/mf6/gwf/dis.py +++ b/flopy4/mf6/gwf/dis.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, dim, field @@ -44,31 +44,31 @@ class Dis(Package): block="griddata", default=1.0, dims=("ncol",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) delc: NDArray[np.float64] = array( block="griddata", default=1.0, dims=("nrow",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) top: NDArray[np.float64] = array( block="griddata", default=1.0, dims=("nrow", "ncol"), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) botm: NDArray[np.float64] = array( block="griddata", default=0.0, dims=("nlay", "nrow", "ncol"), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) idomain: Optional[NDArray[np.int64]] = array( block="griddata", default=1, dims=("nlay", "nrow", "ncol"), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) nodes: int = dim( coord="node", diff --git a/flopy4/mf6/gwf/drn.py b/flopy4/mf6/gwf/drn.py index 9d89a3a4..222cd999 100644 --- a/flopy4/mf6/gwf/drn.py +++ b/flopy4/mf6/gwf/drn.py @@ -7,7 +7,7 @@ from xattree import xattree from flopy4.mf6.constants import LENBOUNDNAME -from flopy4.mf6.converter import dict_to_array +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_utils import update_maxbound @@ -37,14 +37,14 @@ class Drn(Package): block="period", dims=("nper", "nodes"), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + 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(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) aux: Optional[NDArray[np.float64]] = array( @@ -54,7 +54,7 @@ class Drn(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) boundname: Optional[NDArray[np.str_]] = array( @@ -65,6 +65,6 @@ class Drn(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) diff --git a/flopy4/mf6/gwf/ic.py b/flopy4/mf6/gwf/ic.py index e5190bab..2e3cc40c 100644 --- a/flopy4/mf6/gwf/ic.py +++ b/flopy4/mf6/gwf/ic.py @@ -3,7 +3,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field @@ -16,5 +16,5 @@ class Ic(Package): block="griddata", dims=("nodes",), default=1.0, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/npf.py b/flopy4/mf6/gwf/npf.py index 2601e8d2..a92d6981 100644 --- a/flopy4/mf6/gwf/npf.py +++ b/flopy4/mf6/gwf/npf.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path from flopy4.utils import to_path @@ -54,47 +54,47 @@ class Xt3dOptions: block="griddata", dims=("nodes",), default=0, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k: NDArray[np.float64] = array( block="griddata", dims=("nodes",), default=1.0, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k22: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k33: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle1: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle2: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle3: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) wetdry: Optional[NDArray[np.float64]] = array( block="griddata", dims=("nodes",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/oc.py b/flopy4/mf6/gwf/oc.py index 1e5b0974..432c4a8a 100644 --- a/flopy4/mf6/gwf/oc.py +++ b/flopy4/mf6/gwf/oc.py @@ -1,12 +1,12 @@ from pathlib import Path from typing import Literal, Optional +import attrs import numpy as np -from attrs import Converter, define from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path from flopy4.utils import to_path @@ -14,26 +14,38 @@ @xattree class Oc(Package): - @define(slots=False) + @attrs.define(slots=False) class Format: - fmt_kw: str = field(default="print_format") - columns: int = field(default=10) - width: int = field(default=11) - digits: int = field(default=4) + fmt_kw: str = attrs.field(default="print_format") + columns: int = attrs.field(default=10) + width: int = attrs.field(default=11) + digits: int = attrs.field(default=4) format: Literal["exponential", "fixed", "general", "scientific"] = field(default="general") - @define(slots=False) - class Steps: - all: bool = field(default=True) - first: bool | None = field(default=None) - last: bool | None = field(default=None) - steps: list[int] | None = field(default=None) - frequency: int | None = field(default=None) + @attrs.define(slots=False) + class PrintSaveSetting: + saverecord: list["Oc.SaveRecord"] = attrs.field(default=list) + printrecord: list["Oc.PrintRecord"] = attrs.field(default=list) + + @attrs.define(slots=False) + class SaveRecord: + save: Literal["save"] = attrs.field(init=False, default="save") + rtype: str = attrs.field() + steps: "Oc.Steps" = attrs.field() - @define(slots=False) - class Period: - rtype: str = field() - steps: "Oc.Steps" = field() + @attrs.define(slots=False) + class PrintRecord: + print: Literal["print"] = attrs.field(init=False, default="print") + rtype: str = attrs.field() + steps: "Oc.Steps" = attrs.field() + + @attrs.define(slots=False) + class Steps: + all: bool = attrs.field(default=True) + first: bool | None = attrs.field(default=None) + last: bool | None = attrs.field(default=None) + steps: tuple[int, ...] | None = attrs.field(default=None) + frequency: int | None = attrs.field(default=None) budget_file: Optional[Path] = path( block="options", converter=to_path, default=None, inout="fileout" @@ -51,26 +63,33 @@ class Period: block="period", default=None, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=attrs.Converter(structure_array, takes_self=True, takes_field=True), ) save_budget: Optional[NDArray[np.str_]] = array( dtype=np.dtypes.StringDType(), block="period", default=None, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=attrs.Converter(structure_array, takes_self=True, takes_field=True), ) print_head: Optional[NDArray[np.str_]] = array( dtype=np.dtypes.StringDType(), block="period", default=None, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=attrs.Converter(structure_array, takes_self=True, takes_field=True), ) print_budget: Optional[NDArray[np.str_]] = array( dtype=np.dtypes.StringDType(), block="period", default=None, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=attrs.Converter(structure_array, takes_self=True, takes_field=True), + ) + perioddata: Optional[NDArray[np.object_]] = array( + dtype=PrintSaveSetting, + block="period", + default=None, + dims=("nper",), + converter=attrs.Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/rch.py b/flopy4/mf6/gwf/rch.py index 6e62cd45..03ed0e76 100644 --- a/flopy4/mf6/gwf/rch.py +++ b/flopy4/mf6/gwf/rch.py @@ -7,7 +7,7 @@ from xattree import xattree from flopy4.mf6.constants import LENBOUNDNAME -from flopy4.mf6.converter import dict_to_array +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_utils import update_maxbound @@ -38,7 +38,7 @@ class Rch(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) aux: Optional[NDArray[np.float64]] = array( @@ -48,7 +48,7 @@ class Rch(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) boundname: Optional[NDArray[np.str_]] = array( @@ -59,6 +59,6 @@ class Rch(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) diff --git a/flopy4/mf6/gwf/sto.py b/flopy4/mf6/gwf/sto.py index 4c995e68..c9f45ef6 100644 --- a/flopy4/mf6/gwf/sto.py +++ b/flopy4/mf6/gwf/sto.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path from flopy4.utils import to_path @@ -28,29 +28,29 @@ class Sto(Package): block="griddata", dims=("nodes",), default=0, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) ss: NDArray[np.float64] = array( block="griddata", dims=("nodes",), default=1e-5, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) sy: NDArray[np.float64] = array( block="griddata", dims=("nodes",), default=0.15, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) steady_state: Optional[NDArray[np.bool_]] = array( block="period", dims=("nper",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) transient: Optional[NDArray[np.bool_]] = array( block="period", dims=("nper",), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/wel.py b/flopy4/mf6/gwf/wel.py index e04df63d..e2b97049 100644 --- a/flopy4/mf6/gwf/wel.py +++ b/flopy4/mf6/gwf/wel.py @@ -7,7 +7,7 @@ from xattree import xattree from flopy4.mf6.constants import LENBOUNDNAME -from flopy4.mf6.converter import dict_to_array +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_utils import update_maxbound @@ -42,7 +42,7 @@ class Wel(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) aux: Optional[NDArray[np.float64]] = array( @@ -52,7 +52,7 @@ class Wel(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) boundname: Optional[NDArray[np.str_]] = array( @@ -63,6 +63,6 @@ class Wel(Package): "nodes", ), default=None, - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), on_setattr=update_maxbound, ) diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index fcce122b..2c3ac4de 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -7,7 +7,7 @@ from numpy.typing import NDArray from xattree import ROOT, xattree -from flopy4.mf6.converter import dict_to_array +from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, dim, field @@ -27,19 +27,19 @@ class PeriodData: block="perioddata", default=1.0, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) nstp: NDArray[np.int64] = array( block="perioddata", default=1, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) tsmult: NDArray[np.float64] = array( block="perioddata", default=1.0, dims=("nper",), - converter=Converter(dict_to_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) def to_time(self) -> ModelTime: diff --git a/test/test_codec.py b/test/test_codec.py index f0747efb..cb641240 100644 --- a/test/test_codec.py +++ b/test/test_codec.py @@ -1,5 +1,7 @@ from pprint import pprint +import pytest + from flopy4.mf6.codec import dumps, loads from flopy4.mf6.converter import COMPONENT_CONVERTER @@ -53,15 +55,24 @@ def test_dumps_ic(): pprint(loaded) +@pytest.mark.xfail(reason="nested type unstructuring not yet supported") def test_dumps_oc(): from flopy4.mf6.gwf import Oc oc = Oc( + dims={"nper": 1}, budget_file="test.bud", head_file="test.hds", save_head={0: "all"}, save_budget={0: "all"}, - dims={"nper": 1}, + perioddata={ + 0: Oc.PrintSaveSetting( + printrecord=[ + Oc.PrintRecord("head", Oc.Steps(all=True)), + Oc.PrintRecord("budget", Oc.Steps(all=True)), + ], + ) + }, ) dumped = dumps(COMPONENT_CONVERTER.unstructure(oc)) @@ -69,6 +80,8 @@ def test_dumps_oc(): print(dumped) assert "save head all" in dumped assert "save budget all" in dumped + assert "print head all" in dumped + assert "print budget all" in dumped assert dumped loaded = loads(dumped)