write format adjustments

Author: mjreno

docs/examples/twri.py

@@ -9,7 +9,7 @@
 import flopy4
 
 
-def plot_contourf(head, workspace):
+def plot_head(head, workspace):
     import matplotlib.pyplot as plt
 
     # Plot head results
@@ -101,7 +101,6 @@ def plot_contourf(head, workspace):
 # Uniform recharge on the top layer
 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)
 
@@ -197,7 +196,9 @@ def plot_contourf(head, workspace):
 )
 
 # Plot head results
-plot_contourf(head, workspace)
+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)
@@ -243,18 +244,19 @@ def plot_contourf(head, workspace):
 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]
+gwf.chd.remove(chd)
+# del gwf.chd[0]
 del gwf.drn[0]
 del gwf.wel[0]
 del gwf.rch[0]
 
 # add array based inputs
-gwf.chdg = [chdg]
+# gwf.chdg = [chdg]
+gwf.chd = [chdg]
 gwf.drng = [drng]
 gwf.welg = [welg]
 gwf.rcha = [rcha]
@@ -263,7 +265,6 @@ def plot_contourf(head, workspace):
 workspace = Path(__file__).parent / "twri2"
 workspace.mkdir(parents=True, exist_ok=True)
 sim.workspace = workspace
-sim.__attrs_post_init__()
 
 sim.write()
 sim.run()
@@ -275,4 +276,4 @@ def plot_contourf(head, workspace):
 )
 
 # Plot head results
-plot_contourf(head, workspace)
+plot_head(head, workspace)

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/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__()

flopy4/mf6/converter/unstructure.py

@@ -188,16 +188,16 @@ def _unstructure_block_param(
             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:
+        case xr.DataArray():
             has_spatial_dims = any(
                 dim in field_value.dims for dim in ["nlay", "nrow", "ncol", "ncpl", "nodes"]
             )
             if has_spatial_dims:
                 field_value = _hack_structured_grid_dims(
                     field_value,
-                    structured_grid_dims=value.parent.data.dims,  # type: ignore
+                    structured_grid_dims=value.data.dims,  # type: ignore
                 )
-            if block_name == "period":
+            if "nper" in field_value.dims and 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():
@@ -209,15 +209,13 @@ def _unstructure_block_param(
                     }
             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 or "readasarrays" in xatspec.attrs:
-        # return _unstructure_array_component(value)
         return _unstructure_array_component(value)
     else:
         return _unstructure_component(value)
@@ -273,7 +271,6 @@ 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,
@@ -331,7 +328,7 @@ def _unstructure_component(value: Component) -> dict[str, Any]:
 
     # 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))
+    blocks = dict(sorted(blocks.items(), key=block_sort_key))
 
     # total temporary hack! manually set solutiongroup 1.
     # TODO still need to support multiple..

flopy4/mf6/gwf/__init__.py

@@ -100,6 +100,7 @@ def budget(self):
     chd: list[Chd] = field(block="packages")
     chdg: list[Chdg] = field(block="packages")
     # chd: List[Union[Chd, Chdg]] = field(block="packages", factory=List)
+    # chd: List[Union[Chd, Chdg]] = field(block="packages")
     drn: list[Drn] = field(block="packages")
     drng: list[Drng] = field(block="packages")
     rch: list[Rch] = field(block="packages")

flopy4/mf6/gwf/chdg.py

@@ -2,12 +2,9 @@
 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
@@ -36,7 +33,6 @@ class Chdg(Package):
             "nodes",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
         on_setattr=update_maxbound,
     )
     aux: Optional[NDArray[np.float64]] = array(
@@ -46,6 +42,5 @@ class Chdg(Package):
             "nodes",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
         on_setattr=update_maxbound,
     )

flopy4/mf6/gwf/dis.py

@@ -75,9 +75,15 @@ class Dis(Package):
         scope="gwf",
         init=False,
     )
+    ncpl: int = dim(
+        coord="lnode",
+        scope="gwf",
+        init=False,
+    )
 
     def __attrs_post_init__(self):
         self.nodes = self.ncol * self.nrow * self.nlay
+        self.ncpl = self.ncol * self.nrow
         super().__attrs_post_init__()
 
     def to_grid(self) -> StructuredGrid:

flopy4/mf6/gwf/drng.py

@@ -2,12 +2,9 @@
 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
@@ -37,7 +34,6 @@ class Drng(Package):
             "nodes",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
         on_setattr=update_maxbound,
     )
     cond: Optional[NDArray[np.float64]] = array(
@@ -47,7 +43,6 @@ class Drng(Package):
             "nodes",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
         on_setattr=update_maxbound,
     )
     aux: Optional[NDArray[np.float64]] = array(
@@ -57,6 +52,5 @@ class Drng(Package):
             "nodes",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
         on_setattr=update_maxbound,
     )

flopy4/mf6/gwf/rcha.py

@@ -2,16 +2,11 @@
 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
 
 
@@ -31,16 +26,13 @@ class Rcha(Package):
     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",
@@ -49,8 +41,6 @@ class Rcha(Package):
             "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",
@@ -59,6 +49,4 @@ class Rcha(Package):
             "ncpl",
         ),
         default=None,
-        # converter=Converter(structure_array, takes_self=True, takes_field=True),
-        # on_setattr=update_maxbound,
     )