chdg baseline

Author: mjreno

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:

flopy4/mf6/constants.py

@@ -2,5 +2,5 @@
 
 MF6 = "mf6"
 FILL_DEFAULT = np.nan
-FILL_DNODATA = 1e30
+FILL_DNODATA = 3.0e30
 LENBOUNDNAME = 40

flopy4/mf6/converter/unstructure.py

@@ -7,6 +7,7 @@
 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
@@ -144,7 +145,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]:
+    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]:
     from flopy4.mf6.constants import FILL_DNODATA
 
     blockspec = dict(sorted(value.dfn.blocks.items(), key=block_sort_key))  # type: ignore
@@ -167,67 +250,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[f"period {kper + 1}"][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():

flopy4/mf6/gwf/__init__.py

@@ -8,6 +8,7 @@
 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.ic import Ic
@@ -21,7 +22,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", "Ic", "Npf", "Oc", "Rch", "Sto", "Wel"]
 
 
 def convert_grid(value):
@@ -79,6 +80,7 @@ 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")
+    chdg: list[Chdg] = field(block="packages")
     wel: list[Wel] = field(block="packages")
     drn: list[Drn] = field(block="packages")
     rch: list[Rch] = field(block="packages")

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,
+    )

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"