add welg, rcha baseline

Author: mjreno

docs/examples/twri.py

@@ -8,6 +8,21 @@
 
 import flopy4
 
+
+def plot_contourf(head, workspace):
+    import matplotlib.pyplot as plt
+
+    # Plot head results
+    plt.figure(figsize=(10, 6))
+    head.isel(layer=0, time=0).plot.contourf()
+    plt.title("Filled Contour Plot TWRI Head")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    plt.grid(True)
+    plt.savefig(workspace / "head.png", dpi=300, bbox_inches="tight")
+    plt.close()
+
+
 # Timing
 time = flopy4.mf6.utils.time.Time.from_timestamps(
     ["2000-01-01", "2000-01-02", "2000-01-03", "2000-01-04"]
@@ -28,7 +43,7 @@
 grid = flopy4.mf6.utils.grid.StructuredGrid(
     nlay=nlay, nrow=nrow, ncol=ncol, top=top, botm=bottom, delr=delr, delc=delc, idomain=idomain
 )
-dims = {"nper": nper, **dict(grid.dataset.sizes)}  # TODO: temporary
+dims = {"nper": nper, "ncpl": nrow * ncol, **dict(grid.dataset.sizes)}  # TODO: temporary
 
 # Grid discretization
 # TODO: xorigin, yorigin
@@ -84,9 +99,10 @@
 )
 
 # Uniform recharge on the top layer
-rch_rate = np.stack(np.full((nlay, nrow, ncol), flopy4.mf6.constants.FILL_DNODATA))
+rch_rate = np.full((nlay, nrow, ncol), flopy4.mf6.constants.FILL_DNODATA)
 rate = np.repeat(np.expand_dims(rch_rate, axis=0), repeats=nper, axis=0)
-rate[0, 0, :] = 3.0e-8
+# rate[0, 0, :, :] = 3.0e-8
+rate[0, 0, ...] = 3.0e-8
 rch = flopy4.mf6.gwf.Rch(recharge=rate.reshape(nper, -1), dims=dims)
 
 # Output control
@@ -103,24 +119,24 @@
 # Wells scattered throughout the model
 wel_q = -5.0
 wel_nodes = [
-    [2, 4, 10, -5.0],
-    [1, 3, 5, -5.0],
-    [1, 5, 11, -5.0],
-    [0, 8, 7, -5.0],
-    [0, 8, 9, -5.0],
-    [0, 8, 11, -5.0],
-    [0, 8, 13, -5.0],
-    [0, 10, 7, -5.0],
-    [0, 10, 9, -5.0],
-    [0, 10, 11, -5.0],
-    [0, 10, 13, -5.0],
-    [0, 12, 7, -5.0],
-    [0, 12, 9, -5.0],
-    [0, 12, 11, -5.0],
-    [0, 12, 13, -5.0],
+    [2, 4, 10],
+    [1, 3, 5],
+    [1, 5, 11],
+    [0, 8, 7],
+    [0, 8, 9],
+    [0, 8, 11],
+    [0, 8, 13],
+    [0, 10, 7],
+    [0, 10, 9],
+    [0, 10, 11],
+    [0, 10, 13],
+    [0, 12, 7],
+    [0, 12, 9],
+    [0, 12, 11],
+    [0, 12, 13],
 ]
 wel = flopy4.mf6.gwf.Wel(
-    q={"*": {(layer, row, col): wel_q for layer, row, col, wel_q in wel_nodes}},
+    q={"*": {(layer, row, col): wel_q for layer, row, col in wel_nodes}},
     dims=dims,
 )
 
@@ -181,4 +197,82 @@
 )
 
 # Plot head results
-head.isel(layer=0, time=0).plot.contourf()
+plot_contourf(head, workspace)
+
+# update simulation with array based inputs
+LAYER_NODATA = np.full((nrow, ncol), flopy4.mf6.constants.FILL_DNODATA, dtype=float)
+GRID_NODATA = np.full((nlay, nrow, ncol), flopy4.mf6.constants.FILL_DNODATA, dtype=float)
+
+head = np.repeat(np.expand_dims(GRID_NODATA, axis=0), repeats=nper, axis=0)
+for i in range(nrow):
+    for k in range(nlay - 1):
+        head[0, k, i, 0] = 0.0
+chdg = flopy4.mf6.gwf.Chdg(
+    print_input=True,
+    print_flows=True,
+    save_flows=True,
+    head=head.reshape(nper, -1),
+    dims=dims,
+)
+
+# Drain in the center left of the model
+elev = np.repeat(np.expand_dims(GRID_NODATA, axis=0), repeats=nper, axis=0)
+cond = np.repeat(np.expand_dims(GRID_NODATA, axis=0), repeats=nper, axis=0)
+for j in range(9):
+    elev[0, 0, 7, j + 1] = elevation[j]
+    cond[0, 0, 7, j + 1] = conductance
+drng = flopy4.mf6.gwf.Drng(
+    print_input=True,
+    print_flows=True,
+    save_flows=True,
+    elev=elev.reshape(nper, -1),
+    cond=cond.reshape(nper, -1),
+    dims=dims,
+)
+
+# well
+q = np.repeat(np.expand_dims(GRID_NODATA, axis=0), repeats=nper, axis=0)
+for layer, row, col in wel_nodes:
+    q[0, layer, row, col] = wel_q
+welg = flopy4.mf6.gwf.Welg(
+    q=q.reshape(nper, -1),
+    dims=dims,
+)
+
+# recharge
+recharge = np.repeat(np.expand_dims(LAYER_NODATA, axis=0), repeats=nper, axis=0)
+recharge[0, ...] = 3.0e-8
+# recharge[0, 0, 0] = 3.0e-7
+# print(recharge)
+# rch = flopy4.mf6.gwf.Rcha(irch=1, recharge=recharge.reshape(nper, -1), dims=dims)
+rcha = flopy4.mf6.gwf.Rcha(recharge=recharge.reshape(nper, -1), dims=dims)
+
+# remove list based inputs
+del gwf.chd[0]
+del gwf.drn[0]
+del gwf.wel[0]
+del gwf.rch[0]
+
+# add array based inputs
+gwf.chdg = [chdg]
+gwf.drng = [drng]
+gwf.welg = [welg]
+gwf.rcha = [rcha]
+
+# create new workspace
+workspace = Path(__file__).parent / "twri2"
+workspace.mkdir(parents=True, exist_ok=True)
+sim.workspace = workspace
+sim.__attrs_post_init__()
+
+sim.write()
+sim.run()
+
+# Load head results
+head = flopy4.mf6.utils.open_hds(
+    workspace / f"{gwf.name}.hds",
+    workspace / f"{gwf.name}.dis.grb",
+)
+
+# Plot head results
+plot_contourf(head, workspace)

flopy4/mf6/binding.py

@@ -34,7 +34,7 @@ def _get_binding_type(component: Component) -> str:
                 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[0:3].upper()}6"
                 return f"{cls_name.upper()}6"
 
         def _get_binding_terms(component: Component) -> tuple[str, ...] | None:

flopy4/mf6/converter/unstructure.py

@@ -190,7 +190,7 @@ def _unstructure_block_param(
             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"]
+                dim in field_value.dims for dim in ["nlay", "nrow", "ncol", "ncpl", "nodes"]
             )
             if has_spatial_dims:
                 field_value = _hack_structured_grid_dims(
@@ -216,19 +216,14 @@ def _unstructure_block_param(
 
 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)
+    if "readarraygrid" in xatspec.attrs or "readasarrays" in xatspec.attrs:
+        # return _unstructure_array_component(value)
+        return _unstructure_array_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]:
+def _unstructure_array_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))
@@ -243,7 +238,6 @@ def _unstructure_grid_component(value: Component) -> dict[str, Any]:
     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] = {}
@@ -279,6 +273,7 @@ def _unstructure_component(value: Component) -> dict[str, Any]:
     data = xattree.asdict(value)
 
     # create child component binding blocks
+    # TODO: this can add certain block keys out of order
     blocks.update(_make_binding_blocks(value))
 
     # process blocks in order, unstructuring fields as needed,
@@ -287,7 +282,6 @@ def _unstructure_component(value: Component) -> dict[str, Any]:
     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] = {}
@@ -335,6 +329,10 @@ def _unstructure_component(value: Component) -> dict[str, Any]:
         if perioddata := blocks.get("perioddata", None):
             blocks["perioddata"] = {"perioddata": xr.Dataset(perioddata)}
 
+    # TODO: this fixes out of order blocks (e.g. model namefile) from
+    # blocks.update() child binding call above
+    # blocks = dict(sorted(blocks.items(), key=block_sort_key))
+
     # total temporary hack! manually set solutiongroup 1.
     # TODO still need to support multiple..
     if "solutiongroup" in blocks:

flopy4/mf6/gwf/__init__.py

@@ -16,14 +16,31 @@
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
 from flopy4.mf6.gwf.rch import Rch
+from flopy4.mf6.gwf.rcha import Rcha
 from flopy4.mf6.gwf.sto import Sto
 from flopy4.mf6.gwf.wel import Wel
+from flopy4.mf6.gwf.welg import Welg
 from flopy4.mf6.model import Model
 from flopy4.mf6.spec import field, path
 from flopy4.mf6.utils import open_cbc, open_hds
 from flopy4.utils import to_path
 
-__all__ = ["Gwf", "Chd", "Chdg", "Dis", "Drn", "Drng", "Ic", "Npf", "Oc", "Rch", "Sto", "Wel"]
+__all__ = [
+    "Gwf",
+    "Chd",
+    "Chdg",
+    "Dis",
+    "Drn",
+    "Drng",
+    "Ic",
+    "Npf",
+    "Oc",
+    "Rch",
+    "Rcha",
+    "Sto",
+    "Wel",
+    "Welg",
+]
 
 
 def convert_grid(value):
@@ -82,10 +99,13 @@ def budget(self):
     sto: Sto | None = field(block="packages", default=None)
     chd: list[Chd] = field(block="packages")
     chdg: list[Chdg] = field(block="packages")
+    # chd: List[Union[Chd, Chdg]] = field(block="packages", factory=List)
     drn: list[Drn] = field(block="packages")
     drng: list[Drng] = field(block="packages")
     rch: list[Rch] = field(block="packages")
+    rcha: list[Rcha] = field(block="packages")
     wel: list[Wel] = field(block="packages")
+    welg: list[Welg] = field(block="packages")
     output: Output = attrs.field(
         default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
     )

flopy4/mf6/gwf/rcha.py

@@ -0,0 +1,64 @@
+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 Rcha(Package):
+    multi_package: ClassVar[bool] = True
+    fixed_cell: bool = field(block="options", default=False)
+    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)
+    readasarrays: bool = field(block="options", default=True)
+    save_flows: bool = field(block="options", default=False)
+    tas_filerecord: Optional[Path] = path(
+        block="options", default=None, converter=to_path, inout="filein"
+    )
+    obs_filerecord: Optional[Path] = path(
+        block="options", default=None, converter=to_path, inout="fileout"
+    )
+    # maxbound: Optional[int] = field(block="dimensions", default=None, init=False)
+    irch: Optional[NDArray[np.int64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            "ncpl",
+        ),
+        default=None,
+        # converter=Converter(structure_array, takes_self=True, takes_field=True),
+        # on_setattr=update_maxbound,
+    )
+    recharge: Optional[NDArray[np.float64]] = array(
+        block="period",
+        dims=(
+            "nper",
+            "ncpl",
+        ),
+        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",
+            "ncpl",
+        ),
+        default=None,
+        # converter=Converter(structure_array, takes_self=True, takes_field=True),
+        # on_setattr=update_maxbound,
+    )