support for layer and grid array packages (#262)

Author: mjreno

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
 

docs/examples/twri.py

@@ -8,6 +8,21 @@
 
 import flopy4
 
+
+def plot_head(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,9 @@
 )
 
 # 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
 rch = flopy4.mf6.gwf.Rch(recharge=rate.reshape(nper, -1), dims=dims)
 
 # Output control
@@ -103,24 +118,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 +196,82 @@
 )
 
 # Plot head results
-head.isel(layer=0, time=0).plot.contourf()
+plot_head(head, workspace)
+
+# UPDATE SIM for array based inputs
+
+# 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
+rcha = flopy4.mf6.gwf.Rcha(recharge=recharge.reshape(nper, -1), dims=dims)
+
+# remove list based inputs
+# TODO: show variations on removing packages
+gwf.chd.remove(chd)
+del gwf.drn[0]
+del gwf.wel[0]
+del gwf.rch[0]
+
+# add array based inputs
+# TODO: needs type checking and list consolidation support (see comments in gwf init)
+gwf.chd = [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.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_head(head, workspace)

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

flopy4/mf6/codec/writer/filters.py

@@ -10,7 +10,7 @@
 
 from flopy4.mf6.constants import FILL_DNODATA
 
-ArrayHow = Literal["constant", "internal", "external", "layered internal"]
+ArrayHow = Literal["constant", "internal", "external", "layered constant", "layered internal"]
 
 
 def array_how(value: xr.DataArray) -> ArrayHow:
@@ -29,6 +29,14 @@ def array_how(value: xr.DataArray) -> ArrayHow:
     if value.ndim <= 2:
         return "internal"
     if value.ndim == 3:
+        layer_const = True
+        for layer in range(value.shape[0]):
+            val_layer = value.isel(nlay=layer)
+            if val_layer.max() != val_layer.min():
+                layer_const = False
+                break
+        if layer_const:
+            return "layered constant"
         return "layered internal"
     raise ValueError(f"Arrays with ndim > 3 are not supported, got ndim={value.ndim}")
 

flopy4/mf6/codec/writer/templates/macros.jinja

@@ -35,7 +35,7 @@
 {{ inset }}CONSTANT {{ value|array2const }}
 {% elif how == "layered constant" %}
 {% for layer in value -%}
-{{ inset }}CONSTANT {{ layer|array2const }}
+{{ "\n" ~ inset }}CONSTANT {{ layer|array2const }}
 {%- endfor %}
 {% elif how == "layered internal" %}
 {% for layer in value %}

flopy4/mf6/constants.py

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

flopy4/mf6/context.py

@@ -10,9 +10,29 @@
 from flopy4.utils import to_path
 
 
+def update_child_attr(instance, attribute, new_value):
+    """
+    Generalized function to update child attribute (e.g. workspace).
+
+    Args:
+        instance: The model instance
+        attribute: The attribute being set (from attrs on_setattr)
+        new_value: The new value being set
+
+    Returns:
+        The new_value (unchanged)
+    """
+
+    for child in instance.children.values():  # type: ignore
+        if hasattr(child, attribute.name):
+            setattr(child, attribute.name, new_value)
+
+    return new_value
+
+
 @xattree
 class Context(Component, ABC):
-    workspace: Path = field(default=None, converter=to_path)
+    workspace: Path = field(default=None, converter=to_path, on_setattr=update_child_attr)
 
     def __attrs_post_init__(self):
         super().__attrs_post_init__()