add multi period test with storage

Author: mjreno

flopy4/mf6/adapters.py

@@ -264,7 +264,6 @@ def plottable(self):
 
     @property
     def has_stress_period_data(self):
-        # TODO oc returns true? is stress package?
         return "nper" in self._data.dims
 
     def check(self, f=None, verbose=True, level=1, checktype=None):

flopy4/mf6/codec/writer/__init__.py

@@ -16,6 +16,7 @@
 _JINJA_ENV.filters["array_how"] = filters.array_how
 _JINJA_ENV.filters["array_chunks"] = filters.array_chunks
 _JINJA_ENV.filters["array2string"] = filters.array2string
+_JINJA_ENV.filters["data2const"] = filters.data2const
 _JINJA_ENV.filters["data2list"] = filters.data2list
 _JINJA_ENV.filters["data2keystring"] = filters.data2keystring
 _JINJA_TEMPLATE_NAME = "blocks.jinja"

flopy4/mf6/codec/writer/filters.py

@@ -63,6 +63,8 @@ def array_how(value: xr.DataArray) -> str:
     # TODO
     # - detect constant arrays?
     # - above certain size, use external?
+    if value.max() == value.min():
+        return "constant"
     return "internal"
 
 
@@ -151,6 +153,13 @@ def nonempty(arr: NDArray | xr.DataArray) -> NDArray:
     return mask
 
 
+def data2const(value: xr.DataArray):
+    if np.issubdtype(value.dtype, np.integer):
+        return value.max().item()
+    if np.issubdtype(value.dtype, np.floating):
+        return f"{value.max().item():.8f}"
+
+
 def data2list(value: list | xr.DataArray | xr.Dataset):
     """
     Yield record tuples from a list, `DataArray` or `Dataset`.
@@ -234,6 +243,9 @@ def dataset2list(value: xr.Dataset):
                 field_vals.append(field_val.item())
             else:
                 field_vals.append(field_val)
+        for j, v in enumerate(field_vals):
+            if isinstance(field_vals[j], float):
+                field_vals[j] = f"{field_vals[j]:.8f}"
         if has_spatial_dims:
             cellid = tuple(idx[i] + 1 for idx in indices)
             yield cellid + tuple(field_vals)

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

@@ -40,7 +40,7 @@
 {{ inset ~ name.upper() }}{% if "layered" in how %} LAYERED{% endif %}
 
 {% if how == "constant" %}
-    CONSTANT {{ value.item() }}
+    CONSTANT {{ value|data2const -}}
 {% elif how == "layered constant" %}
 {% for layer in value -%}
     CONSTANT {{ layer.item() }}

flopy4/mf6/config.py

@@ -1,3 +1,3 @@
 # TODO use https://environ-config.readthedocs.io/en/stable/?
 
-SPARSE_THRESHOLD = 1000
+SPARSE_THRESHOLD = 100000

flopy4/mf6/converter.py

@@ -87,6 +87,54 @@ def _path_to_record(field_name: str, path_value: Path) -> tuple:
     return (field_name.upper(), "FILEOUT", str(path_value))
 
 
+def _user_dims(value, field_value):
+    # terrible hack to convert flat nodes dimension to 3d structured dims.
+    # long term solution for this is to use a custom xarray index. filters
+    # should then have access to all dimensions needed.
+    dims_ = set(field_value.dims).copy()
+    parent = value.parent  # type: ignore
+    if parent is None:
+        # TODO for standalone packages
+        return field_value
+
+    if "nper" in dims_:
+        dims_.remove("nper")
+        shape = (
+            field_value.sizes["nper"],
+            parent.dims["nlay"],
+            parent.dims["nrow"],
+            parent.dims["ncol"],
+        )
+        dims = ("nper", "nlay", "nrow", "ncol")
+        coords = {
+            "nper": field_value.coords["nper"],
+            "nlay": range(parent.dims["nlay"]),
+            "nrow": range(parent.dims["nrow"]),
+            "ncol": range(parent.dims["ncol"]),
+        }
+    else:
+        shape = (
+            parent.dims["nlay"],
+            parent.dims["nrow"],
+            parent.dims["ncol"],
+        )
+        dims = ("nlay", "nrow", "ncol")
+        coords = {
+            "nlay": range(parent.dims["nlay"]),
+            "nrow": range(parent.dims["nrow"]),
+            "ncol": range(parent.dims["ncol"]),
+        }
+
+    if dims_ == {"nodes"}:
+        field_value = xr.DataArray(
+            field_value.data.reshape(shape),
+            dims=dims,
+            coords=coords,
+        )
+
+    return field_value
+
+
 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]] = {}
@@ -161,47 +209,27 @@ def unstructure_component(value: Component) -> dict[str, Any]:
                     dim in field_value.dims for dim in ["nlay", "nrow", "ncol", "nodes"]
                 )
                 if has_spatial_dims:
-                    # terrible hack to convert flat nodes dimension to 3d structured dims.
-                    # long term solution for this is to use a custom xarray index. filters
-                    # should then have access to all dimensions needed.
-                    dims_ = set(field_value.dims).copy()
-                    dims_.remove("nper")
-                    if dims_ == {"nodes"}:
-                        parent = value.parent  # type: ignore
-                        field_value = xr.DataArray(
-                            field_value.data.reshape(
-                                (
-                                    field_value.sizes["nper"],
-                                    parent.dims["nlay"],
-                                    parent.dims["nrow"],
-                                    parent.dims["ncol"],
-                                )
-                            ),
-                            dims=("nper", "nlay", "nrow", "ncol"),
-                            coords={
-                                "nper": field_value.coords["nper"],
-                                "nlay": range(parent.dims["nlay"]),
-                                "nrow": range(parent.dims["nrow"]),
-                                "ncol": range(parent.dims["ncol"]),
-                            },
-                            name=field_value.name,
-                        )
+                    field_value = _user_dims(value, field_value)
 
                     period_data[field_name] = {
                         kper: field_value.isel(nper=kper)
                         for kper in range(field_value.sizes["nper"])
                     }
                 else:
-                    if (
-                        # TODO: refactor
-                        # field_name == "save_budget"
-                        # or field_name == "save_head"
-                        # or field_name == "print_budget"
-                        # or field_name == "print_head"
-                        np.issubdtype(field_value.dtype, np.str_)
-                    ):
+                    if value.__class__.__name__ == "Oc":
                         period_data[field_name] = {
-                            kper: field_value[kper] for kper in range(field_value.sizes["nper"])
+                            kper: field_value.values[kper]
+                            for kper in range(field_value.sizes["nper"])
+                            if field_value.values[kper] is not None
+                        }
+                    elif np.issubdtype(field_value.dtype, np.str_):
+                        fname = field_name
+                        if value.__class__.__name__ == "Sto":
+                            fname = field_name.replace("_", "-")
+                        period_data[fname] = {
+                            kper: field_value[kper]
+                            for kper in range(field_value.sizes["nper"])
+                            if field_value[kper] is not None
                         }
                     else:
                         if block_name not in period_data:
@@ -212,6 +240,8 @@ def unstructure_component(value: Component) -> dict[str, Any]:
                     if isinstance(field_value, bool):
                         if field_value:
                             blocks[block_name][field_name] = field_value
+                    elif isinstance(field_value, xr.DataArray):
+                        blocks[block_name][field_name] = _user_dims(value, field_value)
                     else:
                         blocks[block_name][field_name] = field_value
 
@@ -223,9 +253,20 @@ def unstructure_component(value: Component) -> dict[str, Any]:
 
         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
+                if isinstance(arr, xr.DataArray):
+                    max = arr.max()
+                    if max == arr.min() and max == FILL_DNODATA:
+                        # don't write empty period blocks unless
+                        # to intentionally reset data
+                        pass
+                    else:
+                        if kper not in period_blocks:
+                            period_blocks[kper] = {}
+                        period_blocks[kper][arr_name] = arr
+                else:
+                    if kper not in period_blocks:
+                        period_blocks[kper] = {}
+                    period_blocks[kper][arr_name] = arr.upper()
 
         for kper, block in period_blocks.items():
             dataset = xr.Dataset(block)

flopy4/mf6/gwf/__init__.py

@@ -13,12 +13,13 @@
 from flopy4.mf6.gwf.ic import Ic
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
+from flopy4.mf6.gwf.sto import Sto
 from flopy4.mf6.gwf.wel import Wel
 from flopy4.mf6.model import Model
 from flopy4.mf6.spec import field
 from flopy4.mf6.utils import open_cbc, open_hds
 
-__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Wel"]
+__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Sto", "Wel"]
 
 
 def convert_grid(value):
@@ -66,11 +67,12 @@ def budget(self):
     nc_filerecord: Optional[Path] = field(block="options", default=None)
     dis: Dis = field(converter=convert_grid, block="packages")
     ic: Ic = field(block="packages")
-    oc: Oc = field(block="packages")
+    oc: Oc = field(block="packages", default=None)
     npf: Npf = field(block="packages")
+    sto: Sto = field(block="packages", default=None)
     chd: list[Chd] = field(block="packages")
-    wel: list[Wel] = field(block="packages")
     drn: list[Drn] = field(block="packages")
+    wel: list[Wel] = field(block="packages")
     output: Output = attrs.field(
         default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
     )

flopy4/mf6/gwf/oc.py

@@ -55,31 +55,31 @@ class Period:
     save_head: Optional[NDArray[np.object_]] = array(
         object,
         block="period",
-        default="all",
+        default=None,
         dims=("nper",),
         converter=Converter(dict_to_array, takes_self=True, takes_field=True),
         format="keystring",
     )
     save_budget: Optional[NDArray[np.object_]] = array(
         object,
         block="period",
-        default="all",
+        default=None,
         dims=("nper",),
         converter=Converter(dict_to_array, takes_self=True, takes_field=True),
         format="keystring",
     )
     print_head: Optional[NDArray[np.object_]] = array(
         object,
         block="period",
-        default="all",
+        default=None,
         dims=("nper",),
         converter=Converter(dict_to_array, takes_self=True, takes_field=True),
         format="keystring",
     )
     print_budget: Optional[NDArray[np.object_]] = array(
         object,
         block="period",
-        default="all",
+        default=None,
         dims=("nper",),
         converter=Converter(dict_to_array, takes_self=True, takes_field=True),
         format="keystring",

flopy4/mf6/gwf/sto.py

@@ -39,13 +39,7 @@ class Sto(Package):
         default=0.15,
         converter=Converter(dict_to_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),
-    )
-    transient: Optional[NDArray[np.bool_]] = array(
+    storage: Optional[NDArray[np.str_]] = array(
         block="period",
         dims=("nper",),
         default=None,

test/test_codec.py

@@ -253,11 +253,11 @@ def test_dumps_drn():
     assert len(period2_lines) == 3
 
     # node elev cond
-    assert "1 1 5 10.0 1.0" in dumped  # Period 1: (0,0,4)
-    assert "2 5 1 8.0 2.0" in dumped  # Period 1: (1,4,0)
-    assert "1 2 2 12.0 1.5" in dumped  # Period 2: (0,1,1)
-    assert "1 3 4 9.0 0.8" in dumped  # Period 2: (0,2,3)
-    assert "2 4 3 7.0 2.2" in dumped  # Period 2: (1,3,2)
+    assert "1 1 5 10.00000000 1.00000000" in dumped  # Period 1: (0,0,4)
+    assert "2 5 1 8.00000000 2.00000000" in dumped  # Period 1: (1,4,0)
+    assert "1 2 2 12.00000000 1.50000000" in dumped  # Period 2: (0,1,1)
+    assert "1 3 4 9.00000000 0.80000000" in dumped  # Period 2: (0,2,3)
+    assert "2 4 3 7.00000000 2.20000000" in dumped  # Period 2: (1,3,2)
     assert "1e+30" not in dumped
     assert "1.0e+30" not in dumped
 
@@ -334,8 +334,8 @@ def test_dumps_wel_with_aux():
 
     assert len(lines) == 2
     # node q aux_value
-    assert "1 2 3 -75.0 1.0" in dumped  # (0,1,2) -> node 8, q=-75.0, aux=1.0
-    assert "2 4 5 -25.0 2.0" in dumped  # (1,3,4) -> node 45, q=-25.0, aux=2.0
+    assert "1 2 3 -75.00000000 1.00000000" in dumped  # (0,1,2) -> node 8, q=-75.0, aux=1.0
+    assert "2 4 5 -25.00000000 2.00000000" in dumped  # (1,3,4) -> node 45, q=-25.0, aux=2.0
     assert "1e+30" not in dumped
     assert "1.0e+30" not in dumped